\section{Introduction}
Multi-stage treatments are becoming more prevalent in medicine. 
Chronic conditions such as cancer, auto-immune disease, HIV, and heart disease, often require multiple treatments over a long period. For instance, cancer patients may receive multiple rounds of chemo- and/or biological therapies \citep{habermann2006, huang2020}. Auto-immune disease patients might receive additional treatments in response to relapses \citep{edwards2006, hogan2013}. Traditionally such chronic diseases have been treated following a one-size-fits-most approach by which the standard of care is selected as the treatment or treatment plan that will likely benefit most patients with a similar condition, i.e., the treatment that is optimal only in the sense of the marginal effects.

However in recent years, a more personalized approach to the treatment of chronic diseases has become of interest. Under this approach patients benefit from treatment rules that incorporate the patient heterogeneity of treatment effects and thereby treatment decisions are based on the individual characteristics of each patient.
Dynamic treatment regimes formalize this data-driven approach to treatment by optimizing the overall outcome of interest and providing treatment suggestions that are based on the patient's information available at each point of the decision making process \citep{murphy2003, kosorok2015, kosorok2019, tsiatis2019}. In part because they utilize all available information, such optimized dynamic treatment regimes are often found to be more beneficial than a traditional set of fixed treatment plans\citep{kidwell2015}. Further, because optimal dynamic treatment regimes focus on the overall or long-term outcome, they are more advantageous than applying multiple single-stage treatment rules that only optimize the stage-level outcomes.

Finding an optimal dynamic treatment regime is a reinforcement learning problem \citep{sutton2018}, wherein an action by an agent, or a physician, changes both the immediate reward and the environment of the next stage and the learner, or the statistician, searches for the rule of actions that brings the most overall reward. Like many statistical learning methods, reinforcement learning problems are often solved using the Q-learning algorithm \citep{watkins1989, zhu2015, zhao2011, zhang2017, qian2011}, which defines the value as the expected cumulative sum of discounted rewards and optimizes the set of rules using backward recursion.


In precision medicine applications, the rewards are often defined in the context of survival outcomes, such as mean survival time or survival probabilities, which pose unique challenges. For survival outcomes, the failure time is often censored, and any approach should take into account the missing information. Several methods have been developed for this purpose, for example, redistribution of the failure probability as in the Kaplan-Meier \citet{efron1967} and \citet{robertson1984} or inverse probability weighting \citep{robins1994, robins1992, wahed2006, orellana2010}. Further, the treatment timing may not be fixed and may even depend on a patient's status or previous treatments having association with the failure time. In such scenarios, the three events at each stage are failure, proceeding to the next treatment, and censoring, which can be restated in a competing risks framework, in which the three dependent events compete and only the earliest event is observed. However, the full joint distribution of the events is not identifiable without knowledge of their dependence structure \citep{tsiatis1975}. Thus, the dependency between the events needs to be carefully considered. 
For example, a cancer patient whose physical condition has improved 
has a longer expected survival time and at the same time may want to take another round of chemotherapy earlier than expected. If this is the case, the survival time and the treatment time have a negative association. 

There are two additional challenges that arise in the analysis of survival outcomes that must be addressed. First, the total number of treatments received may vary among patients. For example, some patients may experience the failure or be censored before taking all planned treatments. 
And finally, there could be different optimization criteria for survival times. For some patients and clinicians, the average survival time is of primary interest, while the six-month survival probability may be of greater concern for others. A method that can address different types of endpoints would be useful in practice.


Several methods regarding the estimation of dynamic treatment regimes for survival outcomes have been proposed in the literature. \cite{goldberg2012} first developed a dynamic treatment regime  estimator for survival outcomes using a Q-learning framework. To address the problems introduced by a different number of treatments as well as by censoring, the authors proposed modifying the survival data so that each observation has the same number of treatment stages without missing values and showed that the value of a policy in the original problem can be represented as an expectation of the modified data. After the data modification, the problem was put into a standard Q-learning framework using inverse probability weighting to account for censoring. However, censoring was assumed to be completely independent of all covariates and event times. The estimated Q-function and the resulting decision rules would be thus subject to bias for data with dependent censoring.

\citet{huang2014} addressed a similar problem using backward recursion. Their motivating problem was a recurrent disease clinical trial where patients receive an initial treatment followed by a salvage treatment if patients experience either treatment resistance or relapse. 
\citet{simoneau2019} extended the dynamic treatment regime estimator for continuous outcomes via weighted least squares \citep{wallace2015} into the censored time-to-event data setting. Both of these methods use the accelerated failure time model for the failure time distribution, which carries a risk of model misspecification. Thus, a more flexible model is needed to allow for more relaxed distributional assumptions.
\cite{wahed2013} proposed a dynamic treatment regime estimator using full specification of the likelihood and \citet{xu2016} developed a Bayesian alternative, where the disease progression dynamic was modeled using a dependent Dirichlet process prior with Gaussian process measure. While the latter approach provides a more flexible modeling framework than the former, the class of regimes for both methods is refined to a fixed number of pre-defined treatment sequences.

Each of the existing methods described above has its limitations. First, not all methods allow for a flexible number of treatment stages and/or treatment levels. For example, the methods in \citet{wahed2013} and \citet{huang2014} were designed for a two-stage decision problem, and the \citet{simoneau2019} approach is restricted to binary treatments.
Second, the outcomes of interest in most methods are limited to the mean survival time. \citet{jiang2017} proposed an optimal dynamic treatment regime estimator that maximizes the $t$-time survival probability. However, none of the methods permit choosing the criterion except \citet{wahed2013}.
Third, parametric models or strong structural assumptions are employed for all of these methods, with the exception of \citet{goldberg2012}. Strong modeling assumptions, such as the accelerated failure time model or the proportional hazards models used by the existing methods, are subject to model-misspecification, may limit the policy class, and, consequently, may invoke loss in the value of the optimal policy.
Finally, the censoring assumptions of \cite{goldberg2012} are restrictive. Censoring is associated with failure time in many medical applications, but it is often reasonable to assume independence of censoring after conditioning on patient historical information. 

\begin{table}
\def~{\hphantom{0}}
\caption{Assumptions of the existing and the proposed methods}{%
\footnotesize
\begin{tabular}{lcccccc}
    \hline
    method & $Q$ & $|\mathcal A^{(q)}|$ & criterion & $T$ & $C$ & $\{\boldsymbol\pi\}$\\
    \hline
    the new method & finite& finite& $E[T\wedge \tau], S(t)$ &  NP & CI & flexible\\ 
    \citet{goldberg2012} & finite & finite & $E[T\wedge \tau]$ & NP & independent& flexible\\
    \citet{huang2014} & 2+ & 2+ & $E[T]$& AFT & CI & linear\\ 
    \citet{jiang2017} & finite & 2 & $S(t)$ & PH & CI & linear\\ 
    \citet{simoneau2019} & finite & 2 & $E[T]$ & AFT& CI & linear\\ 
    \citet{wahed2013} & finite & finite & $E[T]$, $S(t)$ & AFT & CI & fixed\\ 
    \citet{xu2016} & finite & finite & $E[T]$& BNP & CI & fixed\\ 
    \hline
    \end{tabular}
}
\label{tab:methods}
\footnotesize{	
$Q$, the number of stages; $|\mathcal A^{(q)}|$, the number of treatment arms at stage $q$; criterion, the target value to be optimized; $T$, failure time model; $C$, censoring assumption; $\{\boldsymbol\pi\}$, policy class;
2+, can be generalized to more than two;
NP, non-parametric; AFT, accelerated failure time; PH, proportional hazards; BNP, Bayesian non-parametric; 
CI, conditional independence; fixed, a fixed number of distinct treatment rules.}
\end{table}

In this manuscript, we develop a general dynamic treatment regime estimator for censored time-to-failure outcomes that addresses all of the limitations discussed for the existing estimators. Our estimator maximizes either the mean survival time or the survival probability at a certain time using backward recursion. As the objective may not be expressed as the sum of intermediate rewards, a standard Q-learning algorithm is not applicable. Instead, the conditional survival probability information, rather than the summarized Q-values, is appended to the previous stage information. A generalized random survival forest is developed for this task, where survival curves for each individual, instead of the observed survival or censoring time, is fed into the random survival forest. A general implementation of our method is available on CRAN, R package \texttt{dtrSurv}.

The key contributions of our work are the flexibility of the proposed method and the exposition of its theoretical properties. The proposed method affords significantly more flexibility than existing methods. The target value can be either the mean survival or the survival probability at a certain time. The method allows for a flexible number of treatment stages as well as a flexible number of treatment arms at each stage. Censoring can be conditionally independent. The conditional survival probability is modeled nonparametrically using a random survival forest-based algorithm. And finally, the proposed method does not require estimating the joint distribution of failure time and treatment time nor does it require any assumptions on their dependency, though these are allowed. We further show the consistency of the random survival forests and the estimated regime values.


Before we give an outline of the paper, we briefly review random survival forests, an important component of the proposed method. A random survival forest is a nonparametric survival regression estimator that provides weighted Kaplan-Meier curves for each point in the feature space with weights given by adaptive kernel estimates. In general, a random forest is an ensemble of trees each of which are generated by recursively partitioning the data into two daughter nodes. In the \citet{breiman1984} Classification and Regression Trees, or CART, approach, partitions are made by searching across all variables and data midpoints, and the best cut point that maximizes the heterogeneity between the resulting daughter nodes based on the mean differences is taken. In \citet{breiman2001} random forests, \citet{breiman1984} trees are built and averaged based on multiple bootstrap samples considering only a random subset of the variables when splitting each node. This randomization---bootstrapping and random partitioning---brings diversity among trees, and thus the variability of the estimators can be less than that obtained in the CART approach. \citet{geurts2006} added an alternative randomization mode by restricting the candidate threshold values to a randomly chosen cut point for each variable. 

Survival trees (random survival forests) generalize trees (random forests) by allowing censored outcomes as inputs and providing covariate-conditional survival probability estimates as outputs. See \citet{gordon1985} and \citet{segal1988} for survival trees and \citet{hothorn2005}, \citet{rsf}, and \citet{zhu2012} for random survival forests. Due to censoring, heterogeneity between two nodes cannot be ascertained by a simple mean; alternate tests such as the log-rank test \citep{peto1972, gehan1965, tarone1977} must be used. At each terminal node of a survival tree, the survival probability is estimated by the Kaplan-Meier procedure. A random survival forest is obtained by averaging the survival trees.

The remainder of the manuscript is organized as follows. In Section~\ref{s:method}, we give notation and describe the dynamic treatment rule estimator. Consistency of the estimator is derived in Section~\ref{s:theory}, and its performance is illustrated through simulations in Section~\ref{s:sim}. Section~\ref{s:leukemia} is devoted to the application of the proposed method to a two-stage leukemia clinical trial. We conclude the paper by discussing future research directions in Section~\ref{s:discussion}.

\section{The method}
\label{s:method}
\subsection{Data setup and notation}
\label{ss:notation}
We assume that each of $n$ patients can have a maximum of $Q$ visits, or treatment stages, with a maximum study length $\tau > 0$. At the $q$th stage, $q = 1,2,...,Q$, the $i$th patient, among the $n^{(q)}\le n$ available patients, receives treatment $A^{(q)}_i \in \mathcal A^{(q)}$, if he or she has survived and has not dropped out by the beginning of the stage, where $\mathcal A^{(q)}$ is the finite treatment space for the $q$th stage, $q = 1,2,...,Q$. Throughout this manuscript, we often drop the subject index $i$ when it does not cause confusion. Our interest lies in  estimating the survival distribution $S_{\boldsymbol\pi}$ of the overall failure time $T$ of patients if they followed a dynamic treatment regime $\boldsymbol\pi = (\pi^{(1)}, \pi^{(2)}, ... \pi^{(Q)})^\top$, and then finding the `best' rule $\boldsymbol\pi_*$, which is the rule that optimizes a certain criterion $\phi$. For $\phi$, we consider $\phi_\mu(S) = -\int_{t>0} (t\wedge \tau) dS(t)$ and $\phi_{\sigma,t}(S) = S(t)$ for some $0<t\le\tau$.
Without loss of generality, we assume that a prolonged time to event is preferred; thus the objective is to maximize $\phi(S)$ over the regimes.

We allow the time to next treatment $U^{(q)}_i$ of the $i$th patient to depend on the historical information $\boldsymbol H^{(q)}_i\in \mathcal H^{(q)}$  available at the beginning of the $q$th stage and the treatment assignment $A^{(q)}_i$. 
At the $q$th stage, we observe either the $i$th patient's failure, advancement to next treatment, or censoring. The times from the beginning of the stage to each of these events are denoted as $T^{(q)}_i$, $U^{(q)}_i$, and $C^{(q)}_i$, respectively. The length of the stage $X^{(q)}_i$ for the $i$th patient is defined by the minimum of these times, and the hypothetically uncensored stage length is defined as $V^{(q)}_i = T^{(q)}_i \wedge U^{(q)}_i$. And, if $q=Q$, $V^{(Q)}_i = T^{(Q)}_i$. We denote the $q$th stage censoring and treatment indicators as $\delta_i^{(q)} = 1(V_i^{(q)} \le C_i^{(q)})$ and $\gamma_i^{(q)} = 1(T^{(q)}_i \le U^{(q)}_i)$, respectively. Note that when $C^{(q)}_i < V^{(q)}_i$, $\gamma_i^{(q)}$ is not observable. The baseline time at stage $q$ is defined as $B^{(q)}_i = \sum_{k=1}^{q-1}X^{(k)}_i,  q>1$, and $B_i^{(1)}=0$. Let $T_i,$ $C_i$,
$X_i = T_i \wedge C_i,$ and $\delta_i = 1(T_i \le C_i)$ 
denote the overall failure time, overall censoring time, overall observed time, and overall censoring indicator, respectively. We denote $\boldsymbol Z^{(q)}_i \in \mathcal Z^{(q)}$ as the covariate information of the $i$th patient that is newly available at the beginning of the $q$th stage. Thus, $\boldsymbol H^{(q)}_i$ may include the historical information such as $A^{(k)}_i$, $B^{(k)}_i$, $\boldsymbol Z^{(k)}_i$, $k=1,2,..,q-1$ and $\boldsymbol Z^{(q)}_i$; where $d^{(q)}$ denotes the dimension of $\boldsymbol H^{(q)}_i~$ for all $i = 1,2, ..., n^{(q)}; q= 1,2,..,Q$. The number $n^{(q)}$ of patients eligible for treatment assignment at the beginning of stage $q$ is $\sum_{i=1}^n 1(U^{(q)}_i \le T^{(q)}_i \wedge C^{(q)}_i)$. The notation is summarized in Table~\ref{tab:notation} in Appendix 1.

\subsection{Overview of the method}
\label{ss:method_overview}
As is the case in Q-learning, our estimator optimizes the intermediate stage decision assuming that all the later decisions are optimal, and thus backward recursion is used. 
Our method is based on the observation that the remaining life
$L^{(q)}$ at stage $q$ is defined recursively as a convolution of $L^{(q+1)}$ and the current stage length $V^{(q)}$ and that
\begin{align}
S^{(q)}(t\mid \boldsymbol H^{(q)}, A^{(q)}, \delta^{(q)} = 1)
   =& \int S^{(q+1)}(t - V^{(q)}\mid \boldsymbol H^{(q+1)}, A^{(q+1)}, \delta^{(q)} = 1)\nonumber\\
  & \hspace{1cm} dP(V^{(q)}, \boldsymbol H^{(q+1)}, A^{(q+1)} \mid  \boldsymbol H^{(q)}, A^{(q)}, \delta^{(q)} = 1), \label{eq:marginal_nocensor} 
\end{align}
where $S^{(q)} (t\mid \cdot) = 1$ for all $t<0$ and $S^{(q)} (t\mid \gamma^{(q-1)} = 1) = 1(t <0)$, for $q=Q-1, Q-2, ..., 1.$ In other words, assuming no censoring events, the survival probability $S^{(q)}$ of the remaining life $L^{(q)}$ at stage $q$ can be obtained by augmenting the survival probability of the remaining life $L^{(q+1)}$ of the next stage by the current stage length $V^{(q)}$ and then marginalizing over all events that occur later than the beginning of stage $q$. Under the covariate-independent censoring assumption, the contribution of the censored cases is incorporated into the equation by redistributing the probability to the right \citep{efron1967}:
\begin{align}
& S^{(q)}(t\mid \boldsymbol H^{(q)}, A^{(q)})\nonumber\\
   = & \int \bigg[\delta^{(q)} S^{(q+1)}(t - X^{(q)}\mid \boldsymbol H^{(q+1)}, A^{(q+1)})  \nonumber\\
    & ~ ~ ~ ~ ~ + 
   (1-\delta^{(q)})\Big\{1(X^{(q)} > t) + 1(X^{(q)}\le t) \frac{S^{(q)}(t\mid \boldsymbol H^{(q+1)}, A^{(q+1)})}{S^{(q)}(X^{(q)}\mid \boldsymbol H^{(q+1)}, A^{(q+1)})}\Big\} \bigg]\nonumber\\
   &~ ~ ~ \times dP(X^{(q)}, \boldsymbol H^{(q+1)}, A^{(q+1)}, \delta^{(q)} \mid  \boldsymbol H^{(q)}, A^{(q)}). \label{eq:marginal} 
\end{align}
Assuming the standard causal inference assumptions, which we will formally introduce in Subsection~\ref{ss:consist2} (Assumption \ref{a:sutva}--\ref{a:positivity}), we can translate $S^{(q)}(t\mid \boldsymbol H^{(q)}, A^{(q)} = a)$ into the counterfactual quantity $S^{(q)}_a(t\mid \boldsymbol H^{(q)})$.
Thus, the regime is estimated using the following procedure: 1. estimating the conditional survival distribution 
$S^{(q)}(t \mid \boldsymbol H^{(q)}, A^{(q)})$ of the remaining life $L^{(q)}$ for $q=Q$; 2. optimizing the survival distribution over $A^{(q)}$ to have $\hat \pi^{(q)}$ and $\hat S^{(q)}_*(t \mid \boldsymbol H^{(q)})$ for $q=Q$; 3. augmenting the estimated optimal curves by the previous stage lengths, i.e., $\hat S^{(q)}_*(t - X^{(q-1)} \mid \boldsymbol H^{(q)}, A^{(q)})$ for $q=Q$; and, 4. repeating steps 1-3 for $q=Q-1,Q-2,...,1$. This procedure is summarized at a high level in Algorithm~\ref{al:dtrSurv}.


\begin{algorithm}[H]
    \label{al:dtrSurv}
    \SetAlgoLined
    \KwResult{A dynamic treatment regime estimate $\hat{\boldsymbol\pi} = (\hat\pi^{(1)},\hat\pi^{(2)},...,\hat\pi^{(Q)})^\top$}
     \For{ stage $q = Q, Q-1, ..., 1$}{
      Obtain $\hat S^{(q)}(\cdot \mid \boldsymbol H^{(q)}, A^{(q)})$ via generalized random survival forest (Algorithm \ref{al:grsf})\;
      Obtain $\hat \pi^{(q)}(\boldsymbol h) = \arg\max_{a\in\mathcal A^{(q)}} \phi^{(q)}\Big\{\hat S^{(q)}(\cdot\mid \boldsymbol H^{(q)} = \boldsymbol h, A^{(q)} = a) \Big\}$\;
      Obtain $\hat S^{(q)}_*(\cdot \mid \boldsymbol H^{(q)}) = \hat S^{(q)}(\cdot \mid \boldsymbol H^{(q)}, A^{(q)}= \hat \pi^{(q)}(\boldsymbol H^{(q)}) )$\;
     }
     
     \caption{the proposed dynamic treatment regime estimator}
\end{algorithm}

At the beginning of the last stage, there are $n^{(Q)}$ patients who have neither experienced failure nor been censored. The data for these patients are used to estimate the terminal stage covariate-conditional survival probability $S^{(Q)}(\cdot\mid \boldsymbol H^{(Q)}, A^{(Q)})$ for each treatment arm using random survival forests described in detail in Subsection~\ref{ss:grsf}. 
Then the optimal regime for the final stage is estimated by $\hat\pi^{(Q)} = \arg\max_A \phi^{(Q)}\Big\{\hat S^{(Q)}(\cdot, A)\Big\},$ where $\hat S$ is the estimate of $S$ and $\phi^{(q)}, q=2,..,Q$ is defined adaptively for the intermediate and terminal survival probabilities. Specifically, $\phi_\mu^{(q)}\Big\{S^{(q)}(\cdot \mid \boldsymbol H^{(q)})\Big\} = B^{(q)} - \int_{t>X^{(q)}}(t\wedge\tau)dS^{(q)}(t\mid \boldsymbol H^{(q)})$ and $\phi_{\sigma, t}^{(q)}\Big\{S^{(q)}(\cdot\mid \boldsymbol H^{(q)})\Big\} =S^{(q)}(t - B^{(q)} \mid \boldsymbol H^{(q)})$.


We denote the estimated optimized survival distribution of the $q$th stage as $\hat S^{(q)}_* = \hat S^{(q)}\{\cdot, \hat\pi^{(q)}(\cdot)\}$.  The optimized survival information of the final stage is carried back to the previous stage. Specifically, the $i$th patient who is available at the $(Q-1)$st stage and has a stage length $X_i^{(Q-1)}$ is given a probabilistic augmentation so that the resulting survival probability of remaining life at the second to the last stage is $\hat S^{(Q-1)}_*(t - X_i^{(Q-1)}\mid \boldsymbol H^{(Q)}_i)$. For those who have already experienced failure or censoring by the end of the stage, no augmentation is required, as either their remaining life has been determined already during the stage or the censoring indicator and the stage length contain all the necessary information.

Next, the conditional distribution of remaining life from the beginning of the $(Q-1)$st stage is estimated by the random survival forests based on the $n^{(Q-1)}$ patients. The distinction between the last stage and all the other stages is that given data, the patients' remaining life conditioning on the current historical information is observed in a scalar form at the last stage, while the remaining life for the other stages is not observed but is given in the form of a stochastic process. So a modification to the random survival forest estimator is required and is discussed in Subsection~\ref{ss:grsf}. Then estimation of $S^{(q)}, \pi^{(q)}, $ and $S^{(q)}_*$ is recursively performed until we reach the first stage and obtain $\hat {\boldsymbol \pi} = (\hat \pi^{(1)}, \hat \pi^{(2)}, ... , \hat \pi^{(Q)})^\top$.  

\subsection{A few aspects of the optimization: Backward recursion and composite values}
\label{ss:backward}
We discuss a few aspects of the optimization problem.
First, we go over why backward recursion can be used in optimizing the truncated survival mean values ($\phi_\mu$) or survival probabilities ($\phi_{\sigma,t}$).
For some criteria, such as the median $\phi_{med}(S) = S^{-1}(0.5)$, this backward recursion does not guarantee an overall optimum. 
However, for the mean survival time $\phi_\mu$ and the survival probability $\phi_{\sigma,t}$ at time $t$, backward recursion is a legitimate method according to Proposition~\ref{t:backward}, which is proved in Appendix 2.

\begin{proposition}[Legitimacy of backward recursion]
\label{t:backward}
Pick either $\phi_\mu$ or $\phi_{\sigma,t}$ for $\phi$.
Let $\{\boldsymbol\pi\}$ be an arbitrary class of decision rules and
 $\boldsymbol\pi_* = (\pi^{(1)}_*,\pi^{(2)}_*,...,$ $\pi^{(Q)}_*)^\top$ 
satisfy  $\phi\Big\{S(\cdot\mid \boldsymbol H^{(1)}, \boldsymbol\pi_*)\Big\} \ge \phi\Big\{S(\cdot\mid \boldsymbol H^{(1)}, \boldsymbol\pi)\Big\}$ for all $\boldsymbol\pi\in \{\boldsymbol\pi\}.$
Then for any $q=1,2, ..., Q-1,$
$$\pi^{(q)}_*(\boldsymbol H^{(q)}) = \arg\max_{a^{(q)}\in\mathcal A^{(q)}}\phi^{(q)}\Big\{S^{(q)}(\cdot\mid \boldsymbol H^{(q)}, \pi^{(q+1)}_*,\pi^{(q+2)}_*,...,\pi^{(Q)}_*)\Big\}.$$
\end{proposition}

Second, we study an alternative to the $\phi_{\sigma,t}$ criterion.
When $\phi_{\sigma, t}(S) = S(t)$ is chosen as the optimization criterion, the treatment rule can not discriminate treatment arms at any point later than $t$. For example, when an investigator is interested in the six month survival probability and a patient takes her third round of treatment seven months after her onset of the disease, the optimal dynamic treatment rule could recommend any treatment option randomly, since it does not affect the patient's six month survival probability. Thus a composite criterion of first maximizing $S(t)$ and second maximizing $E[T\wedge \tau]$ could be more beneficial than the simple criterion $\phi_{\sigma, t}$. The composite criterion values $\phi_{\sigma, \mu, t}(S) = \{S(t), E[T\wedge \tau]\}$ are ordered lexicographically; i.e., $\phi_{\sigma, \mu, t}(S_{\boldsymbol \pi_1}) > \phi_{\sigma, \mu, t}(S_{\boldsymbol \pi_2})$ if and only if either of the following is true:
\begin{align*}
    & S_{\boldsymbol \pi_1}(t) > S_{\boldsymbol \pi_2}(t), & \text{or} \\
    & S_{\boldsymbol \pi_1}(t) = S_{\boldsymbol \pi_2}(t) \text{ and } E[S_{\boldsymbol \pi_1}\wedge \tau] > E[S_{\boldsymbol \pi_2}\wedge \tau].
\end{align*}


\subsection{A generalized random survival forest}
\label{ss:grsf}

We develop a generalized random survival forest that takes survival probabilities as the outcome variable. The high level algorithm is provided in Algorithm~\ref{al:grsf}.
Suppose we are interested in estimating the marginal survival probability of a node $\mathcal N$ at stage $q$, where each individual's failure time information is given by the survival probability of the remaining life and the censoring status: $\{(S_i^{(q)}, \delta_i^{(q)})\}_{i=1}^{n_{\mathcal N}}$ where $n_{\mathcal N}$ is the size of  node $\mathcal N$ and $S_i^{(q)}$ is shorthand for $\hat S^{(q+1)}(t - X^{(q)}_i\mid \boldsymbol H^{(q+1)}_i)$. Then we modify the Kaplan-Meier estimator as 
\begin{align}
\hat S^{(q)}(t\mid \mathcal N) &= \prod_{s\ge 0}^t \left\{1 + \frac{\sum_{i\in\mathcal N} \delta_i^{(q)} dS_i^{(q)}(s)}{\sum_{i\in\mathcal N} S_i^{(q)}(s^-)} \right\},
\label{eq:km}
\end{align}
where $s^- = s - \epsilon$ with an arbitrarily small positive value $\epsilon$.
This modified Kaplan-Meier estimator will be used in both tree partitioning and terminal node survival probability estimation.



\begin{algorithm}[H]
\label{al:grsf}
\SetAlgoLined
\KwResult{generalized random survival forest for the $q$th stage $\hat S^{(q)}(\cdot\mid \boldsymbol H^{(q)}, A^{(q)})$}
 Parameters: the number of trees $n_{\text{tree}}$, the minimum terminal node size $n_{\text{min}}$, the minimum number of events in terminal nodes $n_{\text{min.E}}$\;
 Input data: 
 $\{(X^{(q)}_i, \gamma^{(q)}_i, \delta^{(q)}_i, \boldsymbol H^{(q)}_i, A^{(q)}_i)\}_{i=1}^{n^{(q)}}$, ~ $\hat S^{(q+1)}_*(\cdot \mid \boldsymbol H^{(q+1)})$\;
 Stochastic augmentation: $S^{(q)}_i = \hat S^{(q+1)}_*(t-X^{(q)}_i\mid \boldsymbol H^{(q+1)} = \boldsymbol h^{(q+1)}_i)$
 with $S^{(q)}_i(t) = 1(t\le X^{(q)}_i)$ if $q=Q$ or if $\delta^{(q)}_i = 0$\;
 \For{ treatment arm $a \in \mathcal A^{(q)}$}{
     \For{ tree $b = 1, 2, ..., n_{\text{tree}}$}{
      Recursively partition the feature space with the augmented data $\{(S^{(q)}_i, \delta^{(q)}_i, \boldsymbol H^{(q)}_i)\}_{\{i: a^{(q)}_i = a\}}$
      via LR (\ref{eq:lr}) or MD (\ref{eq:md}) until node size $< 2 n_{\text{min}}$ or number of failures in each node $< 2 n_{\text{min.E}}$\;
      Obtain the terminal node survival probability $\hat S^{(q)}_{b,a}(t\mid \boldsymbol H^{(q)})$ via modified Kaplan-Meier (\ref{eq:km})\;
     }
     Obtain the random survival forest estimator $\hat S^{(q)}(\cdot\mid\boldsymbol H^{(q)}, A^{(q)} = a) =\frac 1 {n_{\text{tree}}} \sum_{b=1}^{n_{\text{tree}}} \hat S^{(q)}_{b,a}(t\mid \boldsymbol H^{(q)})$\;
 }
 \caption{the generalized random survival forest estimator}
\end{algorithm}

Now we develop the splitting rules. Many random survival forest methods use the two-sample log-rank test statistic to partition the population into subgroups with a homogeneous survival distribution \citep{rsf, hothorn2005}. When failure time is not observed as a scalar but is given as a probabilistic process, the log-rank test can be generalized to the generalized log-rank statistic \citep{icrf};
\begin{align}
    \label{eq:lr}
    LR(\mathcal N_1, \mathcal N_2) = \frac{\int_0^\tau \frac{Y_2(t){dN_1(t)}+ Y_1(t){dN_2(t)}}{Y(t)}}
    {\sqrt{\int_0^\tau \frac{Y_1(t)Y_2(t)dN(t)( Y(t)- dN(t))}{Y(t)^3}}},
\end{align}
where $\mathcal N_1$ and $\mathcal N_2$ are two candidate daughter nodes, $Y_l(t) = \frac{1}{|\mathcal N_l|}\sum_{i\in \mathcal N_l} S^{(q)}_i(t^-), N_l(t) = 1 - \frac{1}{|\mathcal N_l|}\sum_{i \in \mathcal N_l} S_i(t), l = 1, 2$, $Y(t) = \lambda_{1} Y_1(t) + \lambda_{2} Y_2(t),$ $N(t) = \lambda_{1} N_1(t) + \lambda_{2} N_2(t),$ $\lambda_{l} = \frac{|\mathcal N_l|}{|\mathcal N_1| + |\mathcal N_2|},$ and $|\mathcal N| = \sum_{i=1}^{n^{(q)}} 1(i \in \mathcal N)$. The generalized log-rank splitting rule chooses $(\mathcal N_1, \mathcal N_2)= \arg\max_{\mathcal N_1, \mathcal N_2}LR(\mathcal N_1, \mathcal N_2)$ as the best partition.

When the optimization criterion is the truncated mean survival time $\phi_\mu$, splitting can also be done by finding the split point that maximizes the difference of the truncated mean survival times. Define \begin{align}
    \label{eq:md}
    MD(\mathcal N_1, \mathcal N_2) = \bigg|\int_{t>0} (t\wedge \tau) d\hat S(t\mid \mathcal N_1) - \int_{t>0} (t\wedge \tau) d\hat S(t\mid \mathcal N_2)\bigg|,
\end{align}
where $\hat S$ is the modified Kaplan-Meier estimator defined in Equation~(\ref{eq:km}). The mean difference splitting rule finds $(\mathcal N_1, \mathcal N_2)= \arg\max_{\mathcal N_1, \mathcal N_2}MD(\mathcal N_1, \mathcal N_2)$ as the best split.


The recursive partitioning continues until each terminal node size is less than a predefined threshold $k$. At each terminal node $\mathcal N$ of the $b$th survival tree, the survival probability estimate $\hat S^{(q)}_b(t\mid \mathcal N)$ is given by the modified Kaplan-Meier estimator in (\ref{eq:km}), $b=1,2,...,n_{\text{tree}}$, where $n_{\text{tree}}$ is the number of trees in the forest. The random survival forest estimate is then given by
\begin{align*}
   \hat S^{(q)}(t\mid \boldsymbol h^{(q)}) &= \frac 1{n_{\text{tree}}}\sum_{b=1}^{n_{\text{tree}}}
   \sum_{l=1}^{\#\{\mathcal N_{b,l}\}_l} 1(\boldsymbol h^{(q)} \in \mathcal N_{b,l}) \hat S^{(q)}(t\mid \mathcal N_{b,l}),
\end{align*}
where $\#\{\mathcal N_{b,l}\}_l$ is the size, or the number of terminal nodes, of the $b$th tree.



\section{Theoretical properties}
\label{s:theory}
We show consistency of the proposed dynamic treatment regime estimator $\hat {\boldsymbol \pi}$ using empirical process theory. We first derive that the generalized random survival forest is a uniformly---over the feature space---consistent estimator. The uniformity is required since estimating the optimal treatment policy of the $q$th stage relies on integrating the estimated survival curve of ($q+1$)st stage across the newly accrued information $(\boldsymbol H^{(q+1)}, T^{(q)}, U^{(q)}) \mid \boldsymbol H^{(q)}$ between the two stages. Then we show that the value of the estimated regime is consistent for the value of the optimal regime, where the versions of the value are defined in Subsection~\ref{ss:consist2}.

\subsection{Consistency of the generalized random survival forests}
\label{ss:consist1}
We derive consistency of the generalized random survival forest outside of the context of dynamic treatment regimes. We have the following assumptions on data distribution that are commonly used to achieve consistency of random forests or random survival forests \citep{meinshausen2006, wager2015, cui2017, wager2018}. These assumptions are tailored for use later in the dynamic treatment regime setting by adding stage indicators in Subsection \ref{ss:consist2}.


\begin{assumption}[Lipschitz continuous survival probability and censoring density]
\label{a:lipschitz}
There exist constants $L_S$ and $L_G$ such that
$|S(t\mid \boldsymbol h_1)-S(t\mid \boldsymbol h_2)|\le L_S \|\boldsymbol h_1-\boldsymbol h_2\|_1$ and $|g(t\mid \boldsymbol h_1)-g(t\mid \boldsymbol h_2)|\le L_G \|\boldsymbol h_1-\boldsymbol h_2\|_1$ for all $\boldsymbol h_1, \boldsymbol h_2\in \mathcal H$ and $t\in[0,\tau]$, where $G$ is the censoring survival distribution and $g$ is its derivative with respect to time.
\end{assumption}

\begin{assumption}[Weakly dependent covariate values]
\label{a:covariate}
The covariate information ($\boldsymbol Z$) is given as a $d_Z$-dimensional vector residing in a unit hypercube, i.e., $\boldsymbol Z\in [0, 1]^{d_Z}$, with a density $f_{\boldsymbol Z}$ such that $\zeta^{-1} \le f_{\boldsymbol Z} (\boldsymbol z) \le \zeta$ for all but countably many points $\boldsymbol z\in [0,1]^{d_Z}$ and some constant $\zeta \ge 1$.
\end{assumption}

\begin{remark}
The distributional assumption was made for notational convenience and could be relaxed. For example, the results in this paper still hold for covariates that are a multivariate normal random vector truncated by some constants since random forests are invariant to scale transformation. Also the results are still valid with categorical covariates with a finite number of levels after modification to Assumption \ref{a:lipschitz}: in this setting, the norm is calculated by excluding the categorical covariate coordinates. The case of categorical covariates will be discussed in the proof of Theorem \ref{t:uc}.
\end{remark}


We also have assumptions regarding the tree-based estimator terminal node size and splitting rules. The first assumption forces the terminal node size to be both sufficiently small and sufficiently large.

\begin{assumption}[Terminal node size]
\label{a:terminalnode}
The maximum size $k$ of the terminal nodes grows at the following rate:
$$k \asymp n^\beta, \frac 1 2 < \beta < 1,$$
where $a \asymp b$ implies both $a = \mathcal O(b)$ and $b = \mathcal O(a)$.
\end{assumption}

Next, we give definitions of regular trees and random split trees that are often used for consistency proofs in the random forest literature \citep{meinshausen2006, wager2015, cui2017, wager2018}.

\begin{definition}[random-split and $\alpha$-regular trees and forests]
\label{def:random}
A tree is called a \textit{random-split} tree if each feature is given a minimum probability ($\varphi/d$) of being the split variable at each intermediate node, where $0<\varphi<1$. A tree is \textit{$\alpha$-regular}, if every daughter node has at least $\alpha$ fraction of the training sample in the corresponding parent node. A random forest is called a \textit{random-split}  (\textit{$\alpha$-regular}) forest, if each of its member trees is \textit{random-split}  (\textit{$\alpha$-regular}).
\end{definition}

\begin{assumption}[$\alpha$-regular and random split trees]
\label{a:split}
Trees are $\alpha$-regular and random-split with a constant $0<\varphi<1$.
\end{assumption}



Now we give our uniform consistency results of the generalized random survival forests. This theorem is largely based on the Z-estimation lemma and Donskerness of the tree kernels. A detailed proof is relegated to Appendix 3.

\begin{theorem}[Uniform consistency of a random survival forest]
\label{t:uc}
Suppose Assumptions \ref{a:lipschitz}-\ref{a:split} hold. Let $\tilde S'$ be an estimator of $S'$ that is uniformly consistent in both $\boldsymbol h\in \mathcal H$ and $t\in [0,\tau]$, where $S'$ satisfies $S(t\mid \boldsymbol h) = \int S'(t - (T\wedge U)\mid \boldsymbol H')dP(T,U,\boldsymbol H'\mid \boldsymbol H = \boldsymbol h)$. Then, the generalized random survival forest $\hat S(t\mid \boldsymbol h)$ built based on $\tilde S'$ is uniformly consistent. Specifically,
$$\sup_{t\in[0,\tau], \boldsymbol h\in\mathcal H}|\hat S(t\mid \boldsymbol h) - S(t\mid \boldsymbol h)|\to 0,$$ in probability as $n\to\infty.$
\end{theorem}

Consistency of a random survival forest sequence in a dynamic treatment regime setting requires adaptive versions of Assumptions \ref{a:lipschitz}-\ref{a:terminalnode} and an additional condition that guarantees an arbitrarily large sample size for the last stage with high probability. To achieve this condition, we assume that, at each stage except the last one, each patient receives the next treatment before experiencing failure or censoring with at least a certain minimum probability.

\begin{assumption}[Lipschitz continuous survival probability and censoring density at each stage]
\label{a:lipschitz2}
There exist constants $L_S$ and $L_G$ such that
$|S^{(q)}(t\mid \boldsymbol h_1)-S^{(q)}(t\mid \boldsymbol h_2)|\le L_S \|\boldsymbol h_1-\boldsymbol h_2\|$ and $|g^{(q)}(t\mid \boldsymbol h_1)-g^{(q)}(t\mid \boldsymbol h_2)|\le L_G \|\boldsymbol h_1-\boldsymbol h_2\|$ for all $\boldsymbol h_1, \boldsymbol h_2\in \mathcal H^{(q)}$, $t\in[0,\tau - B^{(q)}]$, and $q=1,2,...,Q$, where $G^{(q)}$ is the censoring survival distribution at the $q$th stage and $g^{(q)}$ is its derivative with respect to time.
\end{assumption}


\begin{assumption}[Weakly dependent historical information at each stage]
\label{a:covariate2}
The patient history information $\boldsymbol H^{(q)}$ at stage $q$ is given as a $d^{(q)}$-dimensional vector lying in a subset $\mathcal H^{(q)}$ of $[0,1]^{d^{(q)}}$, and its joint and marginal densities are bounded so that $\zeta^{-1} \le f_{\boldsymbol H^{(q)}} (\boldsymbol h) \le \zeta$ and $\zeta^{-1} \le f_{\boldsymbol H^{(q)}, j} (\boldsymbol h_j) \le \zeta$ for all but a finite number of $\boldsymbol h\in \mathcal H^{(q)}$ and $j=1,2,...,d^{(q)}$ for some constant $\zeta \ge 1$.
\end{assumption}

Assumption~\ref{a:covariate2} modifies Assumption~\ref{a:covariate} so that the support of the history space can be less than a hyper-cube, since, for example, the baseline survival times, $(B^{(1)} \le B^{(2)} \le B^{(3)} \le ...)$, have a dependent support structure.


\begin{assumption}[Probability of proceeding to the next stage]
\label{a:ProbNext}
There exists a constant $M\in(0,1)$ such that for all $h^{(q)}\in \mathcal H^{(q)}$ and $q= 1,2,...,Q-1$, 
$\textnormal{pr}(U^{(q)} < T^{(q)} \wedge C^{(q)} \mid \boldsymbol h^{(q)}) > M.$
\end{assumption}

\begin{corollary}[Uniform consistency of a random survival forest sequence]
\label{t:ucSeq}
Suppose Assumptions \ref{a:terminalnode}--\ref{a:ProbNext} hold. 
Let $\hat {\boldsymbol S} = (\hat S^{(1)}, ..., \hat S^{(q)}, ..., \hat S^{(Q)})$ be the sequence of the generalized random survival forest estimators of $\boldsymbol S = (S^{(1)}, ..., S^{(q)}, ..., S^{(Q)})$ such that the $q$th stage random survival forest is built based on $\hat S^{(q+1)}$ for $q=1,2,...,Q-1$. Then, 
$$\sup_{t\in[0,\tau], \boldsymbol h\in\mathcal H, q \in\{1,2,...,Q\}}|\hat S^{(q)}(t\mid \boldsymbol  h) - S^{(q)}(t\mid \boldsymbol h)|\to 0,$$ in probability as $n\to\infty.$
\end{corollary}

The proof of Corollary~\ref{t:ucSeq} is given in Appendix 3.

 
\iffalse
    The optimal treatment rule is found by comparing estimated survival curves with different treatment arms based on a certain criterion $\phi$. Although the criterion value $\phi(\hat S_*) $ of the estimated optimal survival probability may converge, the estimated optimal survival probability itself may not, since we assume discrete treatment arm spaces. 
    In Definition~\ref{def:metric} below, we weaken the uniformity of the survival estimator to a $\phi$-metric or $\phi$-norm, and provide Lemma~\ref{t:phi_c}
    
    \begin{definition}[A $\phi$-metric]
    \label{def:metric}
    A $\phi$-metric is defined as $d_\phi(S, \tilde S) = \sup_{h\in\mathcal H}|\phi(S) - \phi(\tilde S)|$ for some functional $\phi$. An estimator $\tilde S_n$ of $S$ is called $\phi$-consistent, if $d_\phi(S, \tilde S_n) \to 0$ in probability (or almost surely), as $n\to \infty$. We further define the $\phi$-norm as $\|\cdot\|_\phi = \sup_{h\in\mathcal H}|\phi(\cdot)|$ for $\phi = \phi_\mu, \phi_{\sigma, t}$.
    \end{definition}
    
    \begin{lemma}[$\phi$-consistency of the generalized survival random forests]
    \label{t:phi_c}
    Suppose Assumptions 1-\ref{a:split} hold. Let $\tilde S'$ be an estimator of $S'_0$ that is $\phi$-consistent for $\phi = \phi_\mu, \phi_{\sigma, t}$, and $t\in[0,\tau]$. 
    Then, the generalized random survival forest $\hat S(t\mid h)$ built based on $\tilde S'$ is $\phi$-consistent.
    \end{lemma}
    The proof is provided in the Appendix. \textcolor{red}{TBD!!!}
\fi


\subsection{Consistency of the dynamic treatment regime estimator}
\label{ss:consist2}
We show consistency of the dynamic treatment regime estimator using the causal inference framework. With this framework, the theoretical properties are applicable to broader settings than sequentially randomized experiments. A counterfactual outcome is defined as an outcome that would have been obtained if a person had a treatment option contrary to his/her actual treatment. To denote the counterfactual outcomes pertaining to a counterfactual treatment or treatment policy, we add a subscript to the corresponding random variables or survival functions. For example, $S_{\pi}^{(q)}(t|\boldsymbol H^{(q)})$ is a counterfactual survival probability of the remaining life if treatment was given according to a treatment rule $\pi$ to a person with history $\boldsymbol H^{(q)}$ at stage $q$.
We assume the following three standard causal inference assumptions \citep{hernan2020, rubin2005, cole2009} known as the stable unit treatment value assumption, sequential ignorability, and positivity.
\begin{assumption}[Stable unit treatment value assumption, or SUTVA]
\label{a:sutva}
Each person's counterfactual dynamics, such as failure time, time to next treatment, and censoring, are not affected by the treatments or history of other patients. Moreover, each treatment has only one version, or if there are multiple versions, their differences are irrelevant. 
\end{assumption}

\begin{assumption}[Sequential ignorability]
\label{a:ignorable}
Treatment assignment is given independently of the counterfactual outcomes, conditionally on the individual's historical information, where the outcomes include all future random quantities such as failure time, time to next treatment, and censoring time. 
\end{assumption}

\begin{assumption}[Positivity]
\label{a:positivity}
For each stage $q=1,2,...,Q$, given historical information, the probability of having each treatment $a\in \mathcal A^{(q)}$ is greater than a constant $L> 0$. Or,
$$\textnormal{pr}(A^{(q)} = a|\boldsymbol H^{(q)} = \boldsymbol h) >L $$
for all $a\in\mathcal A^{(q)}$, $\boldsymbol h$ such that $\textnormal{pr}(\boldsymbol H^{(q)} = \boldsymbol h) > 0.$
\end{assumption}

The following theorem states that the dynamic treatment regime estimator has a value consistent for the value of the optimal regime. We define the value $\Phi$ of a treatment regime $\boldsymbol\pi$ as the criterion value of the survival probability if all individuals in a population follow the treatment regime. We use $\Phi(\boldsymbol\pi) = \phi(S_{\boldsymbol\pi})$ to denote the value of $\boldsymbol \pi$, and $\Phi$ inherits the subscript of $\phi$.

\begin{theorem}[Consistency of the dynamic treatment regime estimator]
\label{t:regime}
Suppose Assumptions~\ref{a:lipschitz2}--\ref{a:positivity} hold, and the optimal dynamic treatment regime estimator ($\hat{\boldsymbol \pi}$) is built based on the generalized random forest grown under Assumptions~\ref{a:terminalnode} and ~\ref{a:split}. 
Then, for each $\phi = \phi_\mu, \phi_{\sigma, t}$,
$$|\Phi(\boldsymbol \pi) - \Phi(\hat{\boldsymbol\pi})|\to 0,$$ in probability as $n\to\infty.$
\end{theorem}
The proof is provided in the Appendix 4.



\section{Simulations}
\label{s:sim}

\subsection{Simulation settings}
\label{ss:sim_setting}
We perform simulations to study the performance of the proposed method and compare those results to that of existing methods. The overall scheme of the simulations follows the pattern in \citet{goldberg2012}, where they mimicked a cancer clinical trial using the tumor size and wellness dynamics. We adapt their data generating mechanism so that dependence among censoring, tumor size, and wellness is incorporated in the mechanism, and deterministic processes are relaxed into more realistic ones.

In the hypothetical trial, patients can go through up to $Q=3$ treatment rounds during the trial with length of the trial, $\tau = 10$. Each patient has tumor size $\phi(t)$ and wellness $\omega(t)$ at time $t\in[0,10]$ that affects the timing of failure or treatment. We use $[t^{(q)}, t^{(q+1)}]$ to denote the $q$th stage. For each stage, patients are given either a more aggressive ($A=1$) or a less aggressive ($A=0$) treatment.

At the beginning of the trial, patients have uniformly distributed tumor sizes and wellness: $\rho(0) \sim U(0.5, 1), \omega(0) \sim U(0.5, 1)$. At each treatment stage, the more aggressive treatment immediately reduces the tumor size more than the less aggressive one. However, the more aggressive treatment worsens patient wellness more than the less aggressive one:
\begin{align*}
\rho(t^{(q)}+) & = \frac{\rho(t^{(q)})}{\omega(t^{(q)}) \{10 - 6 A\}} \vee 0\\
\omega(t^{(q)}+) & =\omega(t^{(q)}) - 2^{A-2}.
\end{align*}

After each treatment ($t^{(q)} < t \le t^{(q+1)}$), the tumor size and wellness follow the Ornstein-Uhlenbeck processes with negative dependence between the two:
\begin{align*}
\frac{d\chi(t)}{dt} = -2\chi(t) + \eta_{\chi}(t),
\end{align*}
where $\chi = \rho, \omega$ and the second term obeys a drifted Wiener process such that
\begin{align*}\frac{d}{dt}\begin{pmatrix} \eta_\rho(t) \\ \eta_\omega(t)\end{pmatrix} \sim N\Bigg(\begin{pmatrix} 0.1\\ 0.05 \end{pmatrix}, \begin{pmatrix} 1 & -0.5 \\
-0.5 & 1\end{pmatrix}\Bigg).
\end{align*} The positive means in the normal distribution imply that on average both tumor size and wellness grow. When tumor size is greater than 1, the patient receives the next treatment. In other words, $t^{(q+1)} = \min_t \arg\max_{t > t^{(q)}} 1(\omega(t) > 1)$, for $q=1,2,..,Q-1$. The patient failure time is governed by the hazard process:
\begin{align}
    \label{eq:sim_failure}
    \lambda_F(c) & = \frac 1 5\Bigg\{\frac{\rho(u)}{\omega(u)}e^{\boldsymbol g(H^{(q)})^\top \boldsymbol \beta_F(A^{(q)})}  + 1\Big(\omega(u) < 0.1\Big)\Bigg\},
\end{align}
where $g(H^{(q)}) = \{1, \log(B^{(q)} + 1), {\boldsymbol Z^{(q)}}^\top\}^\top$, 
$\boldsymbol \beta_F(a)$ is the hazard ratio for the failure time given treatment $a$ that is defined in Table~\ref{tab:scenario}, and
the indicator accelerates the failure time when the wellness drops below the critical level, 0.1. The baseline covariates at the $q$th stage follow a $p-$variate normal distribution conditioning on the previous stage value: $\boldsymbol Z^{(q)} = \frac 1 2 \boldsymbol Z^{(q-1)} + \frac 1 2 \boldsymbol Z_0^{(q)}$, where $\boldsymbol Z_0^{(q)} \sim N_p(\boldsymbol 0_p, 0.8 I_p + 0.2 \cdot 1_p^p), q=1,2,...,Q$, $1_p^p$ is the $p$-dimensional square matrix of ones, and $\boldsymbol Z^{(0)} = \boldsymbol 0_p$.

For each stage, the censoring occurs according to the hazard function,
\begin{align}
    \label{eq:sim_censor}
    \lambda_C(u) & = e^{\boldsymbol g(H^{(q)})^\top \boldsymbol \beta_C(A^{(q)})},
\end{align} 
where the censoring hazard ratio, $\boldsymbol \beta_C(a)$, is also defined in Table~\ref{tab:scenario}.


\begin{table}
    \def~{\hphantom{0}}
    \caption{The failure time and censoring hazard coefficients and dimensions of covariates in Scenarios 1 to 4.}{%
    \scriptsize
        \begin{tabular}{ccccc}
          \hline
          & scenario 1 & scenario 2 & scenario 3 & scenario 4\\
          & (base)  & (small $p$) & (large $p$) & (moderate censoring)\\
          \hline
          $p$ & 5& 2& 10 &5\\
         
          $\boldsymbol\beta_F(1)$ & 
          $(0, -3, 2,2,1, -1,-1)^\top$& 
          $(0, -2, 2, 0)^\top$&
          $(0, -2, \boldsymbol 2_4^\top,- \boldsymbol 1_4^\top, -\boldsymbol 2_2^\top)$&
          $(0, -3, 2,2,1, -1,-1)^\top$\\
         
          $\boldsymbol\beta_F(0)$ & 
          $(0, -1, 1,1,1, 1,1)^\top$&
          $(0, -1, 1, 1)^\top$&
          $(0, -1, \boldsymbol 1_4^\top,\boldsymbol 1_4^\top, \boldsymbol 0_2^\top)$&
          $(0, -1, 1,1,1, 1,1)^\top$\\
         
          $\boldsymbol\beta_C(1)$ & 
          $(-3, .2, \boldsymbol{.2}_3^\top, -\boldsymbol{.2}_2^\top)^\top$&
          $(-3, .2, .2, -.2)^\top$&
          $(-3, .2, \boldsymbol{.2}_4^\top, -\boldsymbol{.2}_3^\top, \boldsymbol 0_3^\top)^\top$&
          $(-2, .2, \boldsymbol{.2}_3^\top, -\boldsymbol{.2}_2^\top)^\top$\\
         
          $\boldsymbol\beta_C(0)$ & 
          $(-3, .1, \boldsymbol{.2}_3^\top,\boldsymbol{.2}_2^\top)^\top$&
          $(-3, .1, .2, .2)^\top$&
          $(-3, .1, \boldsymbol{.2}_4^\top,\boldsymbol{.2}_3^\top, \boldsymbol 0_3^\top)^\top$&
          $(-2, .1, \boldsymbol{.2}_3^\top,\boldsymbol{.2}_2^\top)^\top$\\
          \hline
        \end{tabular}
    }
    \label{tab:scenario}
	\footnotesize{
        $\boldsymbol\beta_F(\cdot)$ and $\boldsymbol\beta_C(\cdot)$, the log hazard ratios of the counterfactual failure and censoring times in (\ref{eq:sim_failure}) and (\ref{eq:sim_censor}), respectively.
	}
\end{table}

The treatment is allocated according to a Bernoulli trial with propensity 
\begin{align}
\label{eq:propensity}
\pi(\boldsymbol H^{(q)}) = \Big\{1 + \exp(-g(\boldsymbol H^{(q)})^\top \boldsymbol \beta_\pi)\Big\}^{-1},
\end{align}
where
the coefficient is $\boldsymbol \beta_\pi = (-1, 1, -\frac 1 2 \boldsymbol 1_p^\top)^\top$ for the observational data setting and $\boldsymbol \beta_\pi = \boldsymbol 0_{p+2}$ for the randomized trial setting. 


Simulations are run with two sample sizes, $n = 300$ and $1000$. Thus, we have in total $16$ distinct factorial settings (four hazard/censoring scenarios, two propensity designs, and two sample sizes). For each setting, we have $n_{rep} = 200$ simulation replicates, and we estimate the optimal treatment rule based on both the $\phi_\mu$ and $\phi_{\sigma,t=5,\mu}$ criteria. We compare the proposed method with the methods in \citet{goldberg2012} and \citet{simoneau2019}. We further compare the results with the observed policy characterized by (\ref{eq:propensity}) and the estimated optimal zero-order model, where all patients are given the same embedded treatment regime, which is the best on average over all embedded regimes. In particular, in implementing the \citet{goldberg2012} method, the linear $Q$-function space and the random forest-based approximation space were separately used. The method of \citet{huang2014} is not included in these simulations because \cite{simoneau2019} can be viewed as a doubly robust extension of their method. After estimation, the values---either the expected truncated survival time or the survival probability at $t=5$---are evaluated by generating an uncensored sample of size $n_{eval} = 10000$ according to the estimated policies. The simulations were complete using \texttt{R} package \texttt{dtrSurv}. For reproducible research, the code is available at https://github.com/hunyong/survQlearn.

\subsection{Simulation results}
\label{ss:sim_result}
In Figure~\ref{fig:sim1} LEFT, the simulation results show how the truncated mean survival time ($\phi_\mu$) of each policy behaves under different settings. For most of the settings, on average, all estimated regimes have greater values than the zero-order model, and the zero-order model has a higher value than the observed policy. This implies that the standard of care improves the overall survival time of patients over the observed policies and that individualized treatment rules can further enhance the outcomes.

Figure~\ref{fig:sim1} LEFT also shows that among the individualized treatment rules, the proposed method outperforms the other methods in many settings. Even when the method exhibits lower or about the same performance as the other methods for the small sample size, it often has higher performance than the others for the larger sample size. See, for example, the high censoring cases (the last row) of either the observational or the RCT setting. In the settings where the dimension of the covariate is high (the third row), the method of \citet{simoneau2019} often has high values, although sometimes estimation is not possible. This might be because of the doubly robustness property of the method and the fact that the true optimal decision rule might be close to a linear function, as suggested by the linear interaction of the stage-wide hazard functions in (\ref{eq:sim_failure}).
The method of \citet{simoneau2019}, however, often does not provide estimates when either the sample size of the terminal stage is small due to censoring or failure or the dimension of covariates is large.


Figure~\ref{fig:sim1} RIGHT shows the simulation results in terms of the survival probability at $t=5$, or $\phi_{\sigma, 5}$.
Note that the estimated regimes for the proposed method are distinct between LEFT and RIGHT but are the same for the other methods.
The proposed estimator shows better performance in terms of the policy values under most settings, giving us similar interpretations as the previous results.

\begin{figure}
\centering
\includegraphics[width = 1\textwidth]{art/Fig1.eps}
\caption{The estimated mean truncated survival time (LEFT) and the estimated survival probability at time $t=5$ (RIGHT) under each estimated policy. 
Each dot represents one simulation replicate. The box plots are the quartile summary of the points and the black dots are the average values of each method. Each row of boxes represents a data generation scenarios, and each column of boxes represents a combination of a study design and a sample size.}
\label{fig:sim1}
\end{figure}




\section{Leukemia data example}
\label{s:leukemia}
We analyze acute myeloid leukemia patient survival time using the proposed and existing methods. We use a data set collected from a clinical trial in which 210 acute myeloid leukemia patients were randomized to a frontline treatment followed by a salvage treatment that was adaptively chosen based on the patient status \citep{wahed2013, xu2016}. One of four initial treatment options was given to each patient with an equal probability:
combination of fludarabine, cytosine arabinoside, and idarubicin (FAI), combination of FAI and all-trans retinoic acid (ATRA) (FAIA), combination of FAI + granulocyte colony stimulating factor (G-CSF) (FAIG), and combination of FAI + ATRA + G-CSF (FAIAG). 
Response to the first stage treatment took one of the following forms: complete remission followed by disease progression, failure (either partial remission or remission with subsequent relapse), or resistance to treatment. Surviving patients that did not attain complete remission were given a salvage treatment based on physician judgement. The salvage treatments were labeled with one of two categories depending on whether or not they contained high dose ara-C (HDAC) or not. We assume that the salvage treatment effects are homogeneous across different versions and that the current data set contains most of the necessary information to achieve Assumption~\ref{a:ignorable}.

The data set is composed of age at entry, type of leukemia characterized by cytogentic abnormality, treatment assignment, response to the initial treatment (failure, complete remission, or resistance),
duration of each treatment stage, and survival and censoring status at each treatment stage.
Because both age and type of leukemia are known to have high power in predicting overall survival time, they are the main factors taken into account by clinicians when prescribing treatment in practice thereby mitigating the potential violation of Assumption~\ref{a:ignorable}.
Patient age at entry ranges from 21 to 87 years and is 61.1 years on average. The patients' cytogenetics abnormalities were grouped into three categories according to their prognosis \citep{wahed2013}: 65 poor, 79 intermediate, and 66 good. Out of 210 patients, 198 experienced death during the study, and 12 were censored. The overall observed survival time ($T\wedge C$) was on average 414 days and had quartiles of 45, 184, and 443 days. We set $\tau = 450$ days.
Figure~\ref{fig:dynamics} overviews the dynamics of the study.

\begin{figure}
    \centering
    \includegraphics[width = 0.7\textwidth]{art/Fig3.eps}
    \caption{Treatment stages and patient status with the number of instances in the parentheses.}
    \label{fig:dynamics}
\end{figure}

We derive the optimal treatment rules using the proposed method, the \citet{goldberg2012} method, the \citet{simoneau2019} method, and the zero-order model. The value of the observed policy---the combination of randomization and physician's judgement---is estimated and is compared with the other rules. As in Section~\ref{s:sim}, the \citet{goldberg2012} method was implemented using two different Q-function approximations: a linear model and a random forest. Because the \citet{simoneau2019} method was implemented using \texttt{R} package \texttt{DTRreg} the current version (1.7) of which allows binary treatment arms only, the method was applied after dichotomizing the initial treatment according to inclusion of all-trans retinoic acid.

We compare the performance of the five models based on cross validation, where $K = 100$ training data sets of size 168 (80\%) are sampled to estimate the treatment rules without replacement and stratified according to the initial treatment, and the remaining test sets are held out for evaluation. The value of each treatment rule is estimated by the inverse probability (censoring) weighting approach as follows.

Define 
\begin{align}
    W_i(\boldsymbol \pi) & = \frac{\prod_{q=1}^2 1(\pi^{(q)}(\boldsymbol H_i^{(q)}) = A_i^{(q)})\delta_i}{ \widehat{ \textnormal{pr}}(A^{(1)}_i)\widehat{\textnormal{pr}}(A^{(2)}_i \mid \boldsymbol H^{(2)}_i)\prod_{q=1}^2\widehat{\textnormal{pr}}(T^{(q)}_i > C^{(q)}_i\mid \text{age}_i)},
\end{align}

where the propensity of the salvage treatment is modeled as $\textnormal{pr}(A^{(2)} = a) = (1 + \exp(-\boldsymbol H^{(2)} \boldsymbol \beta))^{-1},$ $\boldsymbol H^{(2)}$ is a matrix with columns of one,  $B^{(2)}$, $A^{(1)}$, age, cytogenetic leukemia type, and response to initial treatment. Because there was an insufficient number of censored patients in the study, the censoring propensity score is modeled using only the patient age at entry. Then the value of each policy $\boldsymbol \pi$ and each criterion $\phi = \phi_\mu, \phi_{\sigma,t}$ is estimated by 
 
\begin{align*}
    \hat\phi_\mu(S^{\boldsymbol\pi}) &=\frac{\sum_{i:\text{test set}}(T_i\wedge\tau) W_i(\boldsymbol\pi)}{\sum_{i:\text{test set}} W_i(\boldsymbol\pi)}\\
    & \text{and}\\
    \hat\phi_{\sigma,t}(S^{\boldsymbol\pi}) &=\frac{\sum_{i:\text{test set}}1(T_i > t) W_i(\boldsymbol\pi)}{\sum_{i:\text{test set}} W_i(\boldsymbol\pi)},
\end{align*}
where for the second criterion, we use $t=180$ to evaluate the six month survival probabilities. We estimate the value of the observed policy, i.e., a hybrid of randomized rule for the initial treatment and the physician judgement for the salvage treatment, under the same evaluation procedure.

Figure~\ref{fig:leukemia} summarizes the value estimates of the estimated polices. Each box and scatter plot corresponds to a policy, and within each plot, each dot represents a cross-validation replicate of the estimated value (truncated mean survival time (LEFT) or six-month survival probability (RIGHT)) of the corresponding policy. The results show that the proposed method achieves higher values than the other estimated policies on average in terms of both the truncated mean survival time and the six month survival probability. This is true when compared with the zero-order model implying the advantage of the proposed individualized treatment rules over the standard of care. However, the estimated methods do not have higher values than the observed policy value on average. Recalling that the second treatment assignments were done according to physician judgement, this implies that their judgement was fairly successful and that a larger sample size and more patient information---available only to physicians but not included in this data set---could enhance the performance of the methods. Despite the gap, however, the proposed method yields the closest values to the observed policy values among the competing methods.

\begin{figure}
    \centering
        \includegraphics[width = 1\textwidth]{art/leuk.eps}
    \caption{Estimated truncated mean survival time (LEFT) and estimated six-month survival probability (RIGHT) under the dynamic treatment regimes. Each dot represents one cross-validation replicate. The box plots are the quartile summary of the points and the black dots are the average values of each method.}
    \label{fig:leukemia}
\end{figure}



\section{Discussion}
\label{s:discussion}
In this paper, we introduced a new dynamic treatment regime estimator for survival outcomes. The estimator is versatile in the sense that it allows for a flexible number of treatment stages and arms, it gives a non-linear decision rule that maximizes either the mean survival time or the survival probability, and it permits dependent censoring. In this section, we discuss some considerations when using our method and how the method can be further extended.

The proposed estimator assumes that the distribution is well characterized by the treatment stages rather than the natural passage of time. That is, the cohort at each treatment stage might include patients that received the stage treatment on day 10 as well as patients that received the stage treatment on day 1000. Thus, this method is effective when the disease dynamic is considered stationary between stage transitions or when the treatment effect contains relatively stronger signal than that of the baseline disease dynamics.

Another consideration is the choice of the study length, $\tau$, which is chosen so that the survival probabilities are reliably estimated up to $\tau$. For the stability of estimation, the chosen $\tau$ should not be too large; however, reducing $\tau$ results in information loss. If all patients are censored at a certain stage, the survival probability is not estimable. In this case, a smaller $\tau$ should be chosen so that every stage has a sufficient number of observed failure times.



One interesting extension of the proposed method is optimizing a specific quantile of the survival distribution such as the median survival time.
However, this task introduces a unique challenge with respect to the backward recursion. Specifically, the optimal decisions made for later stages may not be optimal once earlier decisions moderate the distribution. Thus, an exhaustive search may be needed to find an optimal policy. An extension of the quantile-optimal dynamic treatment regime estimators developed by \citet{wang2018} and \citet{linn2017} to right-censored data would be interesting future work.

The proposed method assumes an unrestricted policy class. In practice, however, clinicians and patients may prefer understandable and simple rules. A linear rule, for example, is often of interest. By posing a Cox-type proportional hazards assumption on top of our method, the resulting policy class becomes a set of linear functions. The task then reduces to replacing the generalized random survival forest with the generalized Cox model, where the estimated survival probabilities, instead of the observed time with censoring indicators, are entered as the outcome variable.

In establishing the uniform consistency results, we use the fact that  under the given assumptions, the maximum size of the terminal node becomes arbitrarily small in probability uniformly, as argued in \citet{wager2015}. However, the uniform rate of convergence can be very low. As a random-split tree grows deeper and has an increasingly large number of terminal nodes, there exists, with a not-very-small probability, at least one terminal node that is created without enough splits on a certain covariate. To enhance the rate, splitting can potentially be done in a stratified fashion, by assigning a certain minimum number of splits to each variable. The effect of the tree building rules on the rate of convergence would be an interesting future study especially in the presence of noise variables in high-dimensional data settings.






\section*{Acknowledgement}

The authors thank Peter F. Thall, Abdus S. Wahed, and Yanxun Xu for providing the leukemia data, and Donglin Zeng for bringing the composite criterion optimization into the discussion.
This research was funded in part by grant P01 CA142538 from the National Cancer Institute .

\section*{Supplementary material}
\label{SM}
Supplementary Material available at {\it Biometrika}\ online includes the extended proofs of Theorems \ref{t:uc} and \ref{t:regime}.


