An under-developing metaDTR R-package can be found on Github by the same author.
Introduction
“Dynamic treatment regimes (DTRs), also called treatment policies, adaptive interventions or adaptive treatment strategies, were created to inform the development of health-related interventions composed of sequences of individualized treatment decisions. These regimes formalize sequential individualized treatment decisions via a sequence of decision rules that map dynamically evolving patient information to a recommended treatment. An optimal DTR optimizes the expectation of a desired cumulative outcome over a population of interest. An estimated optimal DTR has the potential to:
- improve patient outcomes through adaptive personalized treatment;
- reduce treatment burden and resource demands by recommending treatment only if, when, and to whom it is needed; and
- generate new scientific hypotheses about heterogeneous treatment effects over time both across and within patients.”(Laber et al. 2014)
We consider the development of DTRs using data from Sequential, Multiple, Assignment Randomized Trials (SMART) (Lavori and Dawson 2000; Murphy 2005a; Nahum-Shani et al. 2012). See Figure 1 for a concrete example of a hypothetical SMART design for addiction management that is introduced in Chakraborty and Murphy (2014). “In this trial, each participant is randomly assigned to one of two possible initial treatments: cognitive behavioral therapy (CBT) or naltrexone (NTX). A participant is classified as a non-responder or responder to the initial treatment according to whether s/he does or does not experience more than two heavy drinking days during the next two months. A non-responder to NTX is re-randomized to one of the two subsequent treatment options: either a switch to CBT, or an augmentation of NTX with CBT (CBT+NTX). Similarly, a non-responder to CBT is re-randomized to either a switch to NTX, or an augmentation (CBT+NTX). Responders to the initial treatment receive telephone monitoring (TM) for an additional period of six months. One goal of the study might be to construct a DTR leading to a maximal mean number of non-heavy drinking days over 12 months.” (Chakraborty and Murphy 2014)
Figure 1: Hypothetical SMART design schematic for the addiction management example (an ‘R’ within a circle denotes randomization at a critical decision point)
DTRs offer a framework for operationalizing the adaptive decision making in clinical practice and thereby potentially improving it. They can be conceptualized as decision support systems for the clinicians-one of the key elements of the Chronic Care Model (Wagner et al. 2001).
Before moving to technical details of optimizing DTR, we here discuss a little bit about why not combining best treatment at each stage that are obtained from separate randomized trials. This way does not necessarily give the optimal DTR, due to:
- Delayed Therapeutic Effect
- Treatment I may not appear best initially but may have enhanced long term effectiveness when followed by a particular “maintenance” treatment (Positive Synergies)
- Treatment I may produce a higher proportion of responders but also result in side effects that reduce the variety of subsequent treatments for those that do not respond, Or the burden imposed by treatment I may be sufficiently high so that non-responders are less likely to adhere to subsequent treatments (Negative Synergies)
- Diagnostic Effect or Prescriptive Effect
- Treatment I may produce a lesser proportion of responders that treatment II, but treatment I may elicit symptoms that allow the investigator to better match subsequent treatment to the patient and thus achieve improved primary outcome to the sequence of treatments as compared to initial treatment II
- Cohort Effect or Selection Effect
- Subjects who enroll in, or who remain in or who are adherent in the trial of the initial treatments may be quite different from the subjects in SMART
In this work, we will introduce the overall structure of finding the optimal DTRs using backward induction. Moreover, we incorporate the idea of meta-learners, which are frequently used in causal inference for optimal treatment recommendation and/or heterogeneous treatment effect estimation, to facilitate the estimation of DTRs. We also evaluate the performance of the proposed meta-learner methods with the existing state-of-art methods using the same simulation setting.
Method
Notation and Problem Setup
Consider a multistage decision problem where decisions are made at \(T\). Here we use capital letters to represent random variables and lower case for observations. For \(t = 1,...,T\), let \(A_t\) denote the treatment assigned at the \(t^{th}\) stage, and \(X_t\) be the observations after treatment assignment \(A_{t-1}\) but prior to \(A_{t}\). In particular, \(X_1\) stands for the baseline covariates, while \(X_t, t>1\) are merely a realization of one “potential” set of covariates corresponding to the history treatment sequence \(a_1,...,a_{t-1}\). In other words, the covariates observed at later stage might be impacted by the past treatment assignments. Notably, \(A_j\) might be multi-category or even continuous, but to ease illustration, we mainly focus on binary treatment scenario, i.e., \(A_{t} \in \{0,1\}, t=1,...,T\) in this section, and will discuss the cases with \(K\) treatment options \(A_{t} \in \{1,...,K\}\) in later sections. See Figure 2 for a better understanding of the notation in a two-stage study design.
Figure 2: A two-stage study design
Following \(t^{th}\) treatment, there will be an observed outcome, or so-called “reward” in the reinforcement learning researches, denoted as \(Y_t\), where the larger the better in current setting. Usually, the outcome depends on the whole history information, i.e., all the covariates observed before \(t^{th}\) stage, \(X_1,...,X_t\); all the treatment assignments \(A_1,...,A_t\); and all past outcomes \(Y_1,...Y_{t-1}\). For simplicity in notation, we use \(H_t = (X_1, A_1, Y_1, ..., A_{t-1}, Y_{t-1}, X_t)\) to represent the history information available at stage \(t\).
The goal of DTR is to maximize the total rewards/outcomes: \[ Y^* = \sum_{t=1}^T w_tY_t, \] where \(w_t\) represent the pre-defined weighting factors for outcomes at each stage. Weighting helps to flexibly customize the goal for a specific question. In trial design, usually people adopt equal weights, but in finance, for example, it’s quite natural to choose \(w_t = \left(\frac{1}{1+r_f}\right)^t, t=1,...,T\) to maximize the present value by discounting the future cash flows. \(r_f\) can be the risk-free rate like 10 Year Treasure Rate.
A dynamic treatment regime (DTR) is an ordered set of deterministic decision rules, denoted as \(\mathbf g = (g_1,...,g_T)\), where \(g_t\) is a map from the space of \(H_t\) to the available assignment space \(A_t\), i.e., \(g_t: H_t \longmapsto A_t\). Let \(Y^*(\mathbf g)\) be the outcome given treatments follows the regime \(\mathbf g\), and then the optimal treatment regime \(\mathbf g^\text{opt} = (g_1^\text{opt},...,g_T^\text{opt})\) should satisfy \[\begin{equation} \mathbb E[Y^*(\mathbf g^\text{opt})] \geq \mathbb E[Y^*(\mathbf g)] , \ \ \mathbf g \in \mathcal G \tag{1} \end{equation}\] where \(\mathcal G\) is the class of all feasible regimes.
To find solution for \(\mathbf g\) satisfies equation (1), the most common way is called Q-learning, which is an indirectly method. On the other hand, people could also use directly method, namely, A-learning to seek solution. In later subsections, we will introduce the basic idea of these two methods.
Q-learning
The optimal regime may be determined via dynamic programming, also referred to as backward induction (Bellman 1966; Murphy et al. 2001; Murphy 2003). Define Q-functions, with “Q” for “quality”, as \[ \begin{aligned} Q_T(h_T, a_T) &= \mathbb E[Y_T \mid H_T=h_T, A_T=a_T], \\ Q_t(h_t, a_t) &= \mathbb E[Y_t+\max_{a_{t+1}} Q_{t+1}(h_{t+1}, a_{t+1}) \mid H_{t+1}=h_{t+1}, A_{t+1}=a_{t+1}] \end{aligned} \] for \(t = T-1,...,1\). The optimal treatment at each stage is then determined by \[\begin{equation} g_t^\text{opt}(h_t) = \arg\max_{a_t} Q_t(h_t,a_t), t = 1,...,T. \tag{2} \end{equation}\] Hence, the optimal DTR is \({\mathbf g^\text{opt}}(h_T) = (g_1^\text{opt}(h_1), ..., g_T^\text{opt}(h_T))\). In practice, we need to first obtain \(Q_t()\), then find \(Q_{t-1}()\), and repeat this procedure recursively until \(t=1\). Since optimal DTR is totally dependent on the performance of estimated \(Q_1(),...,Q_T()\), so this approach is named as “Q-learning”. In terms of how to obtain Q-functions, we introduce meta-learners to complete this task. Under the Q-learning umbrella, there are two main approaches, Single- or S-learner and Two- or T-learner.
S-learner
For simplicity, define value function \(V_t(h_t) = \max_{a_t} Q_t(h_t, a_t)\). Also, consider a target loss function \(\ell(q, \tilde q)\), which measures the “loss” between the observed quantity \(q\) and estimated quantity \(\tilde q\). For continuous \(q\), for example, the loss function \(\ell()\) could be mean squared error; for binary outcomes, it could be logit loss or simply Shannon entropy; or if \(q\) is a distribution, then \(\ell()\) can choose Kullback-Leibler divergence. In S-learner, we basically only need to use the base learner to minimize the following loss \[\begin{equation} \hat Q_t(h_t, a_t) = \arg\min_{Q_t} \ell(Y_{t+1} + \hat V_{t+1}(h_{t+1}), Q_t(h_t, a_t)) \tag{3} \end{equation}\] where \(\hat V_{t+1}(\cdot)\) is obtained in the previous step; and then we can estimate \(\hat V_t(h_t) = \max_{a_t} \hat Q_t(h_t, a_t)\). By conducting this procedure recursively, we get \((\hat Q_1,...,\hat Q_T)\) and followed by \(\hat{\mathbf g}^\text{opt}\). In practice, \(Q()\) is a known function of parameter \(\boldsymbol\Theta\) and (3) is basically solving the following optimization problem \[ \hat {\boldsymbol\Theta}_t = \arg\min_{\boldsymbol\Theta_t} \ell(Y_{t+1} + \hat V_{t+1}(h_{t+1}), Q_t(h_t, a_t; \boldsymbol\Theta_t)). \] For simplicity, we here try not to involve too much notation and therefore, we follow the structure of (3) instead hereafter.
T-learner
The only difference between S- and T-learner is, in S-learner, there is only one Q-function at each stage, while in T-learner, we estimate two Q-functions, one for \(a_t=1\) and the other for \(a_t=0\). Specifically, we now have \[ \begin{aligned} \hat Q_{t,1}(h_t) &= \arg\min_{Q_{t,1}} \ell(Y_{t+1} + \hat V_{t+1}(h_{t+1}), Q_{t,1}(h_t) \mid a_t=1) \\ \hat Q_{t,0}(h_t) &= \arg\min_{Q_{t,0}} \ell(Y_{t+1} + \hat V_{t+1}(h_{t+1}), Q_{t,0}(h_t) \mid a_t=0) \end{aligned} \] and the corresponding estimation of value function is thus \[ \hat V_t(h_t) = \max \left(\hat Q_{t,1}(h_t), \hat Q_{t,0}(h_t) \right). \] Notably, though \(Q_t(h_t, a_t)\) is similar to \(Q_{t,a_t}(h_t)\) conceptually, the former one is estimated using whole data, while the latter one only use part of data, i.e., either \(a_t=1\) subset or \(a_t=0\). The rest are identical to S-learner.
To learn more about S- and T-learner, please refer to my another post.
Since Q-learning merely trains the outcomes rather than the treatment effect, so there might be error in determining the optimal treatment (Murphy 2005b). Also, this is the reason why it is known as indirect method. Therefore, another branch of methods are proposed under A-learning framework, which try to train directly on the treatment effect and obtain optimal treatment accordingly.
A-learning
A-learning is short for Advantage learning which is another main approach to seek optimal DTR (Blatt et al. 2004; Moodie et al. 2007; Murphy 2003; Robins 2004). We can define a contrast function in two arm scenario as \[ C_t(h_t) := \mathbb E[Y_t \mid H_t=h_t, A_t=1] - \mathbb E[Y_t \mid H_t=h_t, A_t=0] \] and the optimal treatment in terms of this contrast function can be written as \[\begin{equation} g_t^\text{opt}(h_t) = I(C_t(h_t)>0), t = 1,...,T. \tag{4} \end{equation}\] Remarkably, the contrast function is parallel to the treatment effect \(\tau()\) in causal inference, which can be deemed as seeking optimal treatment regime in an one-stage design. In particular, if \(A_t\) is a random assignment, or unconfoundedness assumption holds, that is, \[ (Y_t^{(1)}, Y_t^{(0)}) \perp A_t \mid H_t \] then we have \[ C_t(h_t) = \mathbb E[Y_t^{(1)} \mid H_t=h_t] - \mathbb E[Y_t^{(0)} \mid H_t=h_t] = \mathbb E[Y_t^{(1)} - Y_t^{(0)} \mid H_t=h_t] \] which is the treatment effect. Here \(Y_t^{(1)}\) and \(Y_t^{(0)}\) are the potential outcome of receiving treatment 1 and 0 at \(t^{th}\) stage, respectively.
The value function for A-learning in terms of contrast function is \[ V_t(h_t) = Q_{t}(h_t, 0) + C_t(h_t)I(C_t(h_t)>0) \] where Q-functions are still involved. To get rid of Q-functions, we can re-write it as \[\begin{equation} V_t(h_t) = V_{t+1}(h_{t+1}) + C_t(h_t)(I(C_t(h_t)>0) - A_t) \tag{5} \end{equation}\] and it is obvious that, in contrast to Q-learning, the optimal DTR of A-learning is fully determined by the performance of \(C_t(), t=1,...,T\). In addition, we would like to provide some intuition of equation (5). If the predicted optimal treatments matches with the realized assignments, i.e., \(g_t^\text{opt}(h_t) = A_t\), then we do not need to make correction on next stage’s value function \(V_{t+1}(h_{t+1})\) since it is already optimized; otherwise, we adjust it by the amount of \(C_t(h_t)\). For instance, if the realized \(A_t = 1\) but our estimation says it should be \(0\), then we correct the value function by the treatment effect difference, which is \(-C_t(h_t)>0\). That will increase the future value functions and enlarge the overall rewards. This idea is inspired from Zhang and Zhang (2018).
Notably, if we estimate \(C_t(h_t)\) by \(\hat C_t(h_t) = \hat Q_{t}(h_t, 1) - \hat Q_{t}(h_t, 0)\), then optimal treatment found by equation (4) is identical to that of Q-learning (2). However, people can also bypass the estimation of Q-functions and directly find estimator for \(C_t(h_t)\). So, ideally, A-learning is suppose to be more robust against model misspecification than Q-learning for consistent estimation of the optimal treatment regime (Schulte et al. 2014). There are a bunch of literature propose different ways to approach the contrast function or similar alternatives (Shi et al. 2018; Zhang et al. 2013; Zhang and Zhang 2018; Zhao et al. 2015). In the next subsection, we propose a novel method to estimate contrast function, which also follows the idea of meta-learner.
X-learner
X-learner is originally proposed for heterogeneous treatment effect estimation (Künzel et al. 2019). It fixes the data inefficiency in T-learner and also can outperform T-learner if unbalanced treatment allocation exists (Künzel et al. 2019). Here we extend X-learner for optimizing DTR, and the main procedure is summarized below:
- Estimate \(\hat Q_{t,1}(h_t)\) and \(\hat Q_{t,0}(h_t)\) as described in T-learner
- Impute missing potential “outcomes” by \(\hat Q_{t,1}(h_t)\) and \(\hat Q_{t,0}(h_t)\). In particular, for subjects with \(A_t=1\), their potential outcomes are predicted by \(\hat Q_{t,0}(h_t;a_t=1)\) and then we can construct pseudo contrast values \(\tilde C_t(h_t; a_t=1) = Y_{t+1} + \hat V_{t+1}(h_{t+1}) - \hat Q_{t,0}(h_t;a_t=1)\) for all subjects with assignments \(A_t=1\). Next, we find the optimal contrast function by minimizing the loss function \(\hat C_{t,1}(h_t) = \arg\min_{C_{t,1}} \ell(\tilde C_t(h_t; a_t=1), C_{t,1}(h_t))\) via some base learner. Similarly, we can find another contrast function \(\hat C_{t,0}(h_t) = \arg\min_{C_{t,0}} \ell(\tilde C_t(h_t; a_t=0), C_{t,0}(h_t))\) for subjects with \(A_t=0\)
- Lastly, combine both contrast functions to achieve the final model. One recommended way is \[ \hat C_t(h_t) = \hat C_{t,1}(h_t)(1-\hat g_t(h_t)) + \hat C_{t,0}(h_t)\hat g_t(h_t) \] where \(\hat g_t()\) can be the estimated propensity score values at \(t^{th}\) stage.
Simulation Studies
Case 1: Random assignment with two treatment options
In the first case, we follow the settings adopted in Zhao et al. (2015) and Zhang and Zhang (2018), which is 3-stage randomized study with only two treatment options. Treatment assignments at each stage \(A_1\), \(A_2\), and \(A_3\) are randomly generated from a Bernoulli\((0.5)\) distribution. Baseline covariates as well as covariates observed at each stage is generated as follows:
\[
\begin{aligned}
X_{1,i} &\sim \mathcal N(45,15^2), \quad i = 1,...,10 \\
X_{2,i} &\sim \mathcal N(1.5X_{1,i}, 10^2), \quad i = 1,...,5 \\
X_{3,i} &\sim \mathcal N(0.5X_{2,i}, 10^2), \quad i = 1,...,5.
\end{aligned}
\]
Outcomes or regrets are generated from
\[
\begin{aligned}
Y_1 &= Y_2 = 0, \\
Y_3 &= 20 - |0.6X_{1,1}-40|\{A_1-I(X_{1,1}>30)\}^2 - |0.8X_{2,1}-60|\{A_2-I(X_{2,1}>40)\}^2\\
& - |1.4X_{3,1}-40|\{A_3-I(X_{3,1}>40)\}^2 + \varepsilon, \quad \text{where } \varepsilon \sim \mathcal N(0, \sigma^2).
\end{aligned}
\]
Here we take \(\sigma=1\) for fair comparison with results in existing literature. Since \(Y_1=Y_2=0\), it means that we only care about the final outcome and intermediate returns are of no concerns. Also, from the formula of \(Y_3\), we can see that the expectation of optimal return should be 20, i.e.,
\[
\mathbb E[Y^*(g^\text{opt})] = 20,
\]
where
\[
\begin{aligned}
g^{opt}_1 &= I(X_{1,1}>30) \\
g^{opt}_2 &= I(X_{2,1}>40) \\
g^{opt}_3 &= I(X_{3,1}>40).
\end{aligned}
\]
In this work, we adopt BART (Chipman et al. 1998, 2002, 2010) as base learner for implementation of S-, T-, and X-learner. BART is a powerful learning tool which enjoys the benefit of tree structure, boosting technique, and prior and likelihood, and it is also quite efficient in computation. We directly call the bart() function in R package dbarts (Dorie et al. 2022).
Example R code: S-learner
We first generate the simulated data as described in the previous section using the following code. The training set has 400 samples and the testing set has 1000 samples. A, X, and Y follow the notation in Method section that stand for the treatment assignments, covariates, and outcomes, respectively, and with a suffix tr means training data, te means testing data.
set.seed(54321)
n.sbj = 400; n.test = 1000; sigma = 1
n.stage = 3
nX1 = 10; nX2 <- nX3 <- 5
## Treatment Assignments A at each stage:
A.tr = matrix(rbinom(n.sbj*n.stage,1,0.5), ncol = n.stage)
## Observed covariates X before each stage:
X1.tr = matrix(rnorm(n.sbj*nX1, 45, 15), ncol = nX1)
X2.tr = matrix(rnorm(n.sbj*nX2, X1.tr[,1:nX2], 10), ncol = nX2)
X3.tr = matrix(rnorm(n.sbj*nX3, X2.tr[,1:nX3], 10), ncol = nX3)
## Outcome/Reward generator:
y.fun <- function(X1, X2, X3, A) {
20 - abs(0.6*X1[,1])*abs(A[,1] - (X1[,1]>30)) -
abs(0.8*X2[,1] - 60)*abs(A[,2] - (X2[,1]>40)) -
abs(1.4*X3[,1] - 40)*abs(A[,3] - (X3[,1]>40))
}
Y.tr = y.fun(X1.tr, X2.tr, X3.tr, A.tr) + rnorm(n.sbj,0,sigma)
## test data:
X1.te = matrix(rnorm(n.test*nX1, 45, 15), ncol = nX1)
X2.te = matrix(rnorm(n.test*nX2, X1.tr[,1:nX2],10), ncol = nX2)
X3.te = matrix(rnorm(n.test*nX3, X2.tr[,1:nX3],10), ncol = nX3)The code of applying S-learner on the simulated data to find the optimal DTR is given below. As there are 3 stages, we estimate the Q-functions for each stage separately in a backward order, i.e., we first find \(Q_3()\), then \(Q_2()\), and \(Q_1()\) at last. As long as we obtain these Q-functions, i.e., Q.S1, Q.S2, and Q.S3 in the following code, the optimal dynamic treatments for any given subjects are determined.
##+++++++++++++++++ Q-learning: S-learner style +++++++++++++++++
### For stage 3 first:
S3.dat = cbind(X1.tr, X2.tr, X3.tr, A.tr[,-3], 1-A.tr[,3])
Q.S3 = bart(x.train = cbind(X1.tr, X2.tr, X3.tr, A.tr), y.train = Y.tr,
x.test = S3.dat, # prediction for potential outcomes
ntree = 200, verbose = FALSE, keeptrees = TRUE)
V3.est = pmax(Q.S3$yhat.test.mean, Q.S3$yhat.train.mean) # find V_3
### Stage 2 then:
S2.dat = cbind(X1.tr, X2.tr, A.tr[,-c(2:3)], 1-A.tr[,2])
Q.S2 = bart(x.train = cbind(X1.tr, X2.tr, A.tr[, -3]), y.train = V3.est,
x.test = S2.dat, # prediction for potential outcomes
ntree = 200, verbose = FALSE, keeptrees = TRUE)
V2.est = pmax(Q.S2$yhat.test.mean, Q.S2$yhat.train.mean) # find V_2
### Stage 1 last:
S1.dat = cbind(X1.tr, 1-A.tr[,1])
Q.S1 = bart(x.train = cbind(X1.tr, A.tr[, 1]), y.train = V2.est,
ntree = 200, verbose = FALSE, keeptrees = TRUE)
### Prediction on Testing data
Y1.1 = colMeans(predict(Q.S1, newdata = cbind(X1.te, 1)))
Y1.0 = colMeans(predict(Q.S1, newdata = cbind(X1.te, 0)))
A1.te = ifelse(Y1.1>Y1.0, 1, 0)
Y2.1 = colMeans(predict(Q.S2, newdata = cbind(X1.te, X2.te, A1.te, 1)))
Y2.0 = colMeans(predict(Q.S2, newdata = cbind(X1.te, X2.te, A1.te, 0)))
A2.te = ifelse(Y2.1>Y2.0, 1, 0)
Y3.1 = colMeans(predict(Q.S3, newdata = cbind(X1.te, X2.te, X3.te, A1.te, A2.te, 1)))
Y3.0 = colMeans(predict(Q.S3, newdata = cbind(X1.te, X2.te, X3.te, A1.te, A2.te, 0)))
A3.te = ifelse(Y3.1>Y3.0, 1, 0)
A.opt.S = cbind(A1.te, A2.te, A3.te)
### Results
Y.origin = y.fun(X1.tr, X2.tr, X3.tr, A.tr)
Y.eva.S = y.fun(X1.te, X2.te, X3.te, A.opt.S)
mean(Y.origin) # rewards from original assignments## [1] -22.92242mean(Y.eva.S) # rewards from optimized assignments## [1] 19.86825mean(Y.eva.S==20) # percentage that receive the best treatment seq.## [1] 0.992Focus on this simulated sample, we can see that in original study design, since subjects are randomized at each stage, the treatment effects are clearly not optimal, with mean outcome -22.92. However, after trained by S-learner, the average outcome following the optimized DTR can reach to 19.87 (the upper bound is 20), with 99.2% subjects receives the best treatment sequence for them.
Example R code: T-learner
With the same dataset, we now try T-learner. The code is shown below. Since we need to train two models for \(A=1\) and \(A=0\) separately, there are Q.S_.1 and Q.S_.0 at each stage representing two Q-functions.
##+++++++++++++++++ Q-learning: T-learner style +++++++++++++++++
### For stage 3 first:
S3.dat = cbind(X1.tr, X2.tr, X3.tr, A.tr[,-3])
Q.S3.1 = bart(x.train = cbind(X1.tr, X2.tr, X3.tr, A.tr[,-3])[A.tr[,3]==1,],
y.train = Y.tr[A.tr[,3]==1],
x.test = S3.dat,
ntree = 200, verbose = FALSE, keeptrees = TRUE)
Q.S3.0 = bart(x.train = cbind(X1.tr, X2.tr, X3.tr, A.tr[,-3])[A.tr[,3]==0,],
y.train = Y.tr[A.tr[,3]==0],
x.test = S3.dat,
ntree = 200, verbose = FALSE, keeptrees = TRUE)
V3.est = pmax(Q.S3.1$yhat.test.mean, Q.S3.0$yhat.test.mean) # find V_3
### Stage 2 then:
S2.dat = cbind(X1.tr, X2.tr, A.tr[,-c(2:3)])
Q.S2.1 = bart(x.train = cbind(X1.tr, X2.tr, A.tr[, -c(2:3)])[A.tr[,2]==1,],
y.train = V3.est[A.tr[,2]==1],
x.test = S2.dat,
ntree = 200, verbose = FALSE, keeptrees = TRUE)
Q.S2.0 = bart(x.train = cbind(X1.tr, X2.tr, A.tr[, -c(2:3)])[A.tr[,2]==0,],
y.train = V3.est[A.tr[,2]==0],
x.test = S2.dat,
ntree = 200, verbose = FALSE, keeptrees = TRUE)
V2.est = pmax(Q.S2.1$yhat.test.mean, Q.S2.0$yhat.test.mean) # find V_2
### Stage 1 last:
S1.dat = cbind(X1.tr)
Q.S1.1 = bart(x.train = cbind(X1.tr)[A.tr[,1]==1,], y.train = V2.est[A.tr[,1]==1],
x.test = S1.dat, ntree = 200, verbose = FALSE, keeptrees = TRUE)
Q.S1.0 = bart(x.train = cbind(X1.tr)[A.tr[,1]==0,], y.train = V2.est[A.tr[,1]==0],
x.test = S1.dat, ntree = 200, verbose = FALSE, keeptrees = TRUE)
### Prediction on Testing data
Y1.1 = colMeans(predict(Q.S1.1, newdata = X1.te))
Y1.0 = colMeans(predict(Q.S1.0, newdata = X1.te))
A1.te = ifelse(Y1.1>Y1.0, 1, 0)
Y2.1 = colMeans(predict(Q.S2.1, newdata = cbind(X1.te, X2.te, A1.te)))
Y2.0 = colMeans(predict(Q.S2.0, newdata = cbind(X1.te, X2.te, A1.te)))
A2.te = ifelse(Y2.1>Y2.0, 1, 0)
Y3.1 = colMeans(predict(Q.S3.1, newdata = cbind(X1.te, X2.te, X3.te, A1.te, A2.te)))
Y3.0 = colMeans(predict(Q.S3.0, newdata = cbind(X1.te, X2.te, X3.te, A1.te, A2.te)))
A3.te = ifelse(Y3.1>Y3.0, 1, 0)
A.opt.T = cbind(A1.te, A2.te, A3.te)
### Results
Y.eva.T = y.fun(X1.te, X2.te, X3.te, A.opt.T)
mean(Y.eva.T)## [1] 19.17569mean(Y.eva.T==20)## [1] 0.921Example R code: X-learner
Lastly, we show how to apply X-learner to find DTR. X-learner is more complicated than previous two learners. It starts with estimating two Q-functions just like T-learner. Then, impute the potential outcomes and find \(\tilde C_t(h_t;a_t=1)\) and \(\tilde C_t(h_t;a_t=0)\), which is stored in tau_.1 and tau_.0, respectively. The contrast functions are estimated by bart() as well. In the end, we combine two contrast functions with propensity score as weights. Since this is a randomized study, the propensity scores are estimated via observed proportion of \(A=1\) at each stage.
##+++++++++++++++++ A-learning: X-learner style +++++++++++++++++
prop = colMeans(A.tr)
### For stage 3 first:
#### first estimate the counterfactual outcomes from T-learner:
S3.dat = cbind(X1.tr, X2.tr, X3.tr, A.tr[,-3])
Q.S3.1 = bart(x.train = cbind(X1.tr, X2.tr, X3.tr, A.tr[,-3])[A.tr[,3]==1,],
y.train = Y.tr[A.tr[,3]==1],
x.test = S3.dat, ntree = 200, verbose = FALSE, keeptrees = TRUE)
Q.S3.0 = bart(x.train = cbind(X1.tr, X2.tr, X3.tr, A.tr[,-3])[A.tr[,3]==0,],
y.train = Y.tr[A.tr[,3]==0],
x.test = S3.dat, ntree = 200, verbose = FALSE, keeptrees = TRUE)
Y3.pred.0 = predict(Q.S3.0, newdata = S3.dat)
tau3.1 = (Y.tr - colMeans(Y3.pred.0))[A.tr[,3]==1]
Y3.pred.1 = predict(Q.S3.1, newdata = S3.dat)
tau3.0 = (colMeans(Y3.pred.1) - Y.tr)[A.tr[,3]==0]
#### then estimate the C()
C.S3.1 = bart(x.train = S3.dat[A.tr[,3]==1,], y.train = tau3.1,
x.test = S3.dat, ntree = 200, verbose = FALSE, keeptrees = TRUE)
C.S3.0 = bart(x.train = S3.dat[A.tr[,3]==0,], y.train = tau3.0,
x.test = S3.dat, ntree = 200, verbose = FALSE, keeptrees = TRUE)
#### the optimal outcome is determined by C()
V3.est = Y.tr + (C.S3.1$yhat.test.mean*(1-prop[3]) + C.S3.0$yhat.test.mean*prop[3])*
((C.S3.1$yhat.test.mean*(1-prop[3]) + C.S3.0$yhat.test.mean*prop[3] > 0) - A.tr[,3])
### Stage 2 then: (re-estimate all Qs now)
#### first estimate new Q functions:
S2.dat = cbind(X1.tr, X2.tr, A.tr[,-c(2:3)])
Q.S2.1 = bart(x.train = cbind(X1.tr, X2.tr, A.tr[, -c(2:3)])[A.tr[,2]==1,],
y.train = V3.est[A.tr[,2]==1],
x.test = S2.dat, ntree = 200, verbose = FALSE, keeptrees = TRUE)
Q.S2.0 = bart(x.train = cbind(X1.tr, X2.tr, A.tr[, -c(2:3)])[A.tr[,2]==0,],
y.train = V3.est[A.tr[,2]==0],
x.test = S2.dat, ntree = 200, verbose = FALSE, keeptrees = TRUE)
#### estimate counterfactual outcomes
Y2.pred.0 = predict(Q.S2.0, newdata = S2.dat)
tau2.1 = (V3.est - colMeans(Y2.pred.0))[A.tr[,2]==1]
Y2.pred.1 = predict(Q.S2.1, newdata = S2.dat)
tau2.0 = (colMeans(Y2.pred.1) - V3.est)[A.tr[,2]==0]
#### estimate the C() for stage 2
C.S2.1 = bart(x.train = S2.dat[A.tr[,2]==1,], y.train = tau2.1,
x.test = S2.dat, ntree = 200, verbose = FALSE, keeptrees = TRUE)
C.S2.0 = bart(x.train = S2.dat[A.tr[,2]==0,], y.train = tau2.0,
x.test = S2.dat, ntree = 200, verbose = FALSE, keeptrees = TRUE)
#### the optimal outcome is determined by As
V2.est = V3.est + (C.S2.1$yhat.test.mean*(1-prop[2]) + C.S2.0$yhat.test.mean*prop[2])*
((C.S2.1$yhat.test.mean*(1-prop[2]) + C.S2.0$yhat.test.mean*prop[2] > 0) - A.tr[,2])
### Stage 1 last:
#### first estimate new Q functions:
S1.dat = cbind(X1.tr)
Q.S1.1 = bart(x.train = cbind(X1.tr)[A.tr[,1]==1,],
y.train = V2.est[A.tr[,1]==1],
x.test = S1.dat, ntree = 200, verbose = FALSE, keeptrees = TRUE)
Q.S1.0 = bart(x.train = cbind(X1.tr)[A.tr[,1]==0,],
y.train = V2.est[A.tr[,1]==0],
x.test = S1.dat, ntree = 200, verbose = FALSE, keeptrees = TRUE)
#### estimate counterfactual outcomes
Y1.pred.0 = predict(Q.S1.0, newdata = S1.dat)
tau1.1 = (V2.est - colMeans(Y1.pred.0))[A.tr[,1]==1]
Y1.pred.1 = predict(Q.S1.1, newdata = S1.dat)
tau1.0 = (colMeans(Y1.pred.1) - V2.est)[A.tr[,1]==0]
#### estimate the C() for stage 1
C.S1.1 = bart(x.train = S1.dat[A.tr[,1]==1,], y.train = tau1.1,
ntree = 200, verbose = FALSE, keeptrees = TRUE)
C.S1.0 = bart(x.train = S1.dat[A.tr[,1]==0,], y.train = tau1.0,
ntree = 200, verbose = FALSE, keeptrees = TRUE)
### Prediction on Testing data (based on Cs now, not Qs)
tau1.1 = colMeans(predict(C.S1.1, newdata = X1.te))
tau1.0 = colMeans(predict(C.S1.0, newdata = X1.te))
A1.te = ifelse(tau1.1*(1-prop[1]) + tau1.0*prop[1] > 0, 1, 0)
tau2.1 = colMeans(predict(C.S2.1, newdata = cbind(X1.te, X2.te, A1.te)))
tau2.0 = colMeans(predict(C.S2.0, newdata = cbind(X1.te, X2.te, A1.te)))
A2.te = ifelse(tau2.1*(1-prop[2]) + tau2.0*prop[2]> 0, 1, 0)
tau3.1 = colMeans(predict(C.S3.1, newdata = cbind(X1.te, X2.te, X3.te, A1.te, A2.te)))
tau3.0 = colMeans(predict(C.S3.0, newdata = cbind(X1.te, X2.te, X3.te, A1.te, A2.te)))
A3.te = ifelse(tau3.1*(1-prop[3]) + tau3.0*prop[3]> 0, 1, 0)
A.opt.X = cbind(A1.te, A2.te, A3.te)
### Results
Y.eva.X = y.fun(X1.te, X2.te, X3.te, A.opt.X)
mean(Y.eva.X)## [1] 19.55917mean(Y.eva.X==20)## [1] 0.959Case 1: Simulation Results
In previous subsections, we only see the results from a single simulated dataset. To have a better understanding of the methods, we apply the meta-learners on 500 simulated datasets and increase the testing sample size to 10,000, which is adopted in the literature. The mean and standard deviation of estimated \(\mathbb E[Y^*(g^\text{opt})]\) are summarized in the Table 1. We can see that when sample size is small, T-learner is the worst comparing to other two methods. This is understandable because it trains each Q-function only on a part of data, which can lead to terrible Q estimations. While for X-learner, though starting from T-learner, fixes its data inefficiency issue. Therefore, its performance is much better than T-learner when sample size is not large enough for estimating Q functions separately. When sample size increases, all methods’ performance improves tremendously. We also attached the simulation results from existing literature (Zhang and Zhang 2018) under the identical simulation setting; see Figure 3. In the table, there are results from Backward Outcome Weighted Learning (BOWL) (Zhao et al. 2015), a doubly robust augmented inverse probability weighted estimator (AIPWE) proposed by Zhang et al. (2013), and C-learning method by Zhang and Zhang (2018). All of these methods can be deemed as A-learning approaches. By comparing the results in Figure 3 to that in Table 1, we can see the superior performance of meta-learners in both mean and sd, especially S-learner under Case 1 setting. It is worth to mention that the results shown in Figure 3 does not actually involve noise covariates. They basically only include three \(X_1\)s, one \(X_2\), and one \(X_3\) in the analysis.
| n = 200 | n = 400 | n = 800 | |
|---|---|---|---|
| S-learner | 17.76 (0.98) | 19.36 (0.31) | 19.53 (0.2) |
| T-learner | 11.89 (2.64) | 18.55 (0.64) | 19.48 (0.18) |
| X-learner | 15.82 (1.04) | 19.3 (0.29) | 19.64 (0.13) |
Figure 3: Simulation results from other literatures with the same simlation setting
Case 2: Two stage rewards
In Case 1, intermediate rewards are set to be 0, while in Case 2, these rewards will be taken into consideration, or in other words, the weights are now \(w_t=1\) rather than \(w_t=0, t\in\{1,2\}\) in Case 1. Additionally, in Case 1 the covariates at later stages are predetermined before the trial starts, i.e., they are independent of all future interventions, which is not quite realistic. The simulation setting we adopted here is borrowed from simulation scenario 2 in Zhao et al. (2015). Specifically, there are 50 baseline covariates that are generated from normal distribution and two intermediate covariates are generated accordingly; see below. \[ \begin{aligned} X_{1,i} &\sim \mathcal N(0, 1), \quad i = 1,...,50 \\ X_{2,1} &\sim I[\mathcal N(1.25X_{1,1}(2A_1-1), 1) > 0] \\ X_{2,2} &\sim I[\mathcal N(-1.75X_{1,2}(2A_1-1), 1) > 0] \end{aligned} \] and the treatment assignments at stage 1 and 2 are randomly generated from \(\{0,1\}\) with equal probability. The outcomes at each stage are generated as follows \[ \begin{aligned} Y_1 &\sim \mathcal N((1+1.5X_{1,3})(2A_1-1), 1) \\ Y_2 &\sim \mathcal N((Y_1+A_1+0.5X_{2,1}-0.5X_{2,2})(2A_2-1), 1) \end{aligned} \]
Case 2: Simulation Results
Similar to Case 1, the sample size of training varies from 100, 200, and 400, and validation set 10,000. We repeat 500 times for each scenario. The mean and variance of the estimated DTR is provided in the Table 2. The optimal mean DTR in this setting is \(3.667\). Unlike in Case 1 simulation study, T-learner yields superior performance comparing to S- and X-learner. Meanwhile, we provide the simulation results in Zhao et al. (2015) for a direct comparison; see scenario 2 in Figure 4. We can find the proposed meta-learner methods perform consistently well especially T-learner under this setting. In addition, the scenario 3 are the results of Case 1 simulation study. If considering both scenarios together, meta-learners’ performance is quite robust across all settings.
| n = 100 | n = 200 | n = 400 | |
|---|---|---|---|
| S-learner | 2.933 (0.497) | 3.081 (0.17) | 3.241 (0.106) |
| T-learner | 3.039 (0.084) | 3.184 (0.077) | 3.384 (0.087) |
| X-learner | 3.006 (0.067) | 3.104 (0.078) | 3.377 (0.102) |
Figure 4: Simulation results for Case 1 (Scenario 3) and Case 2 (Scenario 2) from Zhao (2015)
Discussions
Multi-armed Experiment
In our previous discussions, we merely focus on scenarios with only two treatment options. When there are more than two arms, Q-learning approaches can easily be extended in most cases, it just has more Q-functions at each stage corresponding to each arm, but the value function is still well-defined and readily obtainable. However, most A-learning methods are not straightforward under multi-armed settings. The main reason is, the target quantity, the treatment effect or contrast function, need to be carefully tailored. For example, one potential way of adjustment for A-learning methods is to work on a one-versus-rest schema, or focus on pairwise comparisons between any two arms, i.e., one-versus-one schema.
These conclusions are also true for meta-learners. S- and T-learner method can easily deal with multiple-treatment scenarios but for X-learner, careful adjustments are required.
Non-randomized Design
In general, there are two sources of confounding in DTR, one comes from the unobserved influential covariates, and the other comes from past treatment assignments. For the first one, we can eliminate the impact of unobserved confounders through randomization. The second one is handled by the optimization structure introduced in Method section. When randomization is no longer achievable, we lose the control of unobserved confounding factors, and thus, to make valid inference, we need assumptions like unconfoundedness and overlapping.
Technically, propensity score will play a more important role in non-randomized studies especially for A-learning approaches, as it could provide additional security to the model misspecification, known as double robustness, along with the outcome models. Notably, for Q-learning methods, propensity score is not involved a lot. Though lose doubly robust property, Q-learning methods can achieve better approximation of outcome models, like S- and T-learner methods with powerful base learner. More importantly, in practice, when the treatment numbers are large, estimating propensity score is problematic because overlapping assumption can easily be violated. Since propensity score can have extreme values in these cases, the estimators involved propensity score might be impaired as well.
See Appendix for an example of using metaDTR package to learn the optimal DTR from a simulated observational, multi-armed, multi-stage study.
VS Reinforcement Learning
Clearly, the whole idea of seeking optimal DTR is parallel to reinforcement learning. The key is all about balancing the short-term and long-term returns to reach the optimal rewards eventually. There are some slight differences in notation/terminology though. We know the key elements of reinforcement learning includes: a set of environment and agent states, \(S\), which is history information \(H\) in this work; a set of actions, \(A\), which is equivalent to treatment assignment \(A\) in this work; reward after receiving action, \(R\), which is the outcome \(Y\) here; the goal is to learn policy, \(\pi(s) = \Pr(a_t=1\mid s_t=s)\), which is denoted by \(g_t(h_t)\). And we also have similar idea in construction of \(Q()\) and \(V()\) functions. However, the main difference is in SMART or sequential observational study, it’s a finite-decision problem with extremely small number of decisions to make, if compared with problems in reinforcement learning world. This leads to very different way to solve the problem. In DTR seeking, we do not need to know the transition probability of states, i.e., we do not need to consider the environment’s interactions; and we can use back deduction to solve the problem given finite and limited stages.
References
Appendix
Here we consider a simulation setting that is way more complicated than the cases used in simulation section. It is an observational, multi-armed, and multi-stage study. Consider a two-stage study that at first stage, subjects only take treatment \(\in \{0,1\}\), but in second stage, subject’s treatment assignment will be based on the previous outcomes, i.e., if original control effect (\(A=0\)) is not satisfying, subject will move to treatment arms (\(A=1\)); else, they stay in control arm. If previous treatment effect (\(A=1\)) is good, they have chance to further escalate to higher dose level (\(A \in \{2, 3\}\)); else, switch to control arm. See detailed simulation settings below: \[ \begin{aligned} X_{1,i} &\sim \mathcal N(0, 1), \quad i = 1,...,50 \\ X_{2,1} &\sim I[\mathcal N(1.25X_{1,1}(2A_1-1)+X_{1,2}, 1) > 0] \\ X_{2,2} &\sim I[\mathcal N(-1.75X_{1,2}(2A_1-1)+X_{1,3}, 1) > 0] \\ X_{2,j} &\sim \mathcal N(0, 1), \quad i = 3,...,10 \end{aligned} \] \[ \begin{aligned} \text{logit}(\Pr(A_{1}=0)) &= X_{1,1} + X_{1,2}X_{1,3} + X_{1,4}^2/3 \\ \Pr(A_{2}=0) &= \frac{(1-A_1)Y_1^2 + X_{2,1}}{C} \\ \Pr(A_{2}=1) &= \frac{A_1Y_1^2 + X_{2,1}}{C} \\ \Pr(A_{2}=2) &= \frac{A_1Y_1^2 + X_{2,1}}{C} \\ \Pr(A_{2}=3) &= \frac{A_1Y_1^2/6 + X_{2,1}}{C} \\ C &= \sum_{j=0}^3 \Pr(A_{2}=j) \end{aligned} \]
\[
\begin{aligned}
Y_1 &\sim \mathcal N((1+X_{1,3}-X_{1,2})(2A_1-1)+X_{1,5}, 1) \\
Y_2 &\sim \mathcal N((Y_1+A_1+0.5X_{2,1}-0.5X_{2,2})A_2 + X_{1,5}, 1)
\end{aligned}
\]
The code of applying package metaDTR is provided below:
## solve by package metaDTR:
X = list(X1.tr, X2.tr)
A = list(as.factor(A1.tr), as.factor(A2.tr))
Y = list(Y1.tr, Y2.tr)
DTRs = learnDTR(X = X,
A = A,
Y = Y,
weights = rep(1, length(X)),
baseLearner = c("BART"),
metaLearners = c("S", "T"),
include.X = 1,
include.A = 2,
include.Y = 1,
verbose = FALSE
)
## test data at baseline
X1.te = matrix(rnorm(n.test*nX1, 0, sigma), ncol = nX1)
## make recommendations based on current info
A1.opt = recommendDTR(DTRs, X.new = X1.te)
## Y2/X2 is dependent on A1 (in practice, Y2/X2 should be observed)
Y1.te = Y1.fun(n.test, X1.te, as.numeric(A1.opt$A.opt.S$Stage.1)-1)
X2.te = X2.fun(n.test, X1.te, as.numeric(A1.opt$A.opt.S$Stage.1)-1)
## then keep updating it by providing newly observed info
A2.opt = recommendDTR(DTRs, currentDTRs = A1.opt,
X.new = X2.te, , Y.new = Y1.te)
## original rewards:
mean(Y1.tr+Y2.tr)## [1] 1.66878## optimized rewards (S-learner)
mean(Y.fun(X1.te, X2.te, Y1.te,
as.numeric(A2.opt$A.opt.S$Stage.1)-1,
as.numeric(A2.opt$A.opt.S$Stage.2)-1))## [1] 6.335082## optimized rewards (T-learner)
mean(Y.fun(X1.te, X2.te, Y1.te,
as.numeric(A2.opt$A.opt.T$Stage.1)-1,
as.numeric(A2.opt$A.opt.T$Stage.2)-1))## [1] 6.106891## distribution of optimized A1:
table(A2.opt$A.opt.T$Stage.1)##
## 0 1
## 186 814## distribution of optimized A2:
table(A2.opt$A.opt.T$Stage.2)##
## 0 1 2 3
## 33 39 555 373Clearly, optimal DTR on average yields way better rewards than original assignment that without any personalization. Also, by simulation setting we know that larger A better treatment effect, which is also well reflected in the distribution of recommended treatment assignments at both stage.
Here is another example of applying package metaDTR on simulation Case 1 data. In that case, intermediate covariates can be predetermined before the study. So the optimal DTRs can be obtained immediately from the function recommendDTR(). The results are close to what we have seen in Simulation section.
X = list(X1.tr, X2.tr, X3.tr)
A = list(as.factor(A.tr[,1]), as.factor(A.tr[,2]), as.factor(A.tr[,3]))
Y = list(0, 0, Y.tr)
## training
DTRs = learnDTR(X = X,
A = A,
Y = Y,
weights = rep(1, length(X)),
baseLearner = c("BART"),
metaLearners = c("S", "T"),
include.X = 1,
include.A = 2,
include.Y = 0,
verbose = FALSE
)
## estimate optimal DTR for testing data:
optDTR <- recommendDTR(DTRs, currentDTRs = NULL,
X.new = list(X1.te, X2.te, X3.te))
## S-learner results:
Y.eva.S = y.fun(X1.te, X2.te, X3.te,
matrix(as.numeric(unlist(optDTR$A.opt.S))-1, ncol = 3, byrow = F))
mean(Y.eva.S)## [1] 19.61549## T-learner results:
Y.eva.T = y.fun(X1.te, X2.te, X3.te,
matrix(as.numeric(unlist(optDTR$A.opt.T))-1, ncol = 3, byrow = F))
mean(Y.eva.T)## [1] 19.44285