Title: Reliable Off-Policy Learning for Dosage Combinations
Venue: NeurIPS 2023
Accepted: yes
Keywords: off-policy learning

causal inference

reliable machine learning

medicine

dosaging

normalizing flows
Abstract: Decision-making in personalized medicine such as cancer therapy or critical care must often make choices for dosage combinations, i.e., multiple continuous treatments. Existing work for this task has modeled the effect of multiple treatments independently, while estimating the joint effect has received little attention but comes with non-trivial challenges. In this paper, we propose a novel method for reliable off-policy learning for dosage combinations. Our method proceeds along three steps: (1) We develop a tailored neural network that estimates the individualized dose-response function while accounting for the joint effect of multiple dependent dosages. (2) We estimate the generalized propensity score using conditional normalizing flows in order to detect regions with limited overlap in the shared covariate-treatment space. (3) We present a gradient-based learning algorithm to find the optimal, individualized dosage combinations. Here, we ensure reliable estimation of the policy value by avoiding regions with limited overlap. We finally perform an extensive evaluation of our method to show its effectiveness. To the best of our knowledge, ours is the first work to provide a method for reliable off-policy learning for optimal dosage combinations.

Reliable Off-Policy Learning for Dosage Combinations

Abstract

Decision-making in personalized medicine such as cancer therapy or critical care must often make choices for dosage combinations, i.e., multiple continuous treatments. Existing work for this task has modeled the effect of multiple treatments independently , while estimating the joint effect has received little attention but comes with non-trivial challenges. In this paper, we propose a novel method for reliable off-policy learning for dosage combinations. Our method proceeds along three steps: (1) We develop a tailored neural network that estimates the individualized dose-response function while accounting for the joint effect of multiple dependent dosages. (2) We estimate the generalized propensity score using conditional normalizing flows in order to detect regions with limited overlap in the shared covariate-treatment space. (3) We present a gradient-based learning algorithm to find the optimal, individualized dosage combinations. Here, we ensure reliable estimation of the policy value by avoiding regions with limited overlap. We finally perform an extensive evaluation of our method to show its effectiveness. To the best of our knowledge, ours is the first work to provide a method for reliable off-policy learning for optimal dosage combinations.

1 Introduction

In personalized medicine, decision-making frequently involves complex choices for dosage combinations . For example, in cancer therapy, patients with chemotherapy or immunotherapy receive a combination of two or three drugs, and, for each, the dosage must be carefully personalized [16, 82]. In critical care, medical professionals must simultaneously control multiple parameters for mechanical ventilation, including respiratory rate and tidal volume [20].

One way to find optimal dosages (i.e., multiple continuous treatments) is through randomized controlled trials. Yet, such randomized controlled trials are often impractical due to the complexity of dosing problems or even unethical [62]. Instead, there is a direct need for learning optimal, individualized dosage combinations from observational data (so-called off-policy learning).

Several methods for off-policy learning aim at discrete treatments [1, 3, 4, 14, 23, 37, 40, 74, 76] and are thus not applicable to dosage combinations. In contrast to that, there are only a few methods that aim at off-policy learning for continuous treatments, that is, dosages [6, 10, 41, 44, 87]. To achieve this, some of these methods utilize the individualized dose-response function, which is typically not available and must be estimated from observational data. However, existing works model the effect of multiple continuous treatments independently . Because of that, existing works ignore the dependence structure among dosage combinations, that is, drug-drug interactions. For example, in chemotherapy, the combination of antineoplastic agents with drugs such as warfarin can have adverse effects on patient health [61]. In contrast to that, methods for off-policy learning that account for the joint effect of dosage combinations have received little attention and present the focus of our paper.

Crucially, off-policy learning for dosage combinations comes with non-trivial challenges. (1) Empirically, datasets in medical practice are often high-dimensional, yet of limited sample size. Hence, in practice, there are often sparse regions in the covariate-treatment space, which, in turn, may lead to estimated individualized dose-response functions that have regions with high uncertainty. (2) Mathematically, off-policy learning commonly relies upon the overlap assumption [65]. Yet, certain drug combinations are rarely used in medical practice and should be avoided for safety reasons. Both challenges lead to areas with limited observed overlap in the data, which results in unreliable estimates of the policy value. Motivated by this, we thus develop a method for reliable off-policy learning.

Proposed method: In this paper, we propose a novel method for reliable off-policy learning for dosage combinations. Our method proceeds along three steps. (1) We develop a tailored neural network that estimates the individualized dose-response function while accounting for the joint effect of multiple dependent dosages. Later, we use the network as a plug-in estimator for the policy value. (2) We estimate the generalized propensity score through conditional normalizing flows in order to detect regions with limited overlap in the shared covariate-treatment space. (3) We find the optimal individualized dosage combinations. To achieve that, we present a gradient-descent-ascent-based learning algorithm that solves the underlying constrained optimization problem, where we avoid regions with limited overlap and thus ensure a reliable estimation of the policy value. For the latter step, we use an additional policy neural network to predict the optimal policy for reasons of scalability.

Contributions: 1 (1) We propose a novel method for reliable off-policy learning for dosage combinations. To the best of our knowledge, ours is the first method that accounts for the challenges inherent to dosage combinations. (2) As part of our method, we develop a novel neural network for modeling the individualized dose-response function while accounting for the joint effect of dosage combinations. Thereby, we extend previous works which are limited to mutually exclusive or independent treatments. (3) We perform extensive experiments based on real-world medical data to show the effectiveness of our method.

2 Related Work

Estimating the individualized dose-response function: There are different methods to estimate the outcomes of continuous treatments, that, is the individualized dose-response function. One stream uses the generalized propensity score [26, 28, 29], while others adapt generative adversarial networks [5]. An even other stream builds upon the idea of shared representation learning from binary treatments (e.g., [33, 53, 68, 69, 84]) and extends it to continuous treatments [57, 67]. However, all of the previous methods model the marginal effects of multiple dosages independently and thus without considering the dependence among when multiple treatments are applied simultaneously. Yet, this is unrealistic in medical practice as there are generally drug-drug interactions. To address this, we later model the joint effect of multiple dependent dosages.

There are some approaches for estimating the individualized effects of multi- or high-dimensional treatments, yet not for dosages. For example, some methods target conditional multi-cause treatment effect estimation [55, 58, 59, 73, 78, 86] but are restricted to discrete treatments and can not be directly applied to dosages. Other works use techniques tailored for high-dimensional structured treatments [21, 36, 54]. However, they do not focus on smooth and continuous estimates, or reliable off-policy learning, which is unlike our method.

Off-policy learning for dosages: There is extensive work on off-policy learning for discrete treatments [1, 3, 4, 14, 17, 23, 37, 40, 49, 50, 72, 74, 85]. While, in principle, these methods could be applied to our setting by discretizing the dosages, crucial information about the exact individually prescribed dosage values would be lost. This results in less accurate and less personalized dosage recommendations. However, there are only a few papers that address off-policy learning directly for continuous treatments ('dosages').

For dosages, previous works have derived efficient estimators for policy evaluation and subsequent optimization (e.g., [10, 41]). Other works focus directly on the policy optimization objective to find an optimal dosage [6, 44, 87]. In principle, some of these methods are applicable to settings with multiple dosages where they act as naïve baselines. However, all of them have crucial shortcomings

in that they do not account for multiple dependent dosages and/or that they do not aim at reliable decision-making under limited overlap. Both of the latter are contributions of our method.

Dealing with limited overlap: Various methods have been proposed to address limited overlap for estimating average treatment effects of binary treatments by making use of the propensity score (e.g., [7, 9, 24, 30, 38, 66]). In contrast, there exist only a few works for off-policy learning with limited overlap, typically for binary treatments [39, 42, 70].

Some works deal with the violation of the overlap assumption by leveraging uncertainty estimates for binary [31] and continuous treatment effect estimation [32]. However, these works do not tackle the problem of direct policy learning. In particular for dosage combinations, this is a clear shortcoming, as determining and evaluating policies becomes computationally expensive and infeasible with a growing number of samples and treatment dimensions.

Research gap: Personalized decision-making for optimal dosage combinations must (1) estimate the joint effect of dosage combinations, (2) then perform off-policy learning for such dosage combinations, and (3) be reliable even for limited overlap. However, to the best of our knowledge, there is no existing method for this composite task, and we thus present the first method aimed at reliable off-policy learning for dosage combinations.

3 Problem Formulation

Setting: We consider an observational dataset for n patients with i.i.d. observations { ( x i , t i , y i ) } n i =1 sampled from a population ( X,T,Y ) , where X ∈ X ⊆ R d are the patients' covariates, T ∈ T ⊆ R p is the assigned p -dimensional dosage combination (i.e., multiple continuous treatments), and Y ∈ R is the observed outcome. For example, in a mechanical ventilation setting, X are patient features and measured biosignals, such as age, sex, respiratory measurements, and cardiac measurements; T are the dosages of the ventilation parameters, such as respiratory rate and tidal volume; and Y is patient survival. For simplicity, we also refer to the dosages as treatments.

Weadopt the potential outcomes framework [64], where Y ( t ) is the potential outcome under treatment T = t . In medicine, one is interested in learning a policy π : X → T that maps covariates to suggested dosage combinations. Then, the policy value V ( π ) = E [ Y ( π ( X ))] is the expected outcome of policy π . Our objective is to find the optimal policy π ∗ in the policy class Π that maximizes V ( π ) , i.e.,

In other words, we want to find π ∗ that minimizes the regret given by R ( π ∗ ) = max π ∈ Ω V ( π ) -V ( π ∗ ) , where Ω denotes the set of all possible policies with Π ⊆ Ω .

We further introduce the following notation, which is relevant to our method later. We denote the conditional potential outcome function by µ ( t, x ) = E [ Y ( t ) | X = x ] . For continuous treatments, this is also called individualized dose-response function . Further, the generalized propensity score (GPS) is given by f ( t, x ) = f T | X = x ( t ) , where f T | X = x denotes the conditional density of T given X = x with respect to the Lebesgue measure. The GPS is the stochastic policy used for assigning the observed treatments (i.e., logging / behavioral policy). In our method, both nuisance functions, µ ( t, x ) and f ( t, x ) , are not known and must be later estimated from observational data.

Assumptions: To ensure identifiability of the outcome estimation, the following three standard assumptions are commonly made in causal inference [65]: (1) Consistency : Y = Y ( T ) . Consistency requires that the patient's observed outcome is equal to the potential outcome under the observed treatment. (2) Ignorability : Y ( t ) ⊥ ⊥ T | X, ∀ t ∈ T . Ignorability ensures that there are no unobserved confounders. (3) Overlap : f ( t, x ) > ε , ∀ x ∈ X , t ∈ T and for some constant ε ∈ [0 , ∞ ) . Overlap (also known as common support or positivity) is necessary to ensure that the outcomes for all potential treatments can be estimated accurately for all individuals.

We distinguish between weak overlap for ε = 0 and strong overlap for ε > 0 . Since the estimation of the individualized dose-response function becomes more reliable with increasing ε in the finite sample setting, we refer to strong overlap in the following when we mention overlap, unless explicitly stated otherwise.

Challenges due to limited overlap: Limited overlap introduces both empirical and theoretical challenges. (1) Datasets with high dimensions and/or small sample size are especially prone to have

sparse regions in the covariate-treatment space X × T and thus limited overlap [9]. As a result, the estimated individualized dose-response function will have regions with high uncertainty, which preclude reliable decision-making. (2) It is likely that ε is small for some regions as some drug combinations are rarely prescribed and should be avoided. To address this, we aim at an off-policy method that is reliable, as introduced in Sec. 4.

4 Reliable Off-Policy Learning for Dosage Combinations

In the following, we propose a novel method for reliable off-policy learning for dosage combinations. Our method proceeds along three steps (see Fig. 1): (1) We develop a tailored neural network that estimates the individualized dose-response function while accounting for the joint effect of multiple dependent dosages (Sec. 4.1). Later, we use the network as a plug-in estimator for the policy value. (2) We estimate the generalized propensity score through conditional normalizing flows in order to detect regions with limited overlap in the shared covariate-treatment space (Sec. 4.2). (3) We find the optimal individualized dosage combinations (Sec. 4.3). To achieve that, we present a gradientdescent-ascent-based learning algorithm that solves the underlying constrained optimization problem (Sec. 4.4), where we avoid regions with limited overlap and thus ensure a reliable estimation of the policy value. For the latter step, we use an additional policy neural network to predict the optimal policy for fast inference during deployment.

Target

Method

1.

Estimated dose-response surface

2.

Estimated generalized propensity score

3.

Reliable policy learning

Nuisance functions estimation

Representation

network

MLP

CNF

Legend

Network,

parametrized by

Inputs / outputs with

corresponding distributions

Connections

Prediction head

parameters of network

True optimal dosage

Unreliably predicted dosage

Trajectory of reliable policy

learning

True dose-response surface

4.1 Neural network for estimating the joint effect of multiple dependent dosages

In the first step, we aim to estimate the individualized dose-response function µ ( t, x ) , which is later used as a plug-in estimator for the policy value V ( π ) . However, using a naïve neural network with t and x concatenated as inputs can lead to diminishing effect estimates of t when x is high-dimensional [68]. In addition, in our setting with multiple continuous treatments, we want to yield smooth estimates of the dose-response surface to ensure better generalization when interpolating between sparse observations. To address this, we thus develop a tailored dosage combination network (DCNet) in the following.

Architecture of DCNet: Our DCNet builds upon the varying coefficient network (VCNet) [56] but extends it to learn the joint effect of multiple dependent dosages. DCNet consists of (i) a representation network ϕ ( x ) and (ii) a prediction head h η ( t ) ( ϕ ) , where h is a network with d η different parameters η ( t ) = ( η 1 ( t ) , . . . , η d η ( t ) ) T . The parameters η ( t ) are not fixed but depend on the dosage t . This is crucial to ensure that the influence of t does not diminish in a setting with high-dimensional covariates x .

To handle multiple dosages ( t (1) , . . . , t ( p ) , a naïve way [2] would be to incorporate a separate prediction head for each treatment with corresponding parameters η ( i ) ( t ( i ) ) . However, this would be ineffective for our task, as it captures only the marginal effect of a dosage but not the joint effect of the dosage combination. Instead, in order to learn the joint effect, we must not only model the marginal effect of a dosage (as in [2]) but also drug-drug interactions. We achieve that through a custom prediction head.

Prediction head in DCNet: In DCNet, we use a tailored prediction head, which is designed to model the joint effect of dosage combinations. For this, we leverage a tensor product basis (see, e.g., [48, 81]) to estimate smooth interaction effects in the dose-response surface. As a result, we can directly incorporate all p dosages simultaneously into the parametrization of the network h by using only one head containing the joint dosage information. This is unlike other networks that use p different independent heads [2] and thus cannot capture the joint effect but only marginal effects.

We define each scalar parameter η j ( t ) of our prediction head as

where K 1 , . . . , K p are the number of basis functions ψ ( i ) k i of the p dosages t (1) , . . . , t ( p ) , and β j,... are the trainable coefficients for each element of the tensor product. In order to preserve the continuity of the non-linear dose-response surface, we choose ψ ( i ) k i to be polynomial spline basis functions. Hence, we yield the parameter vector η ( t ) = BΨ ( t ) , where

with matrix rows arranged as ( β j, 11 ... 1 , . . . , β j,K 1 1 ... 1 , . . . , β j,K 1 ...K p -1 1 , . . . , β j,K 1 ...K p -1 K p ) , with

and with ψ ( i ) = ( ψ ( i ) 1 , . . . , ψ ( i ) K i ) T , where ⊗ is the Kronecker product. For training DCNet, we minimize the mean squared error loss L µ = 1 n ∑ n i =1 (ˆ µ ( t i , x i ) -y i ) 2 .

Our prediction head has three key benefits: (1) When estimating µ ( t, x ) , the tensor product basis enables us to model non-linear interaction effects among different dosage dimensions while still having a dose-response surface that offers continuity, smoothness, and expressiveness. (2) Having a dose-response surface that is smooth between different dosages is beneficial for our task as it allows for better interpolation between sparse data points. (3) Smoothness is also exploited by our gradient-based learning algorithm later. The reason is that a smooth plug-in estimator ˆ µ ( t, x ) for the estimated policy value ˆ V ( π ) allows us to have more stable gradients and thus more stable parameter updates during learning.

4.2 Conditional normalizing flows for detecting limited overlap

In the second step, we estimate the GPS f ( t, x ) in order to detect limited overlap, which later allows us to avoid regions with high uncertainty and thus to yield a reliable policy.

Why is the GPS relevant for reliable decision-making? Due to sparsity in the data or otherwise limited overlap, the estimates of the individualized dose-response function ˆ µ ( t, x ) from our DCNet may be naturally characterized by regions with high uncertainty, that is, regions in the covariate-treatment space with low GPS. If we would use the estimated ˆ µ ( t, x ) in a naïve manner as a plug-in estimator for policy learning, we may yield a policy where there can be a large discrepancy between the predicted outcome ˆ µ ( t, x ) and the true outcome µ ( t, x ) and, therefore, a large discrepancy between ˆ V ( π ) and V ( π ) . By avoiding regions with a low GPS, we can thus ensure reliable decision-making in safety-critical applications such as medicine.

Estimation of GPS: The GPS (i.e., the logging policy) is typically not known in real-world settings. Instead, we must estimate the GPS from observational data. To this end, we leverage conditional

normalizing flows (CNFs) [75, 80]. CNFs are a fully-parametric generative model built on top of normalizing flows [60, 71] that can model conditional densities p ( y | x ) by transforming a simple base density p ( z ) through an invertible transformation with parameters γ ( x ) that depend on the input x . 2

In our method, we use CNFs with parameters γ ( x ) to estimate the probability density function f ( t, x ) = f T | X = x ( t ) of T conditioned on the patient covariates X = x . For our method, CNFs have several advantages. 3 First, CNFs are universal density approximators [11, 12, 15, 27], making them flexible and able to capture even complex density functions for dosage combinations. Second, CNFs are properly normalized. Third, unlike kernel methods, CNFs are fully parametric, and, once trained, the inference time is constant [52]. This is particularly important later in the third step of our method (Sec. 4.3): For each training step and training sample in the batch, the predicted conditional density has to be evaluated again under the updated policy, and, hence, a fast inference time of density is crucial to achieving a significant speed up when training the policy network.

In our method, we use neural spline flows [15], as they allow for flexible estimation of multimodal densities, and, in combination with masked auto-regressive networks [13], we can extend them to multiple dimensions. This combination allows for fast and exact evaluation of the multidimensional densities. We also set the base density, p ( z ) , to standard normal distribution. The CNFs are trained by minimizing the negative log-likelihood loss L f = -1 n ∑ n i =1 log ˆ f ( t i , x i ) .

In the following section, we use the estimated GPS ˆ f ( t, x ) during policy learning in order to avoid regions with limited overlap and thus ensure reliable decision-making.

4.3 Constrained optimization to avoid regions with limited overlap

Reliable decision-making : Our aim is to find the policy π rel which maximizes the estimated policy value ˆ V ( π ) while ensuring that V ( π ) can be estimated reliably . For this, we suggest to constrain our search space in the shared covariate-treatment space to regions that are sufficiently supported by data, i.e., where overlap is not violated. Formally, we rewrite our objective from Eq. (1) as

where ε is a reliability threshold controlling the minimum overlap required for our policy. In our setting with finite observational data, this yields the constrained optimization problem

Here, we use our DCNet as a plug-in dose-response estimator ˆ µ ( t, x ) , and we further use our conditional normalizing flows to restrict the search space to policies with a sufficiently large estimated GPS ˆ f ( t, x ) .

Neural network for reliable policy learning: To learn the optimal reliable policy, we further suggest an efficient procedure based on gradient updates in the following. Specifically, we train an additional policy neural network to predict the optimal policy. Our choice has direct benefits in practice. Why not just optimize over the dosage combinations for each patient individually? In principle, one could also use methods for constrained optimization and leverage the nuisance functions ˆ µ and ˆ f to directly optimize over the dosage assignments t for each patient, i.e., t rel i = arg max ˜ t i [ ˆ µ ( ˜ t i , x i )] s.t. ˆ f ( ˜ t i , x i ) ≥ ε . However, this does not warrant scalability at inference time, as it requires that one solves the optimization problem for each incoming patient. The complexity of this optimization problem scales with both the number of dosages and the number of samples in the data to be evaluated. This makes treatment recommendations costly, which prohibits real-time

decision-making in critical care. As such, direct methods for optimization are rendered impractical for real-world medical applications. Instead, we suggest to directly learn a policy that predicts the optimal dosage combination for incoming patients.

As the constrained optimization problem in Eq. (6) cannot be natively learned by gradient updates, we must first transform Eq. (6) into an unconstrained Lagrangian problem

where π θ ( x i ) is the policy learner with parameters θ , and λ i are the Lagrange multipliers for sample i . The Lagrangian min-max-objective can be solved by adversarial learning using gradient descentascent optimization (see, e.g., [47] for a background). Details are presented in the next section where we introduce our learning algorithm.

4.4 Learning algorithm

The learning algorithm for our method proceeds along the following steps (see Algorithm 1):

(1) We use our DCNet to estimate the individualized dose-response function ˆ µ ( t i , x i ) for dosage combinations.

(2) We estimate the GPS ˆ f ( t i , x i ) using the CNFs to detect areas with limited overlap.

(3) We plug the estimated ˆ µ ( t, x ) and ˆ f ( t, x ) into the Lagrangian objective from Eq. (7). As the policy learner, we use a multilayer-perceptron neural network ˆ π θ ( x ) to predict the optimal dosage combination t given the patient covariates x . We use stochastic gradient descent-ascent. In detail, we perform adversarial learning with alternating parameter updates. First, in every training step, we perform a gradient descent step to update the parameters θ of the policy network ˆ π θ with respect to the policy loss

where ε is the reliability threshold and where λ = { λ i } n i =1 are trainable penalty parameters. Then, we perform a gradient ascent step to update the penalty parameters λ with respect to L π . Here, each λ i is optimized for each sample i in the training set.

As there is no guarantee for global convergence in the non-convex optimization problem, we perform k = 1 , . . . , K random restarts and select the best run as evaluated on the validation set by

where π ( k ) θ is the learned policy in run k and ✶ {·} denotes the indicator function. As a result, we select the policy π ( j ) θ which maximizes ˆ V ( π ) under the overlap constraint ˆ f ( π ( k ) θ ( x i ) , x i ) ≥ ε, ∀ i . Note that we use the validation loss from Eq. (9) because L π cannot be directly applied for evaluating the performance on the validation set, as λ i can only be optimized for observations in the training set during the policy learning.

5 Experiments

We perform extensive experiments using semi-synthetic data from real-world medical settings to evaluate the effectiveness of our method. Semi-synthetic data are commonly used to evaluate causal inference methods as it ensures that the causal ground truth is available [8, 83] and thus allows us to benchmark the performance.

MIMIC-IV: MIMIC-IV [35] is a state-of-the-art dataset with de-identified health records from patients admitted to intensive care units. Analogous to [67], we aim to learn optimal configurations of mechanical ventilation to maximize patient survival ( Y ). Here, treatments ( T ) involve the configuration of mechanical ventilation in terms of respiratory rate and tidal volume. When filtering for T , we yield a dataset with n = 5476 patients. We select 33 variables as patient covariates X (e.g., age, sex, respiratory measurements, cardiac measurements). Unlike [67] where treatments T are modeled as p mutually exclusive dosages, we model the more realistic setting, where T are dosage combinations with a joint effect. Our choice is well aligned with medical literature according to which the parameters of mechanical ventilation are dependent and must be simultaneously controlled [20]. Further details for MIMIC-IV are in Appendix B.

TCGA: TCGA [79] is a diverse collection of gene expression data from patients with different cancer types. Our aim is to learn the optimal dosage combination for chemotherapy ( T ). In chemotherapy, combining different drugs is the leading clinical option for treating cancer [82]. Let patient survival denote the outcome ( Y ). We select the same version of the dataset as in [5, 67], containing the measurement of 4,000 genes as features X from n = 9659 patients. Unlike previous works [5, 67], we again model T as different, simultaneously assigned dosages with a joint effect. This follows medical practice where drug combinations in cancer therapy must be jointly optimized for each patient profile [16, 82]. Given that TCGA is high-dimensional, we expect it to be an especially challenging setting for evaluation. Further details for TCGA are in Appendix B.

In both datasets, the response Y ∼ N ( µ ( T, X ) , 0 . 5) can be interpreted as the log-odds ratio of patient survival, which is to be maximized. Due to the fact that we work with actual data from medical prac- tice, certain dosage combinations conditional on the patients' features are more likely to be observed. Even other dosage combinations may not be assigned at all as they can cause harm. For our experiments, we introduce a parameter α controlling the dosage bias, which changes the observed policy of dosage assignments from random uniform ( α = 0 ) to completely deterministic ( α = ∞ ). Hence, with increasing α , the areas in the covariate-treatment space will be increasingly prone to limited overlap. We refer to Appendix B for details. When not stated otherwise, we choose α = 2 and p = 2 as the default setting for our experiments.

Algorithm 1: Reliable off-policy learning for dosage combinations

Baselines: To the best of our knowledge, there are no existing methods tailored to off-policy learning for dosage combinations under limited overlap. Hence, we compare our method with the following baselines, which vary different steps in our method and thus allow us to attribute the source of performance gain. Specifically, we vary dose-response estimation ( ˆ µ ( t, x ) ) vs. optimization ( ˆ π ). (1) For estimating ˆ µ ( t, x ) , we use two baselines: (i) a simple multi-layer perceptron ( MLP ) as a standard neural alternative to our DCNet, and (ii) VCNet [56] as the SOTA method for doseresponse estimation for independent treatments with a separate head per treatment dimension. (2) For optimization, we minimize the policy loss -1 n ∑ n i =1 ˆ µ (ˆ π θ ( x i ) , x i ) without constraints for dealing with limited overlap ( naïve ).

Evaluation: For training, we split the data into train / val / test sets (64 / 16 / 20%). We first train the nuisance models ˆ µ ( t, x ) and ˆ f ( t, x ) . We then perform k = 5 random restarts to train ˆ π ( t, x ) . We select the best run with respect to the policy loss on the factual validation dataset. Thereby, we carefully adhere to the characteristics in real-world settings, where one does not have access to the true dose-response function µ ( t, x ) on the validation set for model selection. In our method, we set

the reliability threshold ¯ ε to the 5% -quantile of the estimated GPS ˆ f ( t, x ) of the train set. Further details including hyperparameter tuning are in Appendix E.

For evaluation, we compare the methods on the test set using the regret R (ˆ π ) = max π ∈ Ω V ( π ) -V (ˆ π ) . We report: (1) the regret for the selected policy, which is the policy with the best policy loss on the validation set out of all k policies; (2) the average regret across all k restarts; and (3) the standard deviation to better assess the reliability of the method.

Results: Table 1 compares the performance of our method against the baselines. We make the following observations: (1) DCNet performs similarly to MLP for low-dimensional data such as MIMIC-IV, as expected. However, we find substantial improvements from DCNet over MLP for high-dimensional data such as TCGA, showing the advantage of enforc-

Reported: best / mean / standard deviation (lower = better)

Table 1: Performance against baselines. Regret on test set over k = 5 restarts.

ing smoothness and expressiveness in our prediction head. (2) DCNet outperforms VCNet on both datasets, which demonstrates the importance of modeling the joint effect of dosage combinations. (3) Choosing a policy that is reliable greatly benefits overall performance. For example, the baselines ( + naïve) even perform worse than the observed policy for TCGA, whereas our method ( + reliable) is clearly superior. (4) We also observe that our selection of reliable policies greatly reduces the standard deviation of the observed regrets. This further demonstrates the reliability of our proposed method. Overall, our method works best across both datasets. For example, the regret on the TCGA dataset for the selected policy drops from 2.52 (DCNet + naïve) to only 0.03 (DCNet + reliable), which is a reduction by 98.8%. In sum, this demonstrates the effectiveness of our method.

Comparison to discretization: There exist several methods for off-policy learning for discrete treatments (see Sec. 2), including methods targeted for overlap violations and treatment combinations. In principle, we could also discretize the dosages in our setting and apply such methods. However, we would then lose information about the exact dosage values and could not exploit the continuity of the dose-response surface. To demonstrate the benefits of our method, we benchmark against an 'oracle' discretized baseline where we assume perfect knowledge of the dose-response function and perfect individual treatment assignment.

This serves as an upper bound for all policy learning methods with discretization. We display the results for different granularities of equally-spaced grid discretization in Table 2. We observe that, even in the oracle setting, discretization itself leads to a higher regret than our proposed method, which confirms the advantages of leveraging continuous dosage information.

Robustness checks: We evaluate the robustness of our method across settings with different levels of limited overlap. For this, we vary the dosage bias α

Reported: best / mean / standard deviation (lower = better)

(in Fig. 2) and the number of dosages p (in Fig. 3). Here, we use DCNet for modeling ˆ µ ( t, x ) in both methods, but then vary the policy selection (naïve vs. reliable, as defined above). Our method (DCNet+reliable) shows robust behavior. As desired, it achieves a low regret and a low variation in all settings. In contrast, the naïve method leads to high variation. This must be solely attributed to that our constrained optimization for avoiding regions with limited overlap leads to large performance gains. We also observe another disadvantage of the naïve method: it may select a suboptimal policy out of the different runs due to the fact that the estimation of ˆ µ ( t, x ) is erroneous in regions with limited overlap. Again, this underlines the benefits of our method for reliable off-policy learning of dosage combinations in settings with limited overlap. In Appendix C, we demonstrate the robustness of our method in further setups.

Shown: selected policy (line) and the range over 5 runs (area). The naïve baseline has a large variability across runs while ours is highly robust as expected (i.e., we see no variability).

(a) MIMIC-IV

(b) TCGA

3.0

Observed

Naive

2.5

Reliable (ours)

3

2.0

Regret

Regret

Observed

1.5

Naive

2

Reliable (ours)

1.0

0.5

D

0.0

Q

2

3

4

0

2

3

4

Dosage bias

Dosage bias

(a) MIMIC-IV

(b) TCGA

3.0

-#·

Observed

3.5

Observed

Naive

Naive

2.5

Reliable (ours)

3.0

Reliable (ours)

2.0

2.5

Regret

Regret

2.0

1.5

1.5

1.0

1.0

0.5

0.5

0.0

0.0

2

1

2

3

#Dosage dimensions

#Dosage dimensions

6 Discussion

Limitations: Our method makes an important contribution over existing literature in personalized treatment design for dosage combinations by adjusting for the naturally occurring but non-trivial problems of drug-drug interactions and limited overlap. However, the complexity of real-world data in medical practice can limit the applicability of our method in certain ways. (i) Our method does not account for the ignorability assumption, which can result in biased estimates in the case of unobserved confounders or other missing data. (ii) We apply our method in a static treatment setting, whereas, in several clinical applications, time-series data are available, e.g., sequences of varying dosages, multiple treatment cycles, and right-censored data. Thus, future work could extend our method to account for causal effect estimates in a time-varying setting (e.g., [18, 25, 51]). (iii) Our method relies on the - for medical applications reasonable - implicit assumption of smooth dose-response surfaces. In settings with extremely unsmooth or even stepwise dose-response surfaces, other methods may be more suitable for dose-response estimation. (iv) Our method aims at reliable policy learning by avoiding areas with limited overlap to minimize potential harm. In other applications such as marketing, one may want to explicitly target unreliable regions to acquire new customers.

Broader impact: Our method can have a significant impact in the field of medicine, where reliable dosing recommendations are crucial for effective treatment. Our approach addresses the limitations of previous methods in terms of reliability and estimating joint effects of dosage combinations, and, hence, can improve decision support for clinicians. The reliability of dosage recommendations is particularly important, as incorrect dosing can have serious consequences for patients, such as adverse effects or treatment failure. By addressing this gap, our approach provides a valuable tool for decision support in medicine, empowering healthcare professionals to make more personalized and safe treatment recommendations.

Acknowledgments

SF acknowledges funding from the Swiss National Science Foundation (SNSF) via Grant 186932.

References

Susan Athey and Stefan Wager. 'Policy learning with observational data'. In: Econometrica 89.1 (2021), pp. 133-161.

Alexis Bellot, Anish Dhir, and Giulia Prando. 'Generalization bounds and algorithms for estimating conditional average treatment effect of dosage'. In: arXiv preprint arXiv:2205.14692 (2022).

Andrew Bennett and Nathan Kallus. 'Efficient policy learning from surrogate-loss classification reductions'. In: ICML . 2020.

Andrew Bennett, Nathan Kallus, and Tobias Schnabel. 'Deep generalized method of moments for instrumental variable analysis'. In: NeurIPS . 2019.

Ioana Bica, James Jordon, and Mihaela van der Schaar. 'Estimating the effects of continuousvalued interventions using generative adversarial networks'. In: NeurIPS . 2020.

Guanhua Chen, Donglin Zeng, and Michael R. Kosorok. 'Personalized Dose Finding Using Outcome Weighted Learning'. In: Journal of the American Statistical Association 111.516 (2016), pp. 1509-1521.

W. G. Cochran and Donald B. Rubin. 'Controlling bias in observational studies: A review'. In: Sankhy¯ a: The Indian Journal of Statistics, Series A (1973), pp. 417-446.

Alicia Curth, Changhee Lee, and Mihaela van der Schaar. 'SurvITE: Learning heterogeneous treatment effects from time-to-event data'. In: NeurIPS . 2021.

Alexander D'Amour et al. 'Overlap in observational studies with high-dimensional covariates'. In: Journal of Econometrics 221.2 (2021), pp. 644-654.

Mert Demirer et al. 'Semi-Parametric Efficient Policy Learning with Continuous Actions'. In: NeurIPS . 2019.

Laurent Dinh, David Krueger, and Yoshua Bengio. 'NICE: Non-linear independent components estimation'. In: arXiv preprint arXiv:1410.8516 (2014).

Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. 'Density estimation using Real NVP'. In: International Conference on Learning Representations . 2017.

Hadi Mohaghegh Dolatabadi, Sarah Erfani, and Christopher Leckie. 'Invertible generative modeling using linear rational splines'. In: AISTATS . 2020.

Miroslav Dudík et al. 'Doubly Robust Policy Evaluation and Optimization'. In: Statistical Science 29.4 (2014).

Conor Durkan et al. 'Neural spline flows'. In: NeurIPS . 2019.

Funmilola A. Fisusi and Emmanuel O. Akala. 'Drug combinations in breast cancer therapy'. In: Pharmaceutical Nanotechnology 7.1 (2019), pp. 3-23.

Dylan J. Foster and Vasilis Syrgkanis. 'Orthogonal statistical learning'. In: arXiv preprint arXiv:1901.09036 (2019).

Dennis Frauen et al. 'Estimating average causal effects from patient trajectories'. In: AAAI . 2023.

Scott Fujimoto, David Meger, and Doina Precup. 'Off-policy deep reinforcement learning without exploration'. In: ICML . 2019.

Giacomo Grasselli et al. 'Mechanical ventilation parameters in critically ill COVID-19 patients: a scoping review'. In: Critical Care 25.1 (2021), p. 115.

Shonosuke Harada and Hisashi Kashima. 'GraphITE: Estimating individual effects of graphstructured treatments'. In: CIKM . 2021.

Negar Hassanpour and Russell Greiner. 'Counterfactual regression with importance sampling weights'. In: IJCAI . 2019.

Tobias Hatt, Daniel Tschernutter, and Stefan Feuerriegel. 'Generalizing off-policy learning under sample selection bias'. In: UAI . 2022.

James Heckman, Hidehiko. Ichimura, and Petra. Todd. 'Matching As An Econometric Evaluation Estimator: Evidence from Evaluating a Job Training Programme'. In: The Review of Economic Studies 64.4 (1997), pp. 605-654.

Konstantin Heß et al. 'Bayesian neural controlled differential equations for treatment effect estimation'. In: arXiv preprint arXiv:2310.17463 (2023).

Keisuke Hirano and Guido W. Imbens. 'The propensity score with continuous treatments'. In: Applied Bayesian Modeling and Causal Inference from Incomplete-Data Perspectives . 2004.

Chin-Wei Huang et al. 'Neural autoregressive flows'. In: International Conference on Machine Learning . 2018.

Kosuke Imai and David A. van Dyk. 'Causal inference with general treatment regimes'. In: Journal of the American Statistical Association 99.467 (2004), pp. 854-866.

G. Imbens. 'The role of the propensity score in estimating dose-response functions'. In: Biometrika 87.3 (2000), pp. 706-710.

Edward L. Ionides. 'Truncated Importance Sampling'. In: Journal of Computational and Graphical Statistics 17.2 (2008), pp. 295-311.

Andrew Jesson et al. 'Identifying causal-effect inference failure with uncertainty-aware models'. In: NeurIPS . 2020.

Andrew Jesson et al. 'Scalable sensitivity and uncertainty analysis for causal-effect estimates of continuous-valued interventions'. In: NeurIPS . 2022.

Fredrik D. Johansson, Uri Shalit, and David Sonntag. 'Learning representations for counterfactual inference'. In: ICML . 2016.

Alistair E. W. Johnson et al. 'MIMIC-III, a freely accessible critical care database'. In: Scientific Data 3.1 (2016), p. 160035.

Alistair E. W. Johnson et al. 'MIMIC-IV, a freely accessible electronic health record dataset'. In: Scientific Data 10.1 (2023), p. 1.

Jean Kaddour et al. 'Causal effect inference for structured treatments'. In: NeurIPS . 2021.

Nathan Kallus. 'Balanced policy evaluation and learning'. In: NeurIPS . 2018.

Nathan Kallus. 'Generalized optimal matching methods for causal inference'. In: Journal of Machine Learning Research (2020).

Nathan Kallus. 'More efficient policy learning via optimal retargeting'. In: Journal of the American Statistical Association 116.534 (2021), pp. 646-658.

Nathan Kallus and Miruna Oprescu. 'Robust and agnostic learning of conditional distributional treatment effects'. In: AISTATS . 2023.

Nathan Kallus and Angela Zhou. 'Policy evaluation and optimization with continuous treatments'. In: AISTATS . 2018.

Nathan Kallus and Angela Zhou. 'Residual unfairness in fair machine learning from prejudiced data'. In: ICML . 2018.

Diederik P. Kingma and Jimmy Ba. 'Adam: A method for stochastic optimization'. In: ICLR . 2015.

Akshay Krishnamurthy, John Langford, and Slivkins, Aleksandrs, Zhang, Chicheng. 'Contextual bandits with continuous actions: smoothing, zooming, and Aaapting'. In: Journal of Machine Learning Research 21.137 (2020), pp. 1-45.

Aviral Kumar et al. 'Conservative Q-learning for offline reinforcement learning'. In: NeurIPS . 2020.

Aviral Kumar et al. 'Stabilizing off-policy Q-learning via bootstrapping error reduction'. In: NeurIPS . 2019.

Tianyi Lin, Chi Jin, and Michael I. Jordan. 'On gradient descent ascent for nonconvex-concave minimax problems'. In: ICML . 2020.

Xihong Lin and Daowen Zhang. 'Inference in Generalized Additive Mixed Models by Using Smoothing Splines'. In: Journal of the Royal Statistical Society. Series B (Statistical Methodology) 61.2 (1999), pp. 381-400.

Haixu Ma, Donglin Zeng, and Yufeng Liu. 'Learning individualized treatment rules with many treatments: A supervised clustering approach using adaptive fusion'. In: NeurIPS (2022).

Haixu Ma, Donglin Zeng, and Yufeng Liu. 'Learning optimal group-structured individualized treatment rules with many treatments'. In: Journal of Machine Learning Research 24.102 (2023), pp. 1-48.

Valentyn Melnychuk, Dennis Frauen, and Stefan Feuerriegel. 'Causal transformer for estimating counterfactual outcomes'. In: ICML . 2022.

Valentyn Melnychuk, Dennis Frauen, and Stefan Feuerriegel. 'Normalizing flows for interventional density estimation'. In: ICML . 2023.

Milan Kuzmanovic, Tobias Hatt, and Stefan Feuerriegel. 'Estimating conditional average treatment effects with missing treatment information'. In: AISTATS . 2023.

Razieh Nabi, Todd McNutt, and Ilya Shpitser. 'Semiparametric causal cufficient dimension reduction of multidimensional treatments'. In: UAI . 2022.

Preetam Nandy, Marloes H. Maathuis, and Thomas S. Richardson. 'Estimating the effect of joint interventions from observational data in sparse high-dimensional settings'. In: The Annals of Statistics 45.2 (2017), pp. 647-674.

Lizhen Nie et al. 'VCNet and functional targeted regularization for learning causal effects of continuous treatments'. In: ICLR . 2021.

Xinkun Nie and Stefan Wager. 'Quasi-oracle estimation of heterogeneous treatment effects'. In: Biometrika 108.2 (2021), pp. 299-319.

Sonali Parbhoo, Stefan Bauer, and Patrick Schwab. 'NCoRE: Neural Counterfactual Representation Learning for Combinations of Treatments'. In: arXiv preprint (2021).

Zhaozhi Qian, Alicia Curth, and Mihaela van der Schaar. 'Estimating multi-cause treatment effects via single cause perturbation'. In: NeurIPS . 2021.

Danilo Rezende and Shakir Mohamed. 'Variational inference with normalizing flows'. In: ICML . 2015.

Rachel P. Riechelmann et al. 'Potential drug interactions and duplicate prescriptions among cancer patients'. In: Journal of the National Cancer Institute 99.8 (2007), pp. 592-600.

James M. Robins, Miguel A. Hernán, and Babette Brumback. 'Marginal structural models and causal inference in epidemiology'. In: Epidemiology 11.5 (2000), pp. 550-560.

Paul R Rosenbaum. 'Model-based direct adjustment'. In: Journal of the American statistical Association 82.398 (1987), pp. 387-394.

Donald B. Rubin. 'Bayesian inference for causal effects: The role of randomization'. In: Annals of Statistics 6.1 (1978), pp. 34-58.

Donald B. Rubin. 'Estimating causal effects of treatments in randomized and nonrandomized studies'. In: Journal of Educational Psychology 66.5 (1974), pp. 688-701.

Donald B. Rubin. 'On the limitations of comparative effectiveness research'. In: Statistics in Medicine 29.19 (2010), pp. 1991-1995.

Patrick Schwab et al. 'Learning counterfactual representations for estimating individual doseresponse curves'. In: AAAI . 2020, p. 5.

Uri Shalit, Fredrik D. Johansson, and David Sontag. 'Estimating individual treatment effect: Generalization bounds and algorithms'. In: ICML . 2017.

Claudia Shi, David M. Blei, and Victor Veitch. 'Adapting neural networks for the estimation of treatment effects'. In: NeurIPS . 2019.

Adith Swaminathan and Thorsten Joachims. 'Counterfactual risk minimization: learning from logged bandit feedback'. In: ICML . 2015.

Esteban G. Tabak and Eric Vanden-Eijnden. 'Density estimation by dual ascent of the loglikelihood'. In: Communications in Mathematical Sciences 8.1 (2010), pp. 217-233.

Shengpu Tang et al. 'Leveraging factored action spaces for efficient offline reinforcement learning in healthcare'. In: NeurIPS (2022).

Akira Tanimoto et al. 'Regret Minimization for Causal Inference on Large Treatment Space'. In: AISTATS . 2021.

Phillip S. Thomas and Emma Brunskill. 'Data-efficient off-Policy policy evaluation for reinforcement Learning'. In: ICML . 2016.

Brian L. Trippe and Richard E. Turner. 'Conditional density estimation with Bayesian normalising flows'. In: arXiv preprint arXiv:1802.04908 (2018).

Daniel Tschernutter, Tobias Hatt, and Stefan Feuerriegel. 'Interpretable off-policy learning via hyperbox search'. In: ICML . 2022.

Xin Wang et al. 'Generalization bounds for estimating causal effects of continuous treatments'. In: NeurIPS . 2022.

Yixin Wang and David M. Blei. 'The blessings of multiple causes'. In: Journal of the American Statistical Association 114.528 (2019), pp. 1574-1596.

John N. Weinstein et al. 'The Cancer Genome Atlas Pan-Cancer analysis project'. In: Nature Genetics 45.10 (2013), pp. 1113-1120.

Christina Winkler et al. 'Learning likelihoods with conditional normalizing flows'. In: arXiv preprint arXiv:1912.00042 (2019).

Simon N. Wood. 'Low-rank scale-invariant tensor product smooths for generalized additive mixed models'. In: Biometrics 62.4 (2006), pp. 1025-1036.

Debra Wu et al. 'Design principles of drug combinations for chemotherapy'. In: Journal of Controlled Release 323 (2020), pp. 36-46.

Liyuan Xu et al. 'Learning deep features in instrumental variable regression'. In: ICLR . 2021.

Yao Zhang, Alexis Bellot, and Mihaela van der Schaar. 'Learning overlapping representations for the estimation of individualized treatment effects'. In: AISTATS . 2020.

Zhengyuan Zhou, Susan Athey, and Stefan Wager. 'Offline Multi-Action Policy Learning: Generalization and Optimization'. In: Operations Research 71.1 (2023), pp. 148-183.

Hao Zou et al. 'Counterfactual prediction for bundle treatment'. In: NeurIPS . Vol. 33. 2020.

Hao Zou et al. 'Counterfactual prediction for outcome-oriented treatments'. In: ICML . 2022.

A Additional Background

Estimating the dose-response function under dosage bias: Estimating the dose-response function becomes challenging with increasing dosage bias. Although the unbiasedness of counterfactual estimation can be achieved through the minimization of the prediction loss on the factual data, the inherent scarcity of data in continuous treatment settings leads to increased variance in the estimation process (see, e.g., [2]), especially in regions with limited overlap. Consequently, the elevated variance hinders accurate generalization over the whole covariate-treatment space and contributes to a larger generalization error, which, in turn, impedes the reliable estimation of treatment effects. Effectively addressing this issue is crucial to improve the robustness and validity of counterfactual predictions.

In Sec. 2, we introduced various existing methods for estimating the dose-response function, e.g., [5, 26, 28, 29, 57, 67]. Besides a custom neural architecture, some methods use additional techniques to reduce the variance of counterfactual predictions for discrete or continuous treatments. Examples include reweighting [63, 77], importance sampling [22], targeted regularization [56], or balancing representations [2, 33, 67]. In our DCNet, we build on top of the neural architectures of [57, 67] and then develop a method that is capable of modeling the joint effect of dosage combinations.

In theory, we could also adapt additional techniques for variance reduction of the dose-response prediction to our setting of dosage combinations and leverage them in combination with DCNet. However, this may not be beneficial for our goal of reliable policy learning for dosage combinations under limited overlap, which is different from the goal of accurate dose-response estimation over the whole covariate-treatment space. These techniques aim to reduce the variance of the dose-response prediction in scarce regions (i.e., with limited overlap), which, as a consequence, shifts the focus away from regions with sufficient overlap. As we constrain the policy to regions with sufficient overlap anyway by using our reliable policy optimization, we aim for an adequate dose-response estimation in these regions, while regions with limited overlap can be neglected. Hence, in our method for reliable policy learning, techniques for tackling dosage bias during dose-response estimation can be considered more counterproductive than useful. We demonstrate this empirically in Appendix D.

Normalizing flows: Normalizing flows were introduced for expressive variational approximations in variational autoencoders [60, 71] and are based on the principle of transforming a simple base distribution, such as a Gaussian, through a series of invertible transformations to approximate the target data distribution. Let us denote the base distribution as p z ( z ) , where z is a latent variable, and the target data distribution as p x ( x ) , where x represents the observed data. The goal is to find a mapping f such that x = f ( z ) .

The core idea behind normalizing flows is to model the transformation f as a composition of invertible functions f = f K ◦ f K -1 ◦ . . . ◦ f 1 , where K is the number of transformations. The last transformation f K is typically modeled as an affine transformation, and others are invertible nonlinear transformations. The inverse transformation can be easily computed since each f k is invertible.

The change of variables formula allows us to compute the probability density of x in terms of z . If we assume the base distribution p z ( z ) is known and has a simple form (e.g., a Gaussian), we can compute the density of x via

where ∣ ∣ ∣ det ( ∂f ∂z )∣ ∣ ∣ is the determinant of the Jacobian matrix of the transformation f with respect to z . This determinant captures the stretching and squishing of the latent space induced by the transformation, thus allowing us to adjust the density of z to match the density of x .

To perform inference, such as density estimation or sampling, we need to compute the log-likelihood of the observed data x . Given a dataset D = { x 1 , . . . , x n } , the log-likelihood can be computed as the sum of the log-densities of each data point

Conditional normalizing flows: In our method, we leverage conditional normalizing flows (CNFs) [75, 80] to estimate the generalized propensity score (GPS). CNFs model conditional densities

p ( y | x ) by transforming a simple base density p ( z ) through an invertible transformation with parameters γ ( x ) that depend on the input x .

In our method, we use neural spline flows [15] in combination with masked auto-regressive networks [13]. Here, masked auto-regressive networks are used to model the conditional distribution of each variable in the data given the conditioning variable. The network takes as input the conditioning variable and generates each variable in the data sequentially, one at a time, while conditioning on the previously generated variables. This autoregressive property allows for efficient computation and generation of samples from the conditional distribution. For the motivation and advantages of our modeling choice, we refer to Sec. 4.2.

Offline reinforcement learning: A related literature stream covers off-policy evaluation (OPE) in an offline reinforcement learning (ORL) setting. Some works leverage value functions and/or logging policies to achieve reliable policy learning (e.g., [19, 45, 46]). The main differences to our setting are that ORL assumes a Markov decision process and sequential decision-making. In contrast, we focus on non-sequential, off-policy learning from observational data. Also, unlike ORL or standard OPE, we leverage the causal structure of treatment assignment and the dose-response function to return causal estimands. This allows us not only to learn decisions but also causal estimates for potential outcomes. This is crucial for medical practice, where physicians would like to reason over different treatments rather than blindly following a fully automated system.

B Data Description

MIMIC-IV: MIMIC-IV 4 is the latest version of the publicly available dataset of The Medical Information Mart for Intensive Care [35] and contains de-identified health records of patients admitted to intensive care units (ICUs). The dataset includes comprehensive clinical data, such as vital signs, laboratory measurements, medications, diagnoses, procedures, and more. In line with the approach in [67], our primary objective is to acquire knowledge about the optimal configurations of mechanical ventilation that can effectively maximize patient survival ( Y ). Note that some works (e.g., [5, 67]) used the previous version of the MIMIC dataset (MIMIC-III [34]), whereas we use the most current one to offer realistic, comparable, and meaningful findings.

As treatments ( T ), we select the specific settings (dosages) utilized for mechanical ventilation, encompassing parameters such as respiratory rate and tidal volume. By filtering the data for the last available measurements per patient before ventilation, we are able to compile a dataset consisting of n = 5476 patients. We carefully selected 33 patient covariates ( X ) that encompass factors of medical relevance such as age, sex, and various respiratory and cardiac measurements. For detailed identifiers of the selected covariates, we refer to our code.

While [67] adopted a framework in which treatments T were modeled as a set of p mutually exclusive dosages, we opted for a more realistic approach, considering that the configuration of mechanical ventilation involves dosage combinations with a joint effect. This decision aligns closely with existing medical literature, which emphasizes the interdependence of mechanical ventilation parameters and underscores the importance of simultaneous control [20].

TCGA: The Cancer Genome Atlas (TCGA) dataset 5 [79] represents a comprehensive and diverse collection of gene expression data derived from patients diagnosed with various types of cancer. Our research objective revolves around acquiring insights into the optimal combination of dosages for chemotherapy ( T ). In the domain of cancer treatment, the concurrent administration of different drugs has emerged as a prominent clinical approach [82]. In this context, patient survival serves as the designated outcome variable ( Y ).

We utilize the dataset as employed in previous works [5, 67]. Specifically, we use the version from [5], which includes gene expression measurements of the 4,000 genes with the highest variability, serving as the features ( X ), obtained from a cohort of n = 9659 patients. 6 Diverging from the approach of the aforementioned works, we once again model T as a combination of different dosages that are simultaneously assigned, taking into account their joint effect. This modeling choice closely aligns with medical practice, where the optimization of drug combinations for cancer therapy necessitates a holistic consideration of each patient's unique profile [16, 82].

We summarize the characteristics of the data in Table 3.

Semi-synthetic data simulation: Our data simulation process is inspired by [5]. We use the observed features x from the real-world datasets. Then, for each treatment dosage j = 1 , . . . , p , we randomly draw two parameter vectors v ( j ) 1 , v ( j ) 2 from a normal distribution N (0 , 1) and normalize them via v ( j ) 1 = v ( j ) 1 / || v ( j ) 1 || and v ( j ) 2 = v ( j ) 2 / || v ( j ) 2 || , respectively. We define η ( j ) = v ( j ) 1 T x 2 v ( j ) 2 T x and

˜ t ( j ) = ( 20 + 20 exp( -η ( j ) ) ) -1 + 0 . 2 . We then sample the j -th dosage of the assigned dosage combination t via

with dosage bias α = 2 , if not specified otherwise. Hence, t ( j ) ∈ [0 , 1] can be interpreted as a percentage or normalized version of the modeled dosages.

We define the dose-response function as

where κ = 1 is an interaction parameter indicating the strength of the joint effect of the different dosages. Finally, we sample the observed outcome from y ∼ N ( µ ( t, x ) , 0 . 5) . In our data simulation, we ensure that ˜ t = ( ˜ t ( j ) ) p j =1 is the optimal dosage combination and that ˜ t additionally is the mode of the joint conditional dosage distribution of t | x with the generalized propensity score as density given by f ( t, x ) = ∏ p j =1 Beta (˜ α, q ( j ) ) .

Taken together, we adapt to a realistic setting in medicine where the dosage assignment is already informed to a certain degree but still can be further optimized. With increasing α , the probability mass around the optimal dosage combinations increases, while, in areas further away from the optimum, overlap becomes more limited, and, hence, the estimation uncertainty of the dose-response function increases. We can interpolate between a random uniform distribution for α = 0 and a completely deterministic dosage assignment for α = ∞ .

C Additional Results

Robustness checks: We now repeat the robustness checks from Sec. 5 also for the MLP baselines. The results are shown in Fig. 4 and Fig. 5. We observe that our method (DCNet+reliable) performs constantly well with respect to low regret and low (to no) variability. This holds true across the different settings. Also, we see that the baseline MLP clearly benefits from our reliable optimization during policy learning. Hence, the constrained optimization that we formulated in Steps 2 and 3 of our method are also a source of gain for existing baselines.

(a) MIMIC-IV

(b) TCGA

4.0

Observed

Observed

MLP + naive

3.5

MLP + naive

MLP + reliable

MLP + reliable

DCNet + naive

3.0

DCNet + naive

DCNet + reliable

DCNet + reliable

3

2.5

Regret

Regret

2.0

2

1.5

1.0

0.5

0

0.0

0

2

3

4

0

2

3

4

Dosagebias

Dosagebias

(a) MIMIC-IV

(b) TCGA

Observed

Observed

3.5

MLP + naive

3.5

MLP + naive

MLP + reliable

MLP + reliable

3.0

DCNet + naive

3.0

DCNet + naive

DCNet + reliable

DCNet + reliable

2.5

2.5

Regret

Regret

2.0

2.0

1.5

1.5

1.0

1.0

0.5

0.5

0.0

0.0

2

3

2

3

#Dosage dimensions

#Dosage dimensions

High-dimensional dosage setting: We perform additional experiments to examine how our method scales to treatments of much higher dimensions. To improve the scalability of our method, we apply multiple prediction heads, each taking 3 treatment dimensions as input. Thereby, we fix the maximum number of jointly interacting dimensions to 3. Our motivation is that applying our method without adjustments to high treatment dimensions would lead to a high computational complexity of our prediction head because we would need to train weights for all possible interactions in the tensor product bases. However, it is often reasonable for real-world medical settings that the effective number of dimensions with joint effects is usually lower-dimensional, which is thus captured by our tailored approach. We report the results in Fig. 6. We observe that, unlike the other baselines, our method (DCNet+reliable) stays robust with low regret and low variation even for high dimensions.

Performance across multiple re-runs: To better assess the robustness of our method, we re-run our experiments for multiple train/val/test splits. Here, we also retrain the nuisance functions ˆ µ and ˆ f for each split. We display the aggregated results for (i) data generated by our standard semi-synthetic

(a) MIMIC-IV

(b) TCGA

5

Observed

Observed

MLP + naive

MLP + naive

MLP+reliable

MLP +reliable

DCNet+naive

DCNet+naive

Regret

3

DCNet+reliable

Regret

3

DCNet+reliable

2

2

2

3

6

9

15

2

3

6

6

15

#Dosage dimensions

#Dosage dimensions

data simulation (Appendix B) in Table 4, and (ii) a more complex setting with a multimodal GPS in Table 5. For (ii), we replace the dosage combination assignment from Eq. (12) with a multimodal mixture density, i.e.,

with

Reported: selected / mean / standard deviation (lower = better)

Reported: selected / mean / standard deviation (lower = better)

We can interpret the results as follows:

(1) The column 'selected' displays the final performance of our method when using the finally recommended policy. In this column, the mean ± std can be considered as the expected regret with Monte-Carlo-CV confidence intervals (including retraining the nuisance functions) and should primarily be used for evaluating the performance and for comparison with different policy learning methods. We observe that our method significantly outperforms the baselines.

(2) The column 'mean' represents the means within the k restarts of the policy learning (i.e., Step 3 in our method). As such, its values can indicate how a randomly selected policy out of the k restarts would perform compared to the 'selected one'. For example, we observe that, for the 'naïve' baselines using a naïve MSE loss for selection, suboptimal policies are oftentimes selected, whereas our method shows robust behavior.

(3) The column 'std' shows the standard deviation within the k restarts of the policy learning (i.e., Step 3 in our method). As such, the values indicate how often similar performing policies are learned. It thus refers to the robustness within the k runs. As a consequence, lower values imply that (i) probably fewer restarts k are needed to find the targeted optimum, and (ii) fewer local optima are learned, which reduces the risk of selecting a suboptimal policy.

Hence, we suggest using the 'selected' column ( mean ± std ) for evaluating the final performance of our method, and the other two columns ('mean' and 'std') to show the inner robustness of our method compared to ablation baselines.

Overall, we demonstrate that the performance of our framework is robust across the tested settings and has low regret and low variation across different re-reruns and different restarts, including the more complex setup.

Sensitivity analysis of the reliability threshold ¯ ε : In our experiments, we set the reliability threshold ¯ ε to the 5% -quantile of the estimated GPS ˆ f ( t, x ) of the train set. This is due to the reason that the GPS is a probability density function over a p -dimensional continuous space and, hence, is hard to interpret intuitively. Unlike discrete probabilities, which are bounded between 0 and 1 and can be interpreted intuitively by humans, the GPS has no upper bound. To this end, we suggest a heuristic to choose ¯ ε dependent on x % -quantiles of the estimated GPS based on the train dataset. As a result, we aim to learn optimal dosages, which have at least the same estimated reliability as x % of the observed data.

In our heuristic, the selection of the quantile x is still subjective and cannot be optimized during hyperparameter tuning, as we assume to have no access to the ground truth µ ( t, x ) . Nevertheless, we can leverage our semi-synthetic data to perform a sensitivity analysis over x for ¯ ε and thus corroborate our choice. For this, we evaluate the final performance of our method on the test data for ¯ ε set to different quantiles of ˆ f over k = 5 runs. The results are displayed in Fig. 7.

We observe that, depending on the dataset, different quantiles for selecting ¯ ε can lead to differences in variability. However, we find that lower values tend to lead to less variability. Also, we find that, even when there is higher variability within the k runs, for every setting, a run with low regret is selected. We attribute this to our validation loss from Eq. 9. Together, the results are in line with our expectations and add to the robustness of our methods.

(a) MIMIC-IV

(b) TCGA

2.5

2.5

2.0

2.0

1.5

1.5

Regret

Regret

1.0

1.0

0.5

0.5

0.0

0.0

0.0

0.2

0.4

0.6

0.8

0.0

0.2

0.4

0.6

0.8

Quantileepsilon

Quantileepsilon

D Insights

Individualized dose-response estimates: In our experiments in Sec. 5 and in Appendix C, we showed a comparison to naïve baselines. Therein, in particular, we mainly demonstrated that the reliable optimization step in our method is an important source of performance gain. We now give additional intuition as to why also our DCNet is a source of performance gain. In the following, we first demonstrate that our DCNet is able to estimate the dose-response function more accurately than the baseline in regions with sufficient overlap. We then discuss why this brings large gains for our reliable optimization. Thereby, we offer an explanation why both DCNet and reliable policy optimization are highly effective when used in combination.

In Table 6, we display the performance of the dose-response estimators. In detail, we report: (i) the mean squared error ('Biased MSE'), calculated on a holdout test dataset which follows the same biased dosage assignment as the train data; (ii) the mean ('MISE', see e.g., [67]) of the integrated squared error calculated on areas with sufficient overlap, i.e. ∑ n i =1 ∫ [0 , 1] p ( µ ( t, x i ) -

ˆ µ ( t, x i )) 2 dt | ˆ f ( t, x ) > ε .; and (iii) the standard deviation of the integrated squared error on areas with sufficient overlap ('SD-ISE'). Intuitively, the MISE and SD-ISE are the expected value and variation of the prediction error of the dose-response estimator over the covariate-treatment space with sufficient overlap. Hence, a low value implies that all suitable dosage combinations with sufficient overlap are estimated properly and that certain, less frequent or non-dominant combinations are not neglected. We observe that our DCNet also performs clearly best wrt. to all metrics, undermining its ability to successfully account for the different drug-drug interactions under dosage bias.

To give further intuition to the above, we select one random observation per dataset and plot the following functions: (i) the estimated GPS, (ii) the oracle individualized dose-response surface, (iii) the estimates using DCNet, and (iv) the estimates using the baseline MLP. We compare two settings, namely with no dosage bias ( α = 0 ) and thus sufficient overlap (in Fig. 8) and with dosage bias ( α = 2 ) and thus clearly limited overlap (in Fig. 9).

The plots can be interpreted as follows. With no dosage bias (Fig. 8), both DCNet and MLP seem to approximate µ adequately (with slightly better estimates of DCNet at the modes at the boundary areas) and are thus suitable for naïve policy learning. However, the findings change in a setting with limited overlap (Fig. 9). Here, the limited overlap is indicated by areas with low estimated GPS ˆ f . Both estimators for ˆ µ show clear deviations to the oracle µ . However, the deviations are smaller for large parts in the case of our DCNet (as compared to the MLP). Evidently, the dose-response surface predicted by the MLP loses its expressiveness, especially for the high-dimensional TCGA dataset, and fails to model the maximum adequately. This leads to erroneous predictions when used as a plug-in estimator for policy learning, which can not be compensated by constraining the search space to regions with sufficient overlap. DCNet, on the other hand, is able to model the dose-response surface more adequately in areas with sufficient overlap due to its custom prediction head ensuring expressiveness. However, in regions with limited overlap, the splines of the prediction head extrapolate and lead to strongly incorrect predictions. Anyway, as our goal is to learn a reliable policy, we aim to exclude the regions with limited overlap and thus high uncertainty (yellow regions) such that they do not affect our final predicted policy.

In conclusion, when now applying our reliable policy optimization on top of our DCNet, we constrain the search space to the reliable regions, in which DCNet outperforms the MLP. Thus, we ensure reliable and accurate predictions for policy learning by using our method.

(a) MIMIC-IV

(b) TCGA

Estimated GPSf

Oracleμ

DCNet

MLP

1.0

4.0

3.5

-3.0

Dosage 2

2.5

0.5

2.0

1.5

1.0

0.5

0.0

-0.0

0.0

0.2

0.4

0.6

0.8

1.0 0.0

0.2

0.40.6

0.8

1.0 0.0

0.2

0.4

0.6

0.8

1.00.0

0.2

0.4

0.6

0.8

1.0

Dosage 1

Dosage 1

Dosage 1

Dosage 1

Estimated GPSf

Oracleμ

DCNet

MLP

3.84

1.0

3.44

3.04

Dosage 2

2.64

2.24

0.5

1.84

1.44

1.04

0.0

0.64

0.0

0.2

0.4

0.6

0.8

1.0 0.0

0.2

0.4

0.6

0.8

1.0 0.0

0.2

0.4

0.6

0.8

1.00.0

0.2

0.4

0.6

0.8

1.0

0.24

Dosage 1

Dosage 1

Dosage 1

Dosage 1

(a) MIMIC-IV

(b) TCGA

Estimated GPSf

Oracleμ

DCNet

MLP

1.0

4.0

3.5

-3.0

Dosage 2

2.5

0.5

2.0

1.5

1.0

0.5

0.0

-0.0

0.0

0.2

0.40.6

0.8

1.0 0.0

0.2

0.40.6

0.8

1.0 0.0

0.2

0.4

0.6

0.8

1.00.0

0.2

0.4

0.6

0.8

1.0

Dosage 1

Dosage 1

Dosage 1

Dosage 1

Estimated GPSf

Oracleμ

DCNet

MLP

3.84

1.0

3.44

3.04

Dosage 2

2.64

2.24

0.5

1.84

1.44

1.04

0.0

0.64

0.0

0.2

0.4

0.6

0.8

1.0 0.0

0.2

0.4

0.6

0.8

1.0 0.0

0.2

0.40.6

0.8

1.00.0

0.2

0.4

0.6

0.8

1.0

0.24

Dosage 1

Dosage 1

Dosage 1

Dosage 1

Computational complexity of policy optimization: To ensure scalability at inference time, we train an additional policy neural network to directly predict the optimal policy, instead of performing optimization for new patients. We explain our motivation in the following: The complexity of a direct optimization at inference time is non-trivial and depends on the complexity of the forward pass of DCNet σ µ and the CNFs σ f , as well as the selected constrained optimization algorithm, and the number of patients to evaluate n eval . While the exact complexity of the direct optimization depends on the specific solver cost σ solve , which, for most algorithms, scales heavily with the dosage dimension p , it is at least O ( n eval × ( σ µ + σ f ) × σ solve ) . This becomes costly for large-scale datasets in practice. In contrast, our policy network is trained only once. Hence, upon deployment, it simply needs operations for evaluation, which is at most O ( n eval × σ net ) , where σ net is the complexity of only the forward pass of the policy network. Hence, it usually holds σ net ≪ σ solve .

E Modeling and Training Details

Dose-response estimator: For training our DCNet ˆ µ ( t, x ) , we follow previous works from the baseline VCNet [2, 56]. We set the representation network ϕ to a multi-layer perceptron (MLP) with two hidden layers with 50 hidden neurons each and ReLU activation function. For the parameters of the prediction head h , we use the same model choices as for ϕ . Additionally, we use B-splines with degree two and place the knots of the tensor product basis at { 1 / 3 , 2 / 3 } p . For the baseline MLP from Sec. 5, we ensure similar flexibility for fair comparison, and we thus select a MLP with four hidden layers with 50 hidden units each and ReLU activation. We train the networks minimizing the mean squared error (MSE) loss L µ = 1 n ∑ n i =1 (ˆ µ ( t i , x i ) -y i ) 2 . For optimization, we use the Adam optimizer [43] with batch size 1000 and train the network for a maximum of 800 epochs using early stopping with a patience of 50 on the MSE loss on the factual validation dataset.

We tune the learning rate within the search space { 0 . 0001 , 0 . 0005 , 0 . 001 , 0 . 005 , 0 . 01 } , and, for evaluation, we use the same criterion as for early stopping.

Conditional normalizing flows: For modeling the GPS ˆ f ( t, x ) , we use conditional normalizing flows (CNFs). Specifically, we use neural spline flows [15] in combination with masked auto-regressive networks [13]. We use a flow length of 1 with 5 equally spaced bins and quadratic splines. For the autoregressive network, we use a MLP with 2 hidden layers and 50 neurons each. Additionally, we use noise regularization with noise sampled from N (0 , 0 . 1) . We train the CNFs by minimizing the negative log-likelihood (NLL) loss L f = -1 n ∑ n i =1 log ˆ f ( t i , x i ) . For optimization, we use the Adam optimizer with batch size 512 and train the CNFs for a maximum of 800 epochs using early stopping with a patience of 50 on the NLL loss on the factual validation dataset.

We tune the learning rate within the search space { 0 . 0001 , 0 . 0005 , 0 . 001 , 0 . 005 , 0 . 01 } and for evaluation we use the same criterion as for early stopping.

Policy learning: After training the nuisance models ˆ µ and ˆ f , we use them as plug-in estimators for our policy learning. For the policy network ˆ π θ , we select a MLP with 2 hidden layers with 50 neurons each and ReLU activation. We set the reliability threshold ¯ ε to the 5% -quantile of the estimated GPS ˆ f ( t, x ) of the train set, when not specified otherwise. We use Adam optimizers for updating θ and λ each, with batch size 512 , and train the network using the gradient descent-ascent optimization objective from Eq. (7) for a maximum of 400 epochs using early stopping with a patience of 20 on the validation loss from Eq. (9) on the factual validation dataset.

We set the learning rate for updating λ to η λ = 0 . 01 and use random search with 10 configurations to tune the learning rate for updating the parameters of the policy network η θ ∈ { 0 . 0001 , 0 . 0005 , 0 . 001 , 0 . 005 , 0 . 01 } and the initial value of the Lagrangian multipliers λ i ∈ [1 , 5] . To evaluate the performance during hyperparameter search, we use the same criterion as for early stopping. Then, after determining the hyperparameters, we perform k = 5 runs to find the optimal policy.

Our above procedure was chosen to ensure a fair and direct comparison to previous literature and to the naïve baselines, which enables us to show the sources of performance gain in our method (see Sec. 5). Note that when applying our method for reliable policy learning for dosage combinations in production, one might consider optimizing more exhaustively also over the different modeling choices of the nuisance estimators and the policy network.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposed a novel method for off-policy learning for continuous dosage combinations that aims to be robust to action space sparsity. To account for drug-drug interactions, they proposed a dosage combination network (DCNet) which leverages a tensor product basis to estimate smooth interaction effects in the dose-response surface, incorporating all dosages simultaneously into parameterization by only one head. The DCNet is updated by minimizing the MSE between the individualized dose-response function (the expected patient outcome of a given treatment in given patient covariates) and patient outcome. Then they use conditional normalizing flows to model the stochastic and multimodal behaviour policy which they call the generalized propensity score (GPS). Finally, they learn a parametric policy in a constrained optimization problem that maximizes individualized dose-response function under the learned policy subject to the probability of the proposed treatment in the GPS being larger than a threshold. This is approximately solved as an unconstrained Lagrangian problem through stochastic gradient descent-ascent with random restarts. Experiments were conducted on two real healthcare datasets. Comparisons to ablated versions (MLP instead of DCNet, without overlap constraint) suggest that the proposed components reduce regret and improve stability.\n\nSTRENGTHS:\n1. proposed a novel method for off-policy learning for continuous dosage combinations\n2. a viable design for constrained policy optimization\n\nWEAKNESSES:\n1. dosage dimensionality in experiments is just 2 or 3, it is hard to say wether the obtained results can generalize to much higher dimensionalities, or treatments intended for different diseases simultaneously\n2. the individualized weights in the proposed DCNet are only trained in an unsupervised fashion, i.e. the training objective does not reflect individual dosages\n3. illegible fonts in figures\n\nQUESTIONS:\n1. do you have to learn an analytical form of the GPS if it is only to indicate overlap?\n2. what is the harm of assuming independent dosage dimensions, which is mentioned several times throughout the paper but never clearly explained\n3. why the single-head individualized dose-response function trained against the overall treatment outcome can account for drug-drug interactions? Is it possible that one drug dominates the others in the model training, or one common combination of dosages dominate less frequent combinations that actually entail more subtle drug-drug interactions? As I understand, the tensor product basis of the proposed DCNet is all inner work of the model, whereas the input, output and training objective are all subject to non-individualized information.\n\nLIMITATIONS:\nplease see above", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-15"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\n(I have more background in offline RL than causal inference, and my description below may include some RL terminologies as the two are closely related.) \n\nThis submission studies the problem of off-policy learning of dosage recommendations in personalized medical decision making, using a causal inference framework. The proposed approach operates in three stages: the first two stages aim to learn the dose-response function (similar to the Q-function) and the logging policy (used to identify poor overlap), whereas the third stage combines the two to learn a reliable policy using constrained optimization (max performance while avoiding poor overlap), solved via a minimax gradient optimization of the Lagrangian form of the objective. Experiments are conducted using semi-synthetic datasets for ventilation (MIMIC-IV) and gene expression (TCGA).\n\nSTRENGTHS:\n- Writing is clear with very organized descriptions of each stage of the proposed approach. \n- Well-motivated problem by real-world medical applications where one needs to determine multiple continuous treatments with dependent effects. \n- Code is provided as part of the supplement (I did not check or run the code).\n\nWEAKNESSES:\n- The contribution of this work should be further clarified. As \"several methods for off-policy learning aim at discrete treatments\" (L27), it may be worth pointing out whether solutions exist for multiple *discrete* treatments (instead of continuous dosages) and explain why those approaches don't directly apply, and how well do we expect those methods to perform for dosage combinations if we discretize the dosages. Furthermore, specific architectures are used for estimating dose-response and GPS, and I'm not entirely convinced why simpler neural networks do not work. \n- Furthermore, there are many recent works on offline RL and OPE that make use of value functions and logging policies to achieve reliable policy learning (e.g., references below, including many for continuous actions). Some discussion about how this work is related and/or different from those approaches would make the paper more complete and resonate with a wider audience. \n- In the experiments, the baselines considered may be more accurately described as ablation studies. These are certainly important for understanding the sources of performance gains, but I think it is important to include more direct baselines. Since \"there are no existing methods tailored to off-policy learning for dosage combinations under limited overlap\" (L329), how well does a standard causal effect estimation method work here? \n\nReferences:\n[1] Off-Policy Deep Reinforcement Learning without Exploration. https://arxiv.org/abs/1812.02900\n[2] Stabilizing Off-Policy Q-Learning via Bootstrapping Error Reduction. https://arxiv.org/abs/1906.00949\n[3] Conservative Q-Learning. https://arxiv.org/abs/2006.04779\n\nQUESTIONS:\n- For medical applications, how common is it that practitioners use exact dosages vs dosage levels? What's the practicality of an exact dosage specification in real-world problems? More discussion about this is appreciated. \n- L131 Why is ε = 0 considered weak overlap instead of no overlap? \n- L183 mentions “polynomial spline basis functions”, a possibly unfamiliar concept to many readers. Can you provide the precise definition? \n- Sec 4.1 is very dense, I still don't quite understand it after several read throughs. It requires a lot of thinking to understand p-dimensions, may be easier to illustrate using  a simple example with two dimensions. \n- L274 training stability: \"as there is no guarantee for global convergence in the non-convex optimization problem\", do you observe any training instability in the experiments and how much variation is there? When selecting the best run over K random restarts, what is K set to in the experiments? \n- L95, L247, Sec 4.3: a neural network policy is trained to approximate the individualized optimal dosage, and is claimed to have a computational advantage. However, no empirical results (runtime of two implementations) or analysis (big-O complexity) is shown for compute scalability. Also, is there any performance loss from using this neural network policy compared to an exact argmax? \n\nMinor:\n- Figures: fonts are too small in all figures\n- L95 \"unfeasible\" -> \"infeasible\"\n- L116 “Ω denotes the set of all possible policy classes” is confusing, I think Ω should be the set of all policies. \n- L246 I think \"max\" in the equation should be argmax\n- L289 \"when filtering for T\" -> what does this mean? \n\n---\n\n**Updates after rebuttal and discussion:**\n\nThank you authors for the response, it resolved most of my questions. I have raised my rating from 5 to 6. I would like to see the final version incorporating the updated limitation section and the details on dataset characteristics.\n\nLIMITATIONS:\nSec 6 Conclusion separately discusses limitations and broader impacts. Though the proposed approach aims to avoid unreliable treatments, I think it is still necessary to acknowledge potential failure scenarios of the proposed approach and discuss what are the possible downstream impacts, especially since this work is focused on the high-stakes decision making of treatment dose recommendation.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-26"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper deals with an important and interesting problem in personalized decision making. In particular, the paper considers the setting when there are multiple continuous treatments available and models the joint effect of dosage combinations. A novel off-policy learning algorithm is proposed to solve the problem using the well-constructed network architecture. Specifically, the paper also considers the issue when the (strong) positivity assumption of the propensity score does not hold. To solve this problem, the paper modifies the proposed algorithm to achieve a reliable estimation of the policy value by avoiding regions with limited overlap. Finally, several experiments and real data analysis are provided to demonstrate the performance of the proposed method.\n\nSTRENGTHS:\n1. The paper has a good presentation to illustrate the target problem, underlining causal framework, proposed methods, and related advantages. The authors have skillfully laid out the paper, providing a clear and concise overview of the problem at hand. The methods proposed by the authors are well-described and thoroughly explained, ensuring that readers can easily grasp the technical aspects of the approach. Furthermore, the paper showcases the advantages of the proposed methods, demonstrating how they address the limitations of existing approaches and offer novel contributions to the field. The well-designed format of presenting the problem, framework, methods, and advantages in a cohesive and coherent manner enhances the overall quality of the paper.\n\n2. The paper provides real data analysis in multiple dosage settings. This empirical analysis adds substantial value to the research as it validates the proposed methods and demonstrates their effectiveness in practical scenarios. Moreover, the multiple dosage settings explored in the analysis indicate the robustness of the proposed methods, making the findings applicable to a wider range of scenarios.\n\nWEAKNESSES:\n1. The paper addresses a specific issue related to limited overlap, which can lead to unreliable estimations of policy value. However, it is worth noting that similar and related problems have been discussed in the literature. To the best of our knowledge, the following papers have explored these problems:\n\n(1). Ma et al. (2023) Learning Optimal Group-structured Individualized Treatment Rules with Many Treatments. Journal of Machine Learning Research.\n\n(2). Ma et al. (2022) Learning Individualized Treatment Rules with Many Treatments: A Supervised Clustering Approach Using Adaptive Fusion. NeurIPS.\n\nAlthough the above papers primarily focus on settings involving discrete treatments, they have also carefully analyzed the limited overlap problem. Specifically, they examine cases where the number of treatments approaches $+\\infty$, which can be treated as one of the special cases of continuous treatment, and propose methods for identifying latent treatment grouping structures (combine treatments or combine levels of doses) to address issues of unbalanced treatment assignment and limited overlap. It would be beneficial to include further discussion and comparisons with these methods in the paper, as they may offer valuable insights and contribute to the overall understanding of the problem.\n\n2. The paper mentions an important term, $\\bar{\\epsilon}$, which acts as a pre-specified threshold controlling the minimum overlap required for the policy estimation when shrinking the treatment-covariate space. It is interesting and important to investigate and provide guidelines for selecting an appropriate value for $\\bar{\\epsilon}$ in the experimental analysis. This could involve sensitivity analysis or exploring different threshold values to evaluate their impact on the estimated policies and their performance metrics. Assessing the robustness of the proposed method to the choice of $\\bar{\\epsilon}$ would provide insights into its practical applicability and generalizability across various scenarios.\n\nQUESTIONS:\n1. In contrast to the approaches mentioned in the provided references, which focus on combining similar treatments and estimating treatment structures, this paper tackles the limited overlap problem in the context of continuous dosage by constraining the search space and shrinking the entire searching space. While the methods may differ in their specific techniques, there could be interesting connections between the two approaches. Exploring and discussing these potential connections could provide valuable insights into how different strategies can address the limited overlap problem across different treatment settings. Such a discussion would contribute to a broader understanding of the problem. In particular, are there any connections between these two methods? That would be interesting to discuss.\n\n2. When reducing the searching space to address the limited overlap problem, there may be concerns about whether the estimated policy with the proposed method leads to sub-optimal values at the population level. Understanding the potential trade-offs between reducing the search space and achieving optimal population-level outcomes may be valuable for both theoretical considerations and practical applications. If reducing the searching space, will the estimated policy with proposed method leads to a sub-optimal value in the population level?\n\nLIMITATIONS:\nPlease see it in Weaknesses section.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces a new methodology to derive a policy to assign individualized dosage combinations relying on retrospectively collected observational data. The method relies on off-policy learning and try to account for the joint effect of multiple dependent dosages. A neural network is used to estimate individualized dose-response functions, accounting for the joint effect of multiple dependent dosages. The policy is then trained, avoiding regions with limited overlap to ensure reliable policy evaluation.\n\nThe developed method is then evaluated using semi-synthetic data relying on MIMIC-IV and TCGA. The robustness of the methodology is also assessed testing different levels of overlap for patients under different treatment regimens. The policy derived using the proposed method achieves low regret and low variance on all the experiments performed.\n\nSTRENGTHS:\nThe paper combines the idea of several recently published papers to derive the global method accounting for joint effects of dosage combinations. \n- A tensor product basis is used to estimate interaction effect and integrated into a model similar to [VCNet](https://arxiv.org/abs/2103.07861);\n- The generalized propensity score is estimated using conditional normalized flows, increasing reliability and avoiding discretization of the problem;\n- A NN is used for policy learning relying on the estimated dose response function and GPS;\n\nA nice protocol is then performed, assessing the performance of the model in the presence of dose combinations with a joint effect.\n\nWEAKNESSES:\nIn the optimization problem of equation (6) (line 237),  the objective function and constraints can be non convex. There is no reason when applying the method of Lagrangian multipliers to have a function that is concave in the Lagrangian multipliers $\\lambda$. The method you rely on, [gradient descent ascent](https://arxiv.org/pdf/1906.00331.pdf) explicitly requires the concavity in $\\lambda$. This can lead to sub optimality and all the nice properties of the method are lost. \n\nGiven this, we can expect that the methodology will induce a lot of variability. If part of it is captured with the random restart proposed, we still expected it to be highly unstable. The results don't reflect that with a surprising std at 0 for the policy derived, whatever the task is. \nThis might  be due to the overly simple ground truth dose response function and GPS. Elements to assess this hypothesis are provided in the Questions. This also might be due to seed tuning. It would be interesting to repeat the split into train, validation and test sets and take the variation in the best underlying regret into account.\n\nQUESTIONS:\nAs mentioned above, the results are only displayed for one split, being potentially prone to seed tuning, it would be interesting to repeat the split into train, validation and test sets and take this variation into account.\n\nGiven that you have access to the ground truth for the dose response function and GPS, could you also provide the error to estimate them in your experiments. This could provide insights on why the observed std is a 0 and motivate a more realistic dose-response function.\n\nIn the TCGA dataset and MIMIC-IV, the outcome Y is taken as the patient's survival. It's very uncommon to get a dataset with full observation of Y and the outcome is often right censored. Did you think of a way to take this into account?\n\nLIMITATIONS:\nT is the assigned p-dimensional dosage combination. In practice most of the time it's a sequence of treatment, varying dosage and number of cycles. The setting in the paper is an advancement over even more simple settings but is far from reflecting the data you can access in clinical routine.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}