Dynamic Treatment Regimes Using Bayesian Additive Regression Trees for Censored Outcomes

Optimizing medical therapy often requires that the treatment be tailored to the individual initially, and that the treatment be adaptive to an individual’s changing characteristics over time. Since individual responses can often be heterogeneous, it is challenging for physicians to customize treatments for individuals based on traditional clinical trial results, which lack the ability to identify subgroups that have different treatment effects and rarely consider successions of treatments. For chronic diseases that can evolve, it is even more important and difficult to choose the best therapy in sequence. To give a simple example, oncologists typically choose an initial immunosuppressant regime for patients with acute myeloid leukemia (AML) who are undergoing allogeneic hematopoietic cell transplantation (AHCT), to prevent a serious potential complication called graft-versus-host disease (GVHD). At the time that such an initial regime fails, a salvage treatment is chosen based on the patient’s prior treatments and responses. Such a multi-stage treatment decision has been summarized as a dynamic treatment regime (DTR) by Murphy (2003). Each decision rule in DTR takes a patient’s individual characteristics, treatment history and possible intermediate outcomes observed up to a certain stage as inputs, and outputs a recommended treatment for that stage.

A number of approaches have been proposed for estimating and optimizing DTRs, including those by Robins (2004), Moodie et al. 2007; Qian and Murphy 2011; Zhao et al. 2015; Krakow et al. 2017; Murray et al. 2018, and Simoneau et al. 2020. Two textbooks and an edited volume have been published on the topic of DTRs (Chakraborty and Moodie (2013); Kosorok and Moodie (2015); Tsiatis et al. (2020)), and a recent paper surveying value-search approaches by Jiang et al. (2019) was the subject of a lively discussion. Bayesian approaches have received relatively little attention in the DTR literature, though exceptions exist (Arjas and Saarela 2010; Saarela et al. 2015, 2016; Rodriguez Duque et al. 2022). However much of the Bayesian DTR methodology is relatively parametric. An exception to this is the Bayesian machine learning (BML) method developed by Murray et al. (2018), which innovatively bridges the gap between Bayesian inferences and dynamic programming methods from machine learning. A key advantage to a Bayesian approach to estimation is the quantification of uncertainty in decision making through the resulting posterior distribution. A second benefit that arises specifically in the BML approach is the highly flexible estimation that is employed, which minimizes the risk of estimation errors due to model mis-specification.

However, the BML method has not yet been adapted to censored outcomes, which is one of the more common types of outcomes in controlling chronic diseases. Motivated by the study of optimal therapeutic choices to prevent and treat GVHD, in this paper, we extend this approach to censored outcomes under the accelerated failure time (AFT) model framework. By modifying the data augmentation step in the BML method, the censored observation times can be imputed in an informative way so that the observed censoring time is well utilized. This extension is implemented using Bayesian additive regression trees (BART); we accomplished the implementation of the proposed AFT-BML approach by developing an R function that utilizes standard BART survival software directly without needing to modify existing (complex) BART software directly. Parallel computing was used to speed up the computational calculations. This R wrapper function can be easily adjusted to accommodate other types of Bayesian machine learning methods.

This paper is organized as follows. In Sect. 2, we briefly review related methods, algorithms, and describe the extended AFT-BML approach for optimizing DTRs for censored outcomes under the accelerated failure time framework. Section 3 presents simulation studies to demonstrate our model performance by comparing it to estimation using Q-learning. An analysis of our motivating dataset of patients diagnosed with AML is given in Sect. 4. Finally, in Sect. 5, we discuss the advantages and disadvantages of our approach and provide some suggestions for future work.

In this section, we applied the proposed method to a retrospective cohort study using registry data collected by the Center for International Blood and Marrow Transplant Research (CIBMTR) (Krakow et al. 2017). There are 4171 patients with complete information in this data who received graft-versus-host disease prophylaxis for their allogeneic hematopoietic cell transplant, which was used to treat their myeloid leukemia, between 1995 and 2007. Some patients were subsequently given a salvage treatment after they developed GVHD and experienced unsuccessful initial treatment. The two stages considered in this study were up-front GVHD prophylaxis treatment and salvage treatment after developing GVHD and failing initial treatment (which is consistently given as steroids). In each stage, patients were assigned one of two treatments, nonspecific highly T-cell lymphodepleting (NHTL) immunosuppressant therapy or standard prophylaxis immunosuppressant. Estimating an optimal DTR to maximize the overall disease-free survival time (DFS) for patients is our primary goal. The primary outcome is time to death, disease persistence, or relapse. The Stage 1 time is defined as the time from graft infusion to diagnosis of steroid-refractory acute GVHD (if they enter Stage 2) or to the primary outcome or last follow-up (if they do not enter Stage 2). The Stage 2 time is defined as the time from starting salvage treatment for steroid refractory GVHD to the primary outcome or last follow up. Among the 13 covariates of interest, time from graft infusion to acute GVHD onset (\(\ge 1\) month, \(<1\) month) and use of \(\ge 4\) immunosuppressors to treat acute GVHD on index form (Yes, No) are only available for those patients who failed at the first treatment. The other covariates include recipient’s age group (\(<10\) years, \(10-39\) years, \(\ge 40\) years), Karnofsky/Lansky performance status at time of transplant (\(\ge 80\%\), \(<80\%\)), disease status at time of transplant (Early, Intermediate, Advanced), donor relationship (Related, Unrelated), donor-recipient sex (female-male, other), graft source (Bone marrow, Peripheral blood, Umbilical cord), human leukocyte antigen (HLA) match (Well-matched, Partially matched, Mismatched), total-body irradiation (Yes, No), cytomegalovirus status (Negative-negative, Donor or recipient positive), conditioning intensity (Myeloblative, RIC/nonmyeloablative), and use of corticosteroids as part of GVHD prophylaxis (No, Yes). The prophylaxis assigned in Stage 1 is also used in fitting the salvage Stage 2 model. A frequency cross table of prophylaxis and salvage treatment assigned is shown in Table 1. The censoring rate in this cohort was 32%.

Both Q-learning and AFT-BML approaches were used to fit this two stage survival data DTR estimation. All the main effects and the two-way interactions between stage-wise treatment and the other covariates are included in Q-learning models, and 1000 nonparametric bootstrap resamples were generated to estimate the uncertainty of the quantities of interest. For nonparametric AFT-BML, with 1000 MCMC posterior samples, the full distribution was available for any predictions. The point estimates of parameters along with bootstrap mean and 95% confidence interval (CI) at each stage were examined for Q-learning. The waterfall plots for the mean differences in DFS on the log time scale under each treatment at each stage for each individual were created for both Q-learning and AFT-BML, as well as the 95% and 50% credible intervals (bootstrap CIs for Q-learning) presented on the same plot. The differences in the median DFS were also explored for both methods, as were the differences in the two-year DFS probabilities.

The analysis results from Q-learning, including point estimates and bootstrap mean, as well as the bootstrap 95% CI, are shown in the Appendix C in Tables 3 and 4. Inspection of the 95% CI for the interaction terms in Table 3 reveals covariate combinations that can be used to identify subgroups where the NHTL or the standard treatment is preferable. For example, assuming all the other covariates are at the reference level, an unrelated donor would benefit more from NHTL than the standard treatment at Stage 1. Similarly in Table 4, when holding the other covariates at the reference level, a patient who received NHTL at Stage 1 would be expected to have a longer DFS time if the standard treatment was given at Stage 2, since the 95% CI of A2.NHTL*A1.NHTL is negative. Intuitively, this might be the case because salvage treatment that is different than the initial treatment which already failed might be expected to be more effective.

In Fig. 4, we present the estimated treatment differences on the log time scale for each stage, along with 95% CI and 50% CI. For AFT-BML, a direct posterior prediction difference for each individual can be calculated from 1000 MCMC posterior samples based on patient-level characteristics. The credible intervals are constructed using the quantiles of posterior samples of each patient. For Q-learning, a standard error can be estimated from the predictions of 1000 bootstrap resamples at an individual level. Using the estimated standard error, the CIs for each patient are calculated in a standard way. A positive difference means NHTL is the preferred treatment for a given stage. The patients are presented in descending order based on the estimated difference, separately for each method.

The results from AFT-BML (Fig. 4a, c) suggest that there is little to be gained by individualizing the treatment in Stage 2 since everyone benefits from the standard treatment, but at Stage 1 there may be significant clinical value in choosing treatment in a personalized fashion to maximize the overall DFS time. In fact, there may be four subgroups that have a distinct difference in expected log event time, in which two groups would benefit from NHTL, one group is indifferent to treatment choice, and one final group that would have longer survival time with the standard treatment. The right panels, showing Fig. 4b, d, provides a similar message using Q-learning: while the standard treatment is the preferred treatment for most (though not all) patients at Stage 2, individualizing the treatment at Stage 1 may lead to important benefits in log DFS time. However, there is a more continuous spectrum of treatment differences in Stage 1 using Q-learning, compared to AFT-BML, and the magnitude of the differences appears larger with Q-learning. Although it is impossible to know why this is the case, we suspect that Q-learning may be overfitting the Stage 1 and Stage 2 models since we have forced all the variables of interest and treatment interaction terms into the model. This can lead to a wider variability in the estimated treatment differences. It may be possible to reduce the potential for overfitting in the Q-learning approach by using a penalized variable selection strategy in the Stage 1 and Stage 2 regression models, but we did not consider this further.

In figure 5, we present the predicted treatment differences on a different time scale, specifically the estimated differences of median DFS time with 95% CI and 50% CI at each stage. Since the event time is assumed to follow a log normal distribution, the exponential of the log normal mean, which transforms the predictions back to the original time scale, is the median rather than the mean. The results on this scale are generally consistent with the log time scale results, though the scale change produces some unusual results. There are some individuals in the middle and edges of Fig. 5a who have wider CIs than their neighbors. These larger CI widths are an artifact of the scale of the treatment difference and the corresponding ordering of the treatment differences. In particular, the wide intervals correspond to individuals whose median survival predictions are higher (under one or both treatments). When doing inference on the median survival using the exponential transformation of the model parameters, the variance increases with the median. As a result, for these cases where the median survival time under one or more treatments is large, the intervals are wider. Furthermore, because we are plotting in order of the median difference, these cases can occur in several places on the plot, as long as one or more of the medians is large. For example, the wide intervals in the middle of the plot occur when the medians under each treatment are both large and of similar magnitude, while the wide intervals on the edges correspond to individuals where the median under one of the treatments is large, but the other is of a different magnitude. If we had instead plotted in order of the median under one of the treatments, you would see the width of the intervals generally get larger with increasing median.

DTRs could also be defined as the optimal treatment rules to maximize each patient’s two-year DFS probability. Figure 6 shows the two-year DFS probability difference between NHTL and standard treatment for all patients sorted in descending order. For Stage 2, all patients are expected to have a higher two-year survival probability if assigned the standard treatment based on AFT-BML method. The Q-learning method agrees with the inferences from AFT-BML except for a very small proportion of patients, although the magnitude of the differences are much bigger for Q-learning. For Stage 1, it is easy to notice that AFT-BML splits the patients into four subgroups, similar to what we saw in Fig. 4a. With Q-learning, however, it is difficult to recognize any clear cut points on the curve. As in Stage 2, the magnitude of the differences are smaller for AFT-BML. The proportions of patients who should have NHTL as the optimal treatment for Stage 1 are consistent between AFT-BML and Q-learning.

To further compare the survival probability predictive performance of AFT-BML vs. Q-learning, we calculate the time dependent area under the ROC curve (Heagerty et al. 2000) with the R package timeROC (Blanche et al. 2013) using the predicted survival time estimated by both AFT-BML and Q-learning as predictors. For the Stage 2 model, the observed time and event indicator can be used directly in calculating the time dependent AUC. For the Stage 1 prediction model (which assumes patients receive optimal treatment in Stage 2), we need to account for not all patients receiving their optimal treatment in Stage 2. To handle this, depending on whether the estimated optimal treatment at Stage 2 was observed or not, the original observation was kept as is, or was censored at the time of entering Stage 2. Since the estimated optimal Stage 2 treatment could be different from AFT-BML to Q-learning, we examined three sets of censored Stage 1 data, including optimal Stage 2 treatment identified by AFT-BML, or Q-learning, or consistently optimal under both Stage 2 models. The time points of interest for Stage 1 are one year, two years, and three years. For Stage 2, only the median and third quartile of the observed time are evaluated. The results from both stages are shown in Table 2. For Stage 2, AFT-BML improves the AUC by \(0.55\%\) at the median, and \(1.87\%\) at the third quartile, indicating that AFT-BML has a better predictive performance at Stage 2. For Stage 1, the time dependent AUC at 1 year from AFT-BML is approximately \(1.7\%\) higher than Q-learning in all three settings. This improvement increases to around \(2.2\%\) as time goes to 2 years and 3 years. It indicates that AFT-BML once again outshines Q-learning at Stage 1 in predictive performance.

To visualize the AFT-BML based DTRs, we applied the ‘fit-the-fit’ method and plotted a single tree as in Logan et al. (2019). Since there is little value in differentiating the treatment for Stage 2, we only focus on the Stage 1 model here. Here the outcome used for the single tree fit is the posterior mean treatment difference of the log survival time, although other outcomes such as median DFS or DFS probabilities at fixed timepoints could also be used as outcomes. The \(R^2\) goodness of fit measure for using a single tree in Fig. 7 to model the Stage 1 posterior mean differences in log mean DFS time predictions is above 90%, indicating that this (highly interpretable) single tree is a reasonable representation of the original AFT-BML model at Stage 1. Values in the nodes are the posterior mean differences in mean log DFS time in months between NHTL and standard treatment. The corresponding 95% CIs are also shown in the same node. The first split (on donor type) indicates that almost all patients receiving unrelated donor transplants would benefit from receiving NHTL as GVHD prophylaxis for their AHCT, while relatively few patients receiving related donor transplants should receive NHTL. As the tree grows, the patients can be divided into four subgroups. After the first split, the two bottom nodes on the left are well apart from each other, as are the two nodes on the right side. These observations agree with Fig. 4a, that identified four subgroups with a distinct posterior mean difference in log DFS times.

In medical practice, it is of important clinical value to select the optimal treatment based on an individual’s characteristics. In many settings, a sequence of optimal treatments is desired, such as for diseases like cancer that can progress. The AFT-BML approach for identifying the DTRs can assist physicians to make sound, data-supported decisions at each stage. Classical parametric approaches, such as Q-learning, are constrained by the necessity of correctly specifying functional forms, including all the interaction terms among covariates and treatment. The Bayesian machine learning approach, in contrast, avoids this restriction and allows the model to adaptively determine the relationships between the outcome and the covariates, facilitating optimal treatment identification. We have extended the Bayesian machine learning approach to censored survival data in an AFT-BML model framework, and also provide parallelizable code for implementation of the computationally intensive algorithm. This wrapper function utilizes standard available BART software without needing to modify the complex underlying BART code. With the simulation studies, we have shown that the AFT-BML approach can achieve almost the same performance as the oracle model for censored outcomes. The results from AFT-BML not only include the optimal treatment and optimal outcomes that classical parametric approaches can provide, but also directly offers the uncertainty measurement for these targets of inference. This extra uncertainty information could be useful in practice since physicians could better assess the confidence they should have in the recommended optimal treatments.

Our comparison to Q-learning in the simulations and the example used a fixed model specification without variable selection. With a larger number of covariates, a penalized AFT model could be used for variable selection in the Q-learning approach. Furthermore, we utilized bootstrapping to estimate the uncertainty of the Q-learning predictions. Alternative approaches, such as penalized Q-learning (Song et al. 2014), could be used to directly provide the uncertainty measurements. However, we are not aware that this approach has been extended to censored data models, such as the AFT model used here.

We compared model performance between AFT-BML and Q-learning in the example by censoring the outcomes of individuals who did not receive optimal subsequent treatment. This was feasible in this dataset because most of the patients received optimal treatment in Stage 2. However, a general strategy of assessing the model performance in the dynamic treatment regimes setting warrants further investigation.

There are some limitations in our approach. We have demonstrated the BML approach for censored data using a parametric log normal AFT-BML model, which has substantial parametric model assumptions even though the functional form of the covariate effects is flexible; other types of Bayesian survival models (Sparapani et al. 2016; Henderson et al. 2020; Linero et al. 2021) could alternatively be used which may require fewer assumptions. In such cases, the methodology described here and our wrapper function can serve as a template for implementing alternative models in a BML DTR framework.

Another limitation is that our AFT-BML approach and wrapper function are currently implemented for the two-stage AFT-BART survival model. We present one possible algorithm for the general setting with more than two stages in the Appendix A, but have not yet implemented it. Modifying our current software implementation to handle more than two stages would introduce some computational challenges. The computational time would increase since we will have more layers of chain burn-in if more than two stages are present, as can be seen from the general algorithm in the appendix. This limitation could be reduced if the multi stage model fittings were embedded together, though this would lose the flexibility of applying BART functions off-the-shelf. Essentially one would need to output the tree structure and terminal node means at the end of one update, and then pass these to the next imputed dataset being analyzed as an initial tree structure and terminal node mean that is being updated.