Censoring-adjusted tree-based policy learning for estimating dynamic treatment regimes with censored outcomes

Abstract

Dynamic Treatment Regimes (DTRs) provide a systematic approach for making sequential treatment decisions that adapt to individual patient characteristics, particularly in clinical contexts where survival outcomes are of interest. We explore ways to learn effective DTRs, from observational data, using Censoring-Adjusted Tree-based Reinforcement Learning (CA-TReL) – a novel framework to address the complexities associated with censored data when estimating optimal DTRs. By enhancing traditional tree-based reinforcement learning methods with augmented inverse probability weighting (AIPW) and censoring-adjusted estimation, CA-TReL delivers robust and interpretable treatment strategies. Through extensive simulation studies and real-world validation using the SANAD epilepsy dataset, CA-TReL achieved statistically better performance relative to several state-of-the-art methods — including OWL, RWL, TBWL, ASCL, and DWsurv — across key metrics such as restricted mean survival time (RMST) and decision-making accuracy. This work represents a step forward in advancing personalized and data-driven treatment strategies across diverse healthcare settings.

Supplementary Information

The online version contains supplementary material available at https://doi.org/10.1038/s41598-026-53856-y.

Introduction

Dynamic Treatment Regimes (DTRs) offer a systematic approach to optimizing sequential treatment decisions by tailoring interventions based on individual patient characteristics and evolving clinical histories^1–3. Widely applicable in managing chronic and progressive diseases, such as cancer and diabetes, DTRs consist of stage-specific decision rules that guide treatment allocation to maximize long-term patient outcomes^4–6. We consider ways to learn DTRs from observational data. Grounded in causal inference⁷(described in the Methods : “Notation and Assumptions” subsection below), the primary goal of learning is to produce a DTR model that can personalize treatment strategies by accounting for the heterogeneity in treatment responses and, in expectation, improves outcomes across the target population.

Here, we measure the effectiveness (reward) of a policy based on the survival time of the individual. For datasets with only uncensored information – that is, where survival outcomes are fully observed for all individuals – various statistical and machine learning techniques – such as Q-learning^8–10 and A-learning^11,12 (which rely on regression-based models to iteratively estimate stage-specific decision rules) have been proposed to estimate optimal DTRs. While often effective, these methods often require strong parametric assumptions and can struggle with scalability in high-dimensional settings. To overcome these limitations, tree-based methods have emerged as a robust alternative due to their flexibility and interpretability. Among these, the Tree-Based Reinforcement Learning (T-RL) algorithm, introduced by Tao et al¹³., embeds augmented inverse probability weighting (AIPW)¹⁴ into the tree structure to estimate optimal treatment rules across multiple stages. T-RL has demonstrated excellent performance on complete datasets, particularly due to its ability to handle complex interactions between covariates and treatments while producing interpretable decision rules, leveraging the inherent transparency and simplicity of decision tree structures.

In the context of censored survival outcomes, there are already several relevant methodological approaches. On the tree-based side, Zhu and Kosorok¹⁵ introduced recursively imputed survival trees (RIST), a nonparametric method for right-censored survival regression that imputes the unobserved event times for censored individuals within a tree ensemble. However, RIST is a survival prediction method and does not estimate treatment decision rules. Building on that idea, Cui et al.¹⁶ extended outcome weighted learning (OWL)¹⁷ to censored survival data by using RIST-generated pseudo-survival rewards and then applying the usual weighted SVM objective; however, their method remains a single-stage individualized treatment rule approach rather than a multistage DTR method. Zhao et al.¹⁸ developed IPCW-based and doubly robust learning methods for censored single-stage treatment-rule estimation, while Zhou et al.¹⁹ refined OWL through residual-based weighting and a smoothed ramp loss, but did not develop a multistage DTR method for censored survival outcomes. Hager et al.²⁰ proposed a two-stage censored-survival DTR method that combines a locally efficient doubly robust augmented inverse-probability-weighted value estimator with an OWL-like weighted SVM classification framework. For multistage censored survival DTRs, Goldberg and Kosorok²¹ proposed censored Q-learning with backward recursion and inverse-probability-of-censoring weighting for flexible numbers of stages. Subsequent work expanded censored survival DTR methodology along several distinct lines. First, dynamic weighted survival modeling (DWSurv) extended the dynamic weighted ordinary least squares (dWOLS)/blip-function framework to right-censored survival outcomes by incorporating inverse-probability-of-censoring adjustment into treatment-based balancing weights and estimating stage-specific treatment effects through a sequence of weighted generalized estimating equations, and later papers focused mainly on its applied use and software implementation^22–24. Second, Zhao et al. extended the shared-parameter DTR literature by proposing censored shared-Q-learning and censored shared-O-learning, thereby adapting both recursive Q-learning and direct O-learning–style weighted classification to the censored setting²⁵. Third, Zhang et al. extended the DWSurv/blip-function line from binary to multicategory treatments by defining a reference-free blip function and estimating it via doubly robust weighted least squares with treatment and censoring weights²⁶. Finally, contrast-learning methods were adapted to censored survival outcomes through pseudo-value substitution, yielding doubly robust weighted-classification procedures with regularization and multi-treatment extensions rather than tree-based or Q-learning-based estimation²⁷.

Although these methods have substantially advanced censored-survival DTR estimation, they rely on paradigms such as tree-based imputation, weighted classification, blip-function regression, or censored Q-learning, rather than directly embedding censoring-adjusted estimation within a multistage tree-based reinforcement-learning framework. To address this gap, we propose Censoring-Adjusted Tree-Based Reinforcement Learning (CA-TReL), an extension of T-RL for right-censored outcomes, where the event time for censored individuals is only known to exceed the observed follow-up time. Building on the robustness and interpretability of T-RL, CA-TReL introduces censoring adjustments through a modified AIPW estimator that incorporates survival and censoring probabilities into the tree-building process. This ensures unbiased estimation of counterfactual outcomes while retaining the flexibility of tree-based methods. By leveraging T-RL’s strengths and addressing its limitations in handling censored data, CA-TReL provides a scalable and interpretable solution for optimizing multi-stage, multi-treatment decision-making under censoring conditions.

The remainder of this paper is organized as follows. The Results section evaluates the performance of CA-TReL through simulation studies and a real-world dataset, comparing it to baseline models including fixed treatment policies (uniform treatment assignment), random treatment assignment policies, and advanced DTR methods, demonstrating its utility in clinical decision-making under censored data conditions. The Discussion section concludes with a discussion of the method’s implications and potential extensions for future research. Finally, the methods section details the notation and assumptions underlying the CA-TReL framework, and its methodology, including the modified purity measures and recursive estimation process.

Results

To thoroughly assess the performance of the proposed method, we performed extensive experiments, of our method, baseline models including fixed treatment policies (uniform treatment assignment), random treatment assignment policies, and five advanced DTR methods — OWL¹⁷, RWL¹⁹, TBWL¹⁶, DWsurv^22,23, and ASCL²⁷– on a synthetic and real-world datasets, using four distinct evaluation metrics, as described in the Methods subsection.

Models for comparison

To evaluate the proposed CA-TReL method, we compared its performance against several approaches in both binary and multi-category treatment settings. Baseline models included fixed treatment policies, where treatments were uniformly assigned to all individuals (e.g., for g = 0, assign treatment 0 to all patients; similarly for g = 1, assign treatment 1, and so on, as used in Table 1). A random treatment assignment policy, where treatments were allocated without considering individual covariates, was also included as a simplistic benchmark.

Table 1.

Performance of several DTR algorithms. The table reports the -restricted survival time, correct decision rate at the first stage (CDR1), average correct decision across all stages (ACDR), and Expected Survival outcome. For each scenario, the best model is highlighted in bold.

Treatment options	Propensity model type	Censoring type	Method		-RMST	CDR1	ACDR	Expected survival outcome
2	True	Exponential	g = 0	61	36.99 ± 0.08	80.14 ± 0.12	66.22 ± 0.04	57.19 ± 0.05
			g = 1		17.17 ± 0.13	19.86 ± 0.02	4.99 ± 0.02	24.21 ± 0.05
			Random		25.52 ± 0.19	49.93 ± 0.92	25.06 ± 0.07	37.86 ± 0.07
			OWL		34.02 ± 1.96	40.50 ± 8.84	35.53 ± 8.17	52.58 ± 4.09
			RWL		39.09 ± 2.14	59.24 ± 10.07	49.88 ± 8.89	63.49 ± 4.49
			TBWL		33.82 ± 3.47	48.58 ± 12.57	41.85 ± 12.38	52.17 ± 6.10
			DWsurv		37.70 ± 0.60	67.43 ± 2.55	58.18 ± 2.16	59.12 ± 1.49
			ASCL		37.09 ± 2.28	59.92 ± 8.29	45.85 ± 6.90	58.94 ± 4.32
			CA-TReL		54.02 ± 0.33	94.49 ± 3.18	95.24 ± 2.35	96.1 ± 1.14
		Conditional	g = 0	58	39.77 ± 0.06	80.14 ± 0.03	68.53 ± 0.03	57.19 ± 0.07
			g = 1		14.41 ± 0.06	19.86 ± 0.03	3.91 ± 0.03	24.21 ± 0.04
			Random		25.34 ± 0.09	49.90 ± 0.05	25.10 ± 0.05	37.82 ± 0.08
			OWL		40.44 ± 3.13	66.10 ± 9.77	60.61 ± 10.20	62.24 ± 3.60
			RWL		38.71 ± 2.46	60.51 ± 8.31	51.96 ± 9.96	62.63 ± 2.87
			TBWL		34.25 ± 5.37	56.69 ± 14.35	48.35 ± 15.79	52.60 ± 6.86
			DWsurv		38.92 ± 0.87	67.71 ± 2.45	59.21 ± 2.82	59.14 ± 1.06
			ASCL		39.29 ± 0.87	65.07 ± 11.58	53.63 ± 11.60	63.28 ± 4.46
			CA-TReL		53.27 ± 0.53	91.43 ± 2.35	91.32 ± 2.06	92.50 ± 1.19
		Uniform	g = 0	5	4.77 ± 0.00	80.14 ± 0.02	69.69 ± 0.06	57.19 ± 0.06
			g = 1		4.49 ± 0.00	19.86 ± 0.02	4.30 ± 0.04	24.21 ± 0.03
			Random		4.63 ± 0.00	50.16 ± 0.10	25.08 ± 0.17	37.83 ± 0.05
			OWL		4.90 ± 0.03	47.49 ± 13.42	46.95 ± 14.37	56.58 ± 5.42
			RWL		4.98 ± 0.00	81.95 ± 5.53	80.35 ± 4.08	81.35 ± 2.92
			TBWL		4.81 ± 0.02	56.25 ± 14.58	46.58 ± 16.10	50.83 ± 7.67
			DWsurv		4.83 ± 0.01	68.41 ± 0.83	61.15 ± 1.31	58.90 ± 0.50
			ASCL		4.98 ± 0.00	86.95 ± 1.70	83.61 ± 1.91	81.28 ± 1.66
			CA-TReL		4.99 ± 0.00	87.08 ± 2.28	90.55 ± 1.93	93.03 ± 1.28
	False	Exponential	g = 0	61	36.99 ± 0.08	80.14 ± 0.02	66.22 ± 0.04	57.19 ± 0.05
			g = 1		17.17 ± 0.13	19.86 ± 0.02	4.99 ± 0.02	24.21 ± 0.05
			Random		25.52 ± 0.20	49.93 ± 0.09	25.06 ± 0.07	37.86 ± 0.07
			OWL		39.74 ± 1.45	72.46 ± 7.36	64.96 ± 7.40	63.92 ± 2.36
			RWL		40.23 ± 1.10	64.11 ± 7.84	53.36 ± 6.37	66.17 ± 2.63
			TBWL		32.30 ± 2.94	44.29 ± 14.28	35.31 ± 14.25	49.51 ± 5.68
			DWsurv		37.73 ± 0.62	67.06 ± 2.56	58.20 ± 2.22	59.12 ± 1.53
			ASCL		35.20 ± 3.16	61.60 ± 13.19	42.53 ± 9.60	55.76 ± 6.10
			CA-TReL		54.33 ± 0.33	95.08 ± 2.86	95.52 ± 2.13	96.05 ± 1.10
		Conditional	g = 0	58	39.77 ± 0.06	80.14 ± 0.03	68.53 ± 0.03	57.18 ± 0.07
			g = 1		14.41 ± 0.06	19.86 ± 0.03	3.91 ± 0.03	24.21 ± 0.04
			Random		25.34 ± 0.09	49.90 ± 0.05	25.10 ± 0.05	37.82 ± 0.08
			OWL		44.94 ± 1.01	78.08 ± 1.91	76.0 ± 2.83	67.21 ± 1.98
			RWL		38.35 ± 2.55	59.51 ± 8.62	51.06 ± 10.14	61.91 ± 2.90
			TBWL		35.21 ± 4.37	57.91 ± 13.69	50.47 ± 14.21	54.79 ± 4.77
			DWsurv		38.92 ± 0.96	67.23 ± 2.72	59.07 ± 3.10	59.08 ± 1.24
			ASCL		38.94 ± 3.09	67.48 ± 11.51	52.29 ± 10.35	63.02 ± 4.42
			CA-TReL		53.28 ± 0.54	92.77 ± 1.89	92.45 ± 1.72	93.07 ± 1.09
		Uniform	g = 0	5	4.77 ± 0.00	80.14 ± 0.02	69.69 ± 0.06	57.19 ± 0.06
			g = 1		4.49 ± 0.00	19.86 ± 0.02	4.30 ± 0.04	24.21 ± 0.03
			Random		4.69 ± 0.00	50.16 ± 0.10	25.08 ± 0.17	37.83 ± 0.05
			OWL		4.85 ± 0.01	80.15 ± 0.02	78.83 ± 2.18	68.10 ± 1.67
			RWL		4.98 ± 0.00	90.03 ± 1.57	86.06 ± 1.00	85.79 ± 0.77
			TBWL		4.79 ± 0.03	68.47 ± 11.68	48.57 ± 13.19	50.46 ± 5.83
			DWsurv		4.84 ± 0.01	68.19 ± 0.79	61.18 ± 1.15	59.15 ± 0.44
			ASCL		4.98 ± 0.00	88.68 ± 1.60	83.55 ± 1.79	81.71 ± 1.75
			CA-TReL		4.99 ± 0.00	82.49 ± 1.34	86.72 ± 0.85	90.75 ± 0.62
3	True	Exponential	g = 0	23	16.98 ± 0.02	80.16 ± 0.07	62.03 ± 0.08	56.21 ± 0.09
			g = 1		7.96 ± 0.02	7.74 ± 0.04	0.00 ± 0.00	8.44 ± 0.01
			g = 2		6.15 ± 0.09	12.11 ± 0.04	4.78 ± 0.04	12.33 ± 0.05
			Random		10.56 ± 0.10	33.32 ± 0.06	11.36 ± 0.05	21.12 ± 0.05
			OWL		17.50 ± 0.46	38.98 ± 16.64	36.29 ± 15.89	50.04 ± 8.59
			RWL		17.30 ± 0.73	60.47 ± 6.54	44.58 ± 4.52	49.64 ± 2.92
			TBWL		14.40 ± 1.99	52.89 ± 16.65	44.67 ± 13.55	46.53 ± 9.83
			DWsurv		15.69 ± 0.32	66.71 ± 3.15	35.58 ± 1.90	43.29 ± 1.29
			ASCL		9.97 ± 1.47	31.02 ± 14.00	13.18 ± 4.17	23.75 ± 5.04
			CA-TReL		21.41 ± 0.25	92.30 ± 1.71	92.87 ± 1.01	92.13 ± 1.60
		Conditional	g = 0	29	22.61 ± 0.08	80.14 ± 0.03	63.45 ± 0.05	56.55 ± 0.08
			g = 1		7.44 ± 0.02	7.74 ± 0.02	0.00 ± 0.00	8.45 ± 0.01
			g = 2		5.60 ± 0.01	12.10 ± 0.09	4.27 ± 0.09	12.36 ± 0.06
			Random		11.58 ± 0.13	33.47 ± 0.06	11.18 ± 0.12	21.25 ± 0.07
			OWL		21.82 ± 2.74	53.73 ± 13.80	52.54 ± 13.42	56.82 ± 8.87
			RWL		19.31 ± 1.10	69.65 ± 3.35	41.38 ± 6.05	43.29 ± 4.47
			TBWL		18.54 ± 2.89	52.76 ± 14.82	40.55 ± 12.79	44.34 ± 8.77
			DWsurv		19.81 ± 0.03	68.34 ± 1.31	37.63 ± 0.58	43.97 ± 0.30
			ASCL		15.73 ± 0.18	60.17 ± 11.93	31.15 ± 8.26	33.92 ± 0.66
			CA-TReL		27.82 ± 0.09	80.84 ± 3.91	82.79 ± 3.03	80.72 ± 2.07
		Uniform	g = 0	5	4.61 ± 0.00	80.14 ± 0.03	62.40 ± 0.02	56.22 ± 0.08
			g = 1		4.15 ± 0.00	7.74 ± 0.04	0.00 ± 0.00	8.45 ± 0.01
			g = 2		2.69 ± 0.00	12.12 ± 0.02	4.8 ± 0.03	12.35 ± 0.02
			Random		3.89 ± 0.00	33.32 ± 0.08	11.11 ± 0.06	21.14 ± 0.07
			OWL		4.71 ± 0.03	66.37 ± 13.75	57.41 ± 11.77	62.61 ± 5.74
			RWL		4.93 ± 0.00	82.61 ± 0.64	79.94 ± 0.61	76.89 ± 0.54
			TBWL		4.33 ± 0.31	60.11 ± 12.25	42.09 ± 14.12	42.91 ± 11.30
			DWsurv		4.66 ± 0.00	67.35 ± 1.41	36.69 ± 0.58	43.45 ± 0.48
			ASCL		4.79 ± 0.01	82.73 ± 0.35	76.83 ± 1.12	70.83 ± 0.77
			CA-TReL		4.96 ± 0.02	86.16 ± 3.50	90.19 ± 2.36	88.98 ± 2.76
	False	Exponential	g = 0	23	17.00 ± 0.03	80.14 ± 0.06	62.01 ± 0.06	56.22 ± 0.07
			g = 1		7.98 ± 0.02	7.74 ± 0.03	0.00 ± 0.00	8.45 ± 0.01
			g = 2		6.12 ± 0.07	12.14 ± 0.03	4.76 ± 0.03	12.33 ± 0.06
			Random		10.58 ± 0.08	33.33 ± 0.05	11.33 ± 0.05	21.15 ± 0.06
			OWL		16.68 ± 1.18	53.09 ± 13.96	51.76 ± 12.63	54.88 ± 9.67
			RWL		16.42 ± 1.39	61.07 ± 6.77	42.09 ± 8.68	47.26 ± 5.66
			TBWL		15.04 ± 1.65	55.15 ± 14.97	42.27 ± 10.16	43.33 ± 7.85
			DWsurv		15.69 ± 0.39	66.61 ± 3.34	35.63 ± 2.05	43.41 ± 1.38
			ASCL		12.79 ± 2.15	50.31 ± 15.26	29.93 ± 12.13	36.68 ± 9.51
			CA-TReL		21.17 ± 0.37	88.39 ± 4.06	88.41 ± 4.44	88.44 ± 4.18
		Conditional	g = 0	29	22.61 ± 0.08	80.14 ± 0.03	63.45 ± 0.05	56.22 ± 0.08
			g = 1		7.44 ± 0.02	7.74 ± 0.02	0.00 ± 0.00	8.45 ± 0.01
			g = 2		5.46 ± 0.14	12.12 ± 0.03	4.28 ± 0.04	12.35 ± 0.04
			Random		11.58 ± 0.13	33.47 ± 0.06	11.18 ± 0.12	21.25 ± 0.07
			OWL		25.62 ± 0.23	78.66 ± 1.08	77.64 ± 0.33	70.84 ± 0.33
			RWL		18.76 ± 1.17	70.43 ± 3.24	36.75 ± 6.46	40.77 ± 4.83
			TBWL		17.06 ± 3.39	56.75 ± 14.33	45.6 ± 10.73	42.4 ± 9.18
			DWsurv		19.68 ± 0.19	67.95 ± 1.41	37.34 ± 0.74	43.79 ± 0.39
			ASCL		14.09 ± 1.16	56.17 ± 10.74	22.77 ± 3.79	28.5 ± 2.92
			CA-TReL		27.87 ± 0.10	80.84 ± 3.91	82.94 ± 2.99	80.88 ± 2.02
		Uniform	g = 0	5	4.61 ± 0.00	80.14 ± 0.03	62.40 ± 0.02	56.22 ± 0.08
			g = 1		4.15 ± 0.00	7.74 ± 0.04	0.00 ± 0.00	8.45 ± 0.01
			g = 2		2.69 ± 0.00	12.12 ± 0.02	4.8 ± 0.03	12.35 ± 0.02
			Random		3.89 ± 0.00	33.32 ± 0.08	11.11 ± 0.06	21.14 ± 0.07
			OWL		4.65 ± 0.13	29.55 ± 11.34	28.06 ± 10.67	38.84 ± 6.57
			RWL		4.91 ± 0.00	79.47 ± 0.69	75.71 ± 0.82	73.67 ± 0.67
			TBWL		4.37 ± 0.11	63.20 ± 11.72	35.63 ± 10.68	35.99 ± 8.07
			DWsurv		4.67 ± 0.00	66.61 ± 1.56	36.36 ± 0.76	43.24 ± 0.65
			ASCL		4.80 ± 0.01	82.99 ± 0.41	76.48 ± 1.22	70.38 ± 0.93
			CA-TReL		4.95 ± 0.02	84.05 ± 3.45	89.31 ± 2.75	87.33 ± 3.04

In addition to these baselines, we compared CA-TReL with five advanced DTR estimation methods. First, Outcome Weighted Learning (OWL)¹⁷ is a value-based approach that formulates treatment assignment as a weighted classification problem, directly optimizing expected outcomes through outcome-weighted support vector machines. For censored survival outcomes, we implemented OWL using pseudo-observations as outcome weights to account for censoring, following the pseudo-value-based strategy described in the ASCL framework²⁷. Second, Residual Weighted Learning (RWL)¹⁹ extends OWL by weighting classification errors with residuals from an outcome model fitted on clinical covariates excluding treatment assignment, rather than with the observed outcomes themselves. This modification reduces sensitivity to simple shifts in the outcome, improves finite-sample stability, and is implemented using a smoothed ramp loss within a penalized learning framework. Third, Tree-Based Weighted Learning (TBWL)¹⁶ extends outcome-weighted learning to right-censored survival outcomes by replacing censored failure times with non-parametrically estimated pseudo-rewards derived from recursively imputed survival trees. It uses the conditional expected survival time to impute only censored observations, and then learns the treatment rule using weighted support vector machines. Fourth, DWsurv^22,23 is a doubly robust method, meaning the estimated treatment effect is consistent if it correctly specifies either the treatment/censoring models (used in the weights) or the outcome (treatment-free) model. It estimates optimal DTRs by solving a series of weighted generalized estimating equations (GEEs). Fifth, Accountable Survival Contrast Learning (ASCL)²⁷ is a recent method that leverages a contrast-learning framework, which is based on the idea of the doubly-robust weighted classification approach¹⁷. By incorporating pseudo-value approaches to address censoring in survival data, ASCL used penalized support vector machines. All five methods have demonstrated strong performance for DTR estimation^{14,16,19,20,22,23,26–28}. We selected these approaches as benchmarks due to their methodological relevance and publicly available implementations, which ensure reproducibility and facilitate direct comparison with our proposed CA-TReL. We used publicly available implementations for reproducibility: OWL and ASCL were obtained from the SurvDTR GitHub repository (https://github.com/taehwa015/SurvDTR), and DWSurv from the authors’ GitHub repository (https://github.com/gabriellesimoneau/Rcode-JASA2019). Building on the OWL implementation, we implemented RWL and TBWL following the procedures described in their original papers, as both methods can be viewed as extensions of the OWL framework. In cases with multi-category treatment options, we adopted a one-vs-one classification strategy for all applicable methods, consistent with the implementation used by the ASCL authors in their simulation studies.

The evaluation metrics included Restricted Mean Survival Time (RMST), which captures the average survival time up to a specified truncation time , making it well-suited for censored data. Correct Decision Rate at the First Stage (CDR1) measured the accuracy of treatment assignments at the initial stage, while Average Correct Decision Rate (ACDR) provided an overall assessment of decision accuracy across all stages²⁷. Counterfactual outcomes were also analyzed to estimate the survival benefits that an individual would have received, if they had followed the treatments recommended by each method¹³.

Synthetic data studies

Given the absence of a true reward function or true policy in existing survival datasets, we generated synthetic datasets by integrating real-world data with simulated censoring mechanisms. Importantly, the survival models only utilize the observed censored times for the subjects affected by those simulated censoring mechanisms. The true underlying information, including the given reward function and true policy used for generating the synthetic data, was exclusively used during the evaluation phase to assess the learned policy’s effectiveness and was not incorporated during the training process. Full details of the synthetic data generation process are provided in the “Synthetic data generation” subsection of the Methods section.

Analysis results

We evaluated CA-TReL against baseline and advanced dynamic treatment regime (DTR) methods across binary and trinary treatment settings under different propensity models, and censoring types. We considered four key performance metrics: restricted mean survival time (RMST), correct decision rate at stage 1 (CDR1), average correct decision rate (ACDR), and expected survival outcome. The evaluation results – presented in Table 1; Fig. 1 – demonstrate the strong empirical performance of the proposed CA-TReL method across various evaluation metrics.

Fig. 1.

Performance comparison of treatment policies across three censoring mechanisms.

In the binary treatment scenarios (i.e., a sequence of binary decisions), CA-TReL consistently outperformed other methods. Under exponential censoring with a correctly specified propensity score model, Table 1 shows that CA-TReL’s results surpassed all other methods – OWL, RWL, TBWL, DWsurv, and ASCL -- in terms of RMST, CDR1, ACDR and expected survival. CA-TReL similarly demonstrated superior performance under conditional censoring with a correct propensity model, and the uniform censoring scenario under a correct propensity model – in all cases, doing better than the other models. We also performed statistical comparisons separately within each simulation setting, where a setting was defined by the treatment-option level, propensity-model specification, and censoring type. Within each setting, we used paired two-sided  t-tests on expected survival to compare CA-TReL with each competing method across the same simulation replications. To account for multiple pairwise comparisons within each setting, we adjusted the resulting p-values using the Holm–Bonferroni procedure. After multiplicity adjustment, CA-TReL remained significantly better than the included competing methods in the reported binary-treatment settings at the family-wise significance level of α = 0.05.

For the trinary treatment settings (i.e., a multistage setting with three treatment options per stage) under correct propensity specification, CA-TReL maintained its robust performance advantage across all three censoring types. It maintained the highest ranking in RMST, CDR1, ACDR, and expected survival, with especially notable advantages under exponential and conditional censoring. Although the performance gap between CA-TReL and other advanced methods narrowed under uniform censoring, the overall results highlight CA-TReL’s robustness, particularly in handling informative censoring conditions where the performance of other methods declined more noticeably.

Figure 1 succinctly illustrates these findings through radar plots, presenting normalized performance metrics averaged across both correct and incorrect propensity model specifications. CA-TReL consistently forms the largest areas within these radar plots, visually emphasizing its balanced and superior performance across diverse synthetic scenarios and censoring mechanisms.

Overall, these results underline the strength of CA-TReL in accurately estimating dynamic treatment regimes and maximizing survival outcomes, regardless of treatment complexity, censoring type, or propensity model specifications. The consistent outperformance across various metrics establishes CA-TReL as a reliable method for personalized treatment planning.

Radar plots show normalized performance across four metrics: RMST, CDR1, ACDR, and Expected Survival. Values are normalized across methods within each metric for interpretability. Results are averaged over correctly and incorrectly specified propensity score models. a Binary treatment setting under three censoring types: Exponential, Conditional, and Uniform. b Trinary treatment setting under the same three censoring types. To help interpret each subfigure, note that each “direction” of the radar plot represents one evaluation metric. Each method is visualized as a quadrilateral connecting its normalized performance on these four metrics. A method that performs well across all metrics will produce a larger and more outward-reaching polygon. For example, in panel (a), CA-TReL forms the most expansive and balanced polygon across all corners, indicating its strong and consistent performance under exponential censoring. In contrast, competing methods show smaller or more uneven polygons, reflecting weaker performance overall or strength concentrated in only a subset of metrics. This visualization provides an intuitive comparison of each method’s multi-objective effectiveness under a specific censoring scenario.

Observational studies

We also applied this proposed method to the SANAD dataset²⁹ to illustrate its utility in estimating DTRs in a real-world clinical setting. We chose the SANAD dataset because SANAD explicitly captures multi-stage clinical decision-making with clearly defined treatment stages and clinically meaningful endpoints, and its use in prior DTR literature allows robust benchmarking. The SANAD study, a randomized controlled trial, evaluated the effectiveness of two antiepileptic drugs — carbamazepine (CBZ) and lamotrigine (LTG) — for patients with partial epilepsy. Treatment failure, defined as the withdrawal of the randomized drug due to unacceptable adverse effects (UAE) or inadequate seizure control (ISC), served as the primary outcome. The dataset includes both longitudinal and survival data for 605 patients, randomly assigned to CBZ () or LTG ().

Our decision to use this dataset and the preprocessing steps were inspired by the Zhang et al. study²⁶, which employed this data to explore DTR estimation. While the original SANAD study²⁹ is a randomized trial, the rate of drug titration — encompassing initial doses and subsequent dose adjustments — was determined by clinicians based on their judgment, effectively making it observational in nature. Zhang et al²⁶. noted that this setup provides a realistic scenario for evaluating DTRs, as the dosing adjustments introduce variability reflective of clinical practice.

Following Zhang et al.‘s approach²⁶, we focused on the first two treatment-decision visits (i.e., = 1 and = 2), although patients could have additional follow-up visits in the raw data. At each subsequent visit ( k>1), the clinician could decrease (), maintain (), or increase () the dose based on the patient’s response to prior treatment. Patients were followed at intervals determined by clinical judgment, leading to variability in follow-up schedules. Due to the limited number of observed events after the third visit, we restricted our analysis to the first two visits to ensure a reliable estimation of optimal treatment regimes. We excluded 31 patients who did not receive treatment (zero initial and subsequent doses) from the analysis, leaving 574 patients, with an event rate of approximately 33%.

Baseline covariates included age, gender, and the presence of a learning disability (recorded as “Yes” or “No”). Longitudinal data encompassed the calibrated dose of medication at each visit, along with the start and end times of treatment intervals. For stage 1 treatment prediction, we used age, gender, the start time of the dose interval, and the presence of a learning disability as input features. For stage 2 treatment prediction, we incorporated all stage 1 features along with the end time of the stage 1 interval and the dose amount used in stage 1.

Analysis results

To assess the performance of our proposed method in an observational setting, we employed an off-policy evaluation (OPE) framework inspired by the approach of Komorowski et al.³⁰. We partitioned the dataset into 80% for training and 20% for testing to facilitate this evaluation. For estimating the expected return of the learned policy, we utilized the Weighted Importance Sampling (WIS) estimator, complemented by a high-confidence evaluation strategy^31,32. Specifically, we generated 1,000 bootstrap samples to construct a 95% confidence interval for the estimated policy value, following the methodology outlined by these works^31,32. This bootstrapping approach provides a robust estimate of the performance of the given policies, accounting for variability in the data.

In the synthetic experiments, we did not use off-policy evaluation as the primary performance assessment because the simulation design provides oracle information unavailable in real datasets, including the true treatment rule and the latent uncensored outcomes of alternative actions. This allows direct evaluation of learned policies through metrics such as RMST, CDR1, ACDR, and expected survival. By contrast, in the real-world SANAD analysis, such oracle quantities are unavailable; therefore, we used off-policy evaluation based on weighted importance sampling to estimate policy value on held-out data.

Accurate estimation of the behavior policy is critical in OPE^33,34 because the importance weights depend directly on the estimated treatment-assignment probabilities. In particular, if these probabilities are poorly calibrated, the resulting weights can become unstable, especially when estimated action probabilities are too close to 0. Recognizing the significance of calibration in this context, we evaluated three modeling techniques — Logistic Regression, Random Forest, and Light Gradient Boosting Machine (LightGBM) — to estimate the behavior policy. Calibration quality was assessed using the Expected Calibration Error (ECE) metric^35,36, as emphasized in recent studies on policy evaluation. We then select the model with the lowest ECE, anticipating it will have well-calibrated probability estimates, thereby enhancing the trustworthiness of the OPE results. This careful approach to behavior policy modeling and confidence-aware evaluation underpins the robustness of our policy value estimates.

The observational off-policy evaluation was then formulated using the estimated behavior policy and a softened evaluation policy. Let denote the behavior policy at stage  k, defined as the conditional probability that the observed clinical process assigns treatment  a given the available treatment history : In practice, this behavior policy was estimated from the training data as using the calibrated behavior-policy model selected by the lowest ECE. In this experiment, for stage 1, the behavior-policy model used age, gender, learning-disability status, and the start time of the first treatment interval. For stage 2, the behavior-policy model was fitted among individuals who entered the second treatment stage and used the observed stage-1 follow-up time, first-visit dose, and stage-1 treatment.

For a learned dynamic treatment regime , we followed Komorowski et al.³⁰ and evaluated a softened version of the learned policy, denoted by . Specifically, at stage  k, the treatment recommended by the learned regime was assigned probability  1-p, while the remaining probability  p was distributed equally among the non-recommended treatment options:

where denotes the set of available treatment options at stage  k, and  p=0.01 was used to avoid zero-probability actions.

The subject-specific importance weight was then computed as.

where for individuals who did not enter the second treatment stage and for individuals who entered the second treatment stage. To improve numerical stability, the final importance weight was clipped to the interval : The self-normalized weighted importance sampling estimate of the mean expected survival under the learned policy was computed on the held-out test set as.

where denotes the observed survival or follow-up outcome for subject  i. We used 1,000 bootstrap resamples of the held-out test set to estimate the distribution of and to report the mean expected survival with corresponding uncertainty estimates.

Using the calibrated behavior policy model and the high-confidence WIS estimator, we evaluated the learned policies from each method on the held-out test set. The WIS estimates, reported as the mean expected survival (in days) from 95% confidence intervals generated using 1,000 bootstrap samples, were 724.25 for OWL, 1076.77 for RWL, 917.28 for TBWL, 1231.40 for ASCL, 1259.54 for DWsurv, and 1285.44 for CA-TReL. In addition, the mean and standard deviation of the 1,000 bootstrap WIS estimates were 591.71 ± 80.93 for OWL,778.29 ± 178.72 for RWL, 737.21 ± 102.94, 717.88 ± 314.61 for ASCL, 767.01 ± 352.47 for DWsurv, and 970.41 ± 270.29 for CA-TReL, respectively. We conducted paired two-sided  t-tests comparing the bootstrap WIS distributions of CA-TReL with each competing method. To account for multiple pairwise comparisons across methods, the resulting  p-values were adjusted using the Holm–Bonferroni procedure, with significance assessed at . After multiplicity adjustment, CA-TReL showed significantly higher estimated WIS policy value than OWL, RWL, TBWL, ASCL, and DWsurv. Since higher WIS values reflect greater expected survival under the learned policy, these results suggest favorable empirical performance for CA-TReL in this observational setting, with CA-TReL yielding the highest estimated survival benefit among the evaluated methods. However, it is important to interpret these findings with caution. Off-policy evaluation is inherently limited by its reliance on accurate behavior policy estimation and sufficient overlap between observed and recommended treatments. Despite minimizing calibration error and using bootstrapping to provide high-confidence estimates, there are still many potential sources of bias that may still affect the validity of the results, including unmeasured confounding and model misspecification.

At the initial visit, CA-TReL’s estimated DTRs suggest prescribing LTG to a greater number of individuals compared to the current treatment strategy. This observation aligns with prior findings^26,37, which highlight LTG’s advantage over CBZ in reducing treatment failures, despite CBZ being the long-standing standard treatment.

Sensitivity analysis under nuisance-model misspecification

We further evaluated the robustness of CA-TReL under four nuisance-model specification settings for the binary treatment setting. Here, the nuisance-models refer to the two auxiliary models used in the CA-TReL estimator: (i) the outcome regression model, used to estimate the conditional outcome mean, and (ii) the censoring model, used to estimate the conditional probability of remaining uncensored. In the simulation setting, we consider these four cases: both the censoring and outcome models are correct, only the outcome model is misspecified, only the censoring model is misspecified, and both models are misspecified. This experiment was conducted under the conditional censoring scenario, which provides a more stringent setting for assessing censoring-model misspecification because it allows censoring to depend on the covariates. In this sensitivity analysis, “correctly specified” means that the working nuisance model used the same covariates and functional structure as the corresponding data-generating mechanism, while “misspecified” means that key covariate-dependent or nonlinear components were intentionally omitted.

For the outcome-regression model, the correctly specified working model used covariates aligned with the simulated outcome-generating mechanisms. At stage 1, the outcome-regression model included baseline creatinine, potassium, and hemoglobin. At stage 2, the outcome-regression model included the observed stage-1 outcome, glucose, as well as the earlier stage-1 features: baseline creatinine, baseline hemoglobin, and baseline potassium. In contrast, the misspecified outcome-regression model intentionally omitted these key outcome-generating covariates and instead used an alternative covariate set consisting of sodium, platelet count, and hematocrit. Thus, the misspecified outcome model did not fully capture the covariate-dependent outcome structure used to generate the simulated survival outcomes.

For the censoring model, the correctly specified working model matched the covariate structure of the conditional censoring mechanism used in data generation. Specifically, the censoring model included baseline creatinine and the absolute deviation of baseline potassium from 4.0. In contrast, the misspecified censoring model intentionally omitted the key censoring-generating covariates, namely baseline creatinine and the potassium-deviation term, and instead used an alternative covariate set consisting of hemoglobin, glucose, sodium, stage-1 treatment, and observed stage-1 follow-up time. This allowed us to evaluate how CA-TReL behaved when the censoring model failed to capture the covariate-dependent censoring structure used in the simulation.

As shown in Table 2, CA-TReL performed best when both models were correctly specified – with respect to all 4 evaluation measures – and showed only modest degradation under single-model misspecification. Performance declined further when both models were misspecified, as expected, but CA-TReL still substantially outperformed the random, fixed-treatment rules and other methods. Overall, these results suggest that CA-TReL is reasonably robust to moderate misspecification of either nuisance component in the binary treatment case. These results are not directly comparable to the CA-TReL results reported in Table 1, because in this sensitivity analysis, we explicitly supplied the nuisance models as correct or misspecified, whereas in Table 1 the nuisance models were estimated from the data rather than being constrained to pre-specified correctly specified or intentionally misspecified working-model structures.

Table 2.

Sensitivity analysis under censoring- and outcome-model misspecification for the binary treatment setting. Performance of CA-TReL under four nuisance-model specification settings. Values are mean ± standard deviation, and the best model for each column is highlighted in bold.

CA-TReL under model specification setting		-RMST	CDR1	ACDR	Expected Survival outcome
Censoring?	Outcome?
Correct	Correct	53.47 ± 0.64	94.74 ± 1.91	93.93 ± 1.83	93.98 ± 1.16
Correct	Wrong	52.96 ± 0.77	91.71 ± 2.92	93.16 ± 1.18	90.03 ± 3.52
Wrong	Correct	53.38 ± 0.64	92.77 ± 1.89	92.45 ± 1.71	93.07 ± 1.09
Wrong	Wrong	52.89 ± 0.75	91.35 ± 2.90	92.95 ± 1.18	89.64 ± 3.49

Discussion and future scope

This work presents CA-TReL, a censoring-adjusted extension of tree-based reinforcement learning (T-RL), that addresses the critical need for robust estimation of dynamic treatment regimes (DTRs) in the presence of censored survival outcomes. By integrating censoring adjustments and leveraging augmented inverse probability weighting (AIPW), CA-TReL is able to account for incomplete outcome observation due to censoring. We validated the model’s performance through rigorous simulations and a real-world clinical dataset (SANAD), which features multi-stage treatment decision-making in epilepsy care. In these evaluations, CA-TReL achieved strong empirical performance relative to established methods — including Outcome Weighted Learning (OWL), Residual Weighted Learning (RWL), Tree-Based Weighted Learning (TBWL), Dynamic Weighted Survival Modeling (DWsurv), and Accountable Survival Contrast Learning (ASCL) — in terms of expected survival outcomes and treatment decision accuracy. While further validation is needed before clinical deployment, these results suggest that CA-TReL may be useful for supporting interpretable, personalized treatment-rule estimation in settings with censored survival outcomes.

Despite its strengths, CA-TReL has certain limitations that warrant further research. First, the framework currently supports discrete treatment options, limiting its applicability to settings where treatments can be categorized into distinct groups. This means it is not designed to cope with real-world scenarios that involve continuous treatment variables, such as drug dosage optimization. Second, treatment selection bias remains a challenge, particularly in observational datasets where treatment assignments may reflect underlying clinician preferences. While CA-TReL incorporates propensity score adjustments to mitigate confounding, residual bias could influence the estimated treatment regimes. The validity of the estimated treatment regimes still depends on standard causal assumptions, including consistency, treatment positivity, conditional exchangeability, and appropriate handling of censoring^4,5,7,21,26. In observational settings, unmeasured confounding and treatment-selection bias may remain, particularly when treatment assignments reflect clinician preferences, disease severity, or other factors not fully captured in the observed history. Incorporating more sophisticated causal inference techniques could further strengthen the model’s robustness and reliability. Third, CA-TReL relies on nuisance models for outcome regression and censoring adjustment, and misspecification of these components may affect performance^18,20,22,26. Although our sensitivity analyses suggest reasonable robustness to moderate nuisance-model misspecification, performance declined when both nuisance components were misspecified. Fourth, it employs a greedy approach when selecting splits, focusing solely on the purity improvement at the current node, which simplifies the tree construction process. While this might produce sub-optimal policies, this limitation is likely modest in the present two-stage setting (=2). An alternative approach to enhance the performance is to incorporate a “lookahead” mechanism³⁸– by evaluating potential future splits and their impact on downstream nodes, this method can identify splits that optimize the tree structure more holistically, potentially leading to better overall treatment decisions. Finally, while the method demonstrated promising results on the SANAD dataset, additional validation across diverse clinical conditions is necessary to generalize its applicability. Future studies could explore CA-TReL’s performance in domains such as cardiovascular diseases, and diabetes management, refining the methodology for specific clinical settings^3,23,61,62.

Conclusions

CA-TReL offers a promising, interpretable framework for estimating dynamic treatment regimes in settings with censored survival outcomes. By integrating censoring-aware adjustments and leveraging the strengths of tree-based methodologies, it provides an effective framework for estimating personalized treatment rules. The model’s statistically supported empirical performance in both simulated and real-world datasets highlights its potential value for treatment-rule estimation under censoring. However, addressing limitations related to continuous treatments, treatment selection bias, nuisance-model misspecification, greedy split selection, and validation across broader clinical settings will be critical for broader adoption. While further methodological development and clinical validation remain necessary before broader adoption, by pursuing these advancements and validating the approach across diverse medical domains, CA-TReL may contribute to more interpretable and data-driven precision-medicine strategies.

Methods

Notation and assumptions

We use the context of coronary artery disease (CAD) as a motivating example to illustrate the dynamic treatment regimes (DTR) framework. CAD is a leading cause of hospitalization and mortality worldwide, presenting a complex clinical challenge that necessitates personalized and adaptive treatment strategies to optimize survival outcomes. We focus on CAD in this study due to its clinical significance, high prevalence, and the availability of detailed patient data in the MIMIC-IV dataset³⁹. In our study, we assume treatment decisions are made over two stages to maximize the long-term outcome (the patient’s overall survival time): Stage 1, focusing on renal and anemia management, followed by stage 2, focusing on glycemic management. These stages reflect two critical aspects of CAD management that affect long-term outcomes. The event of interest is defined as a critical clinical endpoint (e.g., mortality) that could terminate follow-up. It is important to note that this example is designed solely for experimental purposes, and actual treatment stages and options may differ in real-life clinical scenarios.

We assume that data from  N individuals are available. The data from the individuals are independent, identically distributed (iid) copies of covariates (), treatments (), and outcomes () for all individuals, at each stage  k. In the CAD context, X includes patient clinical measurements such as age, creatinine, hemoglobin, potassium, sodium, glucose, platelet count, hematocrit, and white blood cell count. The treatment indicator can represent a multi-categorical or ordinal treatment, with observed values  where | | = denotes the number of available treatment options at the stage  k. For instance, in the CAD context, stage 1 focuses on renal and anemia management, with treatment options including conservative monitoring (observation without intervention), moderate support (medications or dietary adjustments for mild dysfunction), and aggressive intervention (treatment for severe anemia or renal dysfunction). If this stage finishes, the patient moves to Stage 2, which addresses glycemic management, with several treatment options: conservative glycemic management (dietary and lifestyle interventions), pharmacological intervention (glucose-lowering medications), versus aggressive glycemic control (intensive management using insulin or combination therapies). In this framework,  Y represents the overall observed survival time, defined as , where denotes the observed survival time at the stage  k until the next stage or the event of interest. This serves as a quantitative measure for evaluating the effectiveness of treatment policies across the two stages.

Individual data are represented as sequences of covariates, treatments, and outcomes [[, organized by stages, with each stage indicated by the subscript  k. Here, stages refer to distinct points in the treatment process where decisions are made about interventions. Let individuals be identified with a superscript , although this superscript is often omitted for simplicity. For example, instead of explicitly writing [ for the covariates, treatments, and outcomes of the individual  i at the stage  k, we often simplify notation by using , implicitly assuming that the data are indexed by  i when referring to a single individual.

The policy can use the values of the covariates as they are known at the start of stage  k, to decide which treatment to administer. The treatment options for the second stage can be influenced by the treatment given in the first stage, as those second-stage covariates are determined by the treatment , covariates , and outcome observed during the first stage .

Let represent the histories, which include (a function of) covariates, previous treatments, and previous outcomes available at the start of the stage  k, used to guide the treatment decision for that stage. For instance, could be , where is a function summarizing past information at stage  k. For example, perhaps for the first stage and for the second stage, etc. In this paper, we specifically assume to be the identity function, such that for all  k. This means that the history directly incorporates all covariates, treatments, and outcomes from prior stages up to stage  k, without any transformation or aggregation.

In the CAD context, at stage 1, the patient’s history includes baseline covariates such as age, and laboratory values like creatinine and hemoglobin levels. Based on this history, the policy makes a treatment decision , such as prescribing a specific renal management strategy. The underlying stage 1 survival duration is denoted by representing the time until the next decision point, which could be either another treatment stage or terminating the process (perhaps due to an event like kidney failure, or death). If there is no terminal event during stage 1, the policy then always advances to stage 2. Here, the details of that next decision point are influenced by the patient’s clinical progression and the effectiveness of the current treatment. For example, this waiting period for CAD patients could be clinically determined based on the patient’s response to the initial treatment, the stability of their renal function, and other indicators such as hemoglobin levels. At this point, the policy uses the updated history , which incorporates the treatment from stage 1 , the elapsed time during stage 1, and additional covariates such as glucose levels at the end of stage 1 (before the start of stage 2). It can then use this to decide on the next treatment action – e.g., to help guide decisions for glycemic management. Finally, the underlying stage-2 survival duration is denoted by , representing the time until terminating the process (perhaps due to an event like kidney failure). The underlying overall survival time for an individual who reaches stage 2 is .

However, not all patients experience the event of interest within the observation period , resulting in censored outcomes. Censoring can occur, for example, when the study concludes before a patient experiences the event (e.g., when the underlying overall survival time ) or when a patient is lost to follow-up, etc. When , no information on survival is available beyond the study length in the observed data, meaning the outcome of interest is truncated by . The overall censoring time  C marks the point when the observation of an individual’s survival time stops without the event of interest occurring. If the censoring happens during an earlier stage, the patient does not progress to the next stage. Therefore, for some CAD patients, the underlying survival time  T is not fully observed, and the observed outcome is the censored follow-up time  Y.

Formally, let be a random variable that equals 1 if an individual is known to enter stage k and 0 otherwise, with the condition that for all individuals (as we assume that no individual can be censored before entering stage 1). Additionally, if , then for any future stage  j>k, . The duration of survival during each stage  k is represented by , defined as the time interval from the time the individual entered stage k until they either progress to stage  k+1 or experience a terminating event such as failure or death, whichever occurs first. For an individual, this process generates a sequence of covariates, treatments, and outcomes [[, where the observed survival outcome is , with  denoting the underlying survival time, which is not always directly observed, and represents the remaining censoring time at the stage  k, calculated based on the individual’s overall censoring time  C and their progression through the stages. Specifically, is defined as , where represents the cumulative survival time up to stage  k-1. The censoring indicator  specifies whether the event occurred without censoring before the end of the stage  k: means that the full stage- interval is observed before censoring, while means censoring occurs during stage. If censoring occurs at a given stage, no further outcomes, treatments, or covariates are observed for subsequent stages – i.e., this process just terminates. Under this setup, we assume stage-specific conditional independent censoring. Specifically, for each stage, conditional on the observed history and the treatment assigned  at that stage, the remaining censoring time is assumed to be independent of the relevant stage- outcome. The observed outcome is the overall censored follow-up time, which is the sum of the survival times across all stages, .

A Dynamic Treatment Regime (DTR) consists of a series of decision rules , where  G represents the set of all possible treatment regimes. Each policy  g comprises stage-specific decision rule functions . At each stage  k, the decision rule is a function that maps the history  to a treatment, where represents the set of all possible decision rules at the stage  k. This means that provides a treatment recommendation based on the history up to stage  k.

To make the counterfactual notation precise, we distinguish between the full counterfactual survival outcome and the stage-specific continuation outcome used in backward induction. Here, a stage-specific continuation outcome refers to the underlying outcome accumulated from a given stage through the final stage, under a specified treatment at that stage and specified future treatment rules thereafter. Thus, it represents the remaining underlying survival outcome from entry into that stage through the final stage, rather than only the immediate outcome observed during that stage. For a complete dynamic treatment regime , let denote the underlying overall counterfactual survival time that would be observed if an individual followed regime  g throughout the treatment course. This is the counterfactual analogue of the underlying survival time  T, and it is distinct from the observed censored follow-up outcome  Y.

For a fixed treatment sequence , let denote the underlying overall counterfactual survival time that would be observed if the individual received the fixed treatment sequence . For a complete regime  g, this can be written as the counterfactual outcome corresponding to the treatment sequence selected by the regime:

1

This definition clarifies that refers to the overall underlying counterfactual survival outcome under a complete regime, not to the observed censored outcome.

For recursive estimation, we define a stage-specific counterfactual continuation outcome. Let denote the observed treatment history before stage  k, and let denote the optimal future treatment rules after stage  k. We define as the underlying continuation outcome from stage  k through stage  K, given the observed prior treatment history , assigning treatment at stage  k, and following the optimal future rules thereafter. Thus, is the stage-specific counterfactual continuation outcome targeted by the stage- k backward-induction step. In particular, when  k=1, the stage-specific continuation outcome corresponds to the full counterfactual survival outcome under the complete regime.

The goal of an optimal DTR is to maximize the expected patient outcomes over the entire treatment course. The value function for a complete regime  g is The optimal regime, denoted by , is the regime that maximizes this value function: For example, in the case of the aforementioned CAD example, represents the expected underlying survival time if all patients followed regime  g, and is the regime that maximizes expected survival across the population.

Since only one of the counterfactual outcomes is observed for each individual, estimating the full counterfactual distribution requires specific assumptions. To address this challenge, we rely on the following key assumptions^13,26:

• Positivity: This condition, ), ensures that all individuals have a chance of receiving each treatment. In practice, violations commonly arise in the empirical sample when certain treatments are never (or almost never) prescribed to individuals with specific characteristics (e.g., sicker patients rarely receive a more intensive therapy). This can lead to regions of the covariate space where one treatment is effectively unavailable, making estimation of counterfactuals infeasible.

• Consistency: The observed outcome for an individual under a given treatment matches the counterfactual outcome predicted under that treatment.

• No unmeasured confounding: It ensures that, given the observed covariates, the treatment assignment is independent of the potential outcomes, or stated another way, the outcome of a treatment can depend on the observed covariates, but NOT on the decision of which treatment to apply. For instance, for stage 1, ; then for stage 2, ; and so on, where and represent the potential outcomes corresponding to the respective treatment actions. This assumption enables valid causal inferences and is sometimes called ‘strongly ignorable treatment assignment’ or ‘exchangeability’. For the longitudinal setting, this assumption requires that the treatment assignment at a given stage cannot depend on future covariates nor on treatment history. It is also called ‘sequential randomization’. This assumption is frequently the biggest challenge in observational data, where there may be variables related to both treatment decisions and outcomes that are not measured (unobserved confounders). Violations can lead to biased estimates of causal effects.

• No interference assumption: The outcome for any individual should not be influenced by the treatment assignment of another individual. It is also called ‘Stable unit treatment value’. It can be difficult to maintain in real-world clinical settings where doctors’ decisions for one patient often influence subsequent treatments and outcomes.

By adhering to these assumptions, the observed data can be linked to the counterfactual framework, allowing for valid causal inference and identification of optimal DTRs. In addition, we assume a data coverage condition analogous to the concentrability assumption in offline reinforcement learning^40,41— that is, the offline dataset must sufficiently cover the state-action pairs visited by the optimal policy, ensuring the learned policy can be evaluated and optimized reliably from historical data.

Under these causal assumptions, the optimal decision rule at any stage  k is derived recursively, starting from the final stage and moving backward. This is defined as:

Final Stage  K: At the final stage, the previous treatment history has already been observed. Let denote the underlying final-stage continuation outcome that would be observed if treatment were assigned at stage  K. The optimal rule at the last stage  K is

2

Here, denotes expectation with respect to the marginal distribution of patient histories at stage  K, and denotes the number of available treatment options at stage  K. The inner conditional expectation represents the mean underlying final-stage continuation outcome under treatment , given the final-stage history .

Intermediate Stages : For earlier stages, the optimal rule is defined recursively by incorporating the expected continuation outcomes from future stages. We consider , which represents the underlying counterfactual continuation outcome from stage  k through stage  K, given the observed prior treatment history , applying the decision rule at stage  k, and following the optimal future rules thereafter.

At any intermediate stage  k, the optimal decision rule is given by:

3

Here, is the censoring-adjusted pseudo-outcome at stage  k, representing an intermediate estimate of the underlying counterfactual continuation outcome from stage  k onward while accounting for censoring and the optimal decisions at future stages.

Censoring-adjusted tree-based reinforcement learning (CA-TReL)

Censoring-adjusted tree-based reinforcement learning builds on the foundational principles of decision trees and is inspired by the work of Tao et al.¹³, who extended traditional tree-based methods to optimize counterfactual outcomes in dynamic treatment regimes (DTRs). Tao et al.’s framework focuses on complete datasets without censoring, where all outcomes are fully observed. In contrast, CA-TReL adapts this approach to address the complexities of right-censored survival data, making it applicable in settings where incomplete observations are prevalent.

Tree-based learning framework

The construction of decision trees in CA-TReL follows the idea of recursively partitioning the covariate space to create regions that maximize homogeneity based on a “purity” criterion. Unlike standard decision tree methods, such as Classification and Regression Trees (CART)⁴², which rely on directly observed outcomes, CA-TReL focuses on counterfactual mean outcomes. These represent the expected results under hypothetical treatment assignments and serve as the foundation for assessing purity.

At each stage of the tree-growing process, CA-TReL evaluates possible splits by calculating the counterfactual mean outcomes for the resulting child nodes. The purity improvement from a split is determined by comparing the counterfactual outcomes of the child nodes to the parent node. The split that maximizes this improvement is selected, ensuring that the tree partitions the covariate space in a way that optimizes treatment assignment.

Recursive partitioning for counterfactual optimization

To guide the tree-growing process, CA-TReL employs a recursive partitioning framework. Let denote a parent node and  s a potential split at that node. The counterfactual purity measure evaluates the counterfactual mean outcome for the split and without splitting. The optimal split is selected to maximize the purity improvement. This recursive process continues until predefined stopping criteria are met, resulting in a tree structure that optimally partitions the covariate space to guide treatment decisions. The detailed formulation of the purity measure and split-selection objective is provided in the following subsection, ‘Purity measure and split-selection objective’.

Counterfactual framework and recursive optimization

The recursive optimization framework of CA-TReL ensures that decisions at each stage account for their impact on future outcomes. At stage  k, the relevant counterfactual quantity is the stage-specific continuation outcome,

4

where denotes the observed prior treatment history and denotes the optimal future treatment rules. This quantity represents the underlying counterfactual continuation outcome from stage  k through stage  K, if the individual follows the decision rule at stage  k and then follows the optimal rules at all subsequent stages. This recursive structure aligns with Bellman’s principle of optimality and provides a framework for dynamically refining treatment rules based on evolving patient histories.

CA-TReL incorporates pseudo-outcomes at each stage to estimate counterfactual outcomes, particularly in the presence of censoring. These pseudo-outcomes are adjusted to account for incomplete data and serve as substitutes for the unobserved counterfactual outcomes. By integrating these censoring-aware adjustments, as described in the later section on CA-TReL estimation for stage  k, CA-TReL ensures accurate and robust optimization of treatment rules.

Stopping rules to prevent overfitting

CA-TReL employs stopping rules to prevent overfitting, following the framework of Tao et al.¹³, including:

1. Minimum Node Size Constraint: A node is split only if it contains at least  observations, and the resulting child nodes each contain at least observations, where denotes the minimum node size.

2. Purity Improvement Threshold: Splits resulting in a minimal improvement in purity (below a threshold ) are disallowed.

3. Maximum Tree Depth: The tree-growing process is capped at a user-defined maximum depth.

If none of these stopping rules are triggered, the node is split, and the tree continues to grow until the criteria are satisfied.

CA-TReL estimation for the final stage  K

To address the causal objective of optimal DTR estimation at the final stage  K, we propose a censoring-adjusted modification of the augmented inverse probability weighting estimator, which we term CAIPW. This approach builds upon the work of Tao et al¹³. and extends it to handle incomplete observations, ensuring robustness and accuracy in estimating counterfactual outcomes under censoring.

The final-stage CAIPW estimator depends on three nuisance functions: the treatment propensity score, the conditional outcome mean, and the conditional censoring survival function. For each treatment option , let denote the propensity score for receiving treatment at the final stage, conditional on the final-stage history . Let denote the conditional mean of the underlying final-stage continuation outcome if treatment is assigned at the final stage, while the previous treatment history is kept as observed. We also define the conditional censoring survival function as where is the remaining censoring time at the final stage. Therefore, denotes the conditional probability that an individual remains uncensored through the observed final-stage follow-up time , given the observed final-stage history and treatment . The estimated versions of these functions are denoted by , , and , respectively.

The target conditional estimand of the final-stage CAIPW estimator is , the conditional mean of the underlying final-stage continuation outcome under treatment , given the final-stage history . The corresponding population-level target is the marginal mean of this conditional counterfactual continuation outcome. Thus, the estimator targets the underlying counterfactual continuation outcome under treatment , rather than the observed censored follow-up outcome.

The estimator for the counterfactual mean outcome at the final stage is given by

5

where:

• is the censoring indicator.

• is the estimated propensity score for treatment .

• is the censoring survival probability conditional on the history , and .

• represents the conditional mean outcome.

Equation (5) defines the final-stage CAIPW estimating function. Its doubly robust property is stated formally in Supplementary Proposition 2. In particular, under consistency/SUTVA, sequential exchangeability at stage  K, treatment positivity, conditional independent censoring, censoring positivity, and consistent estimation of the censoring survival function, the CAIPW estimating function has conditional expectation, given , equal to the true conditional mean final-stage continuation outcome under treatment , if either the propensity-score model is consistently estimated or the outcome-regression model is consistently estimated. Consequently, the marginal expectation of the CAIPW estimating function equals the marginal mean underlying counterfactual continuation outcome under treatment . Equivalently, the sample average of the final-stage CAIPW estimating function converges to the marginal mean underlying counterfactual continuation outcome under treatment , provided that the censoring survival model is consistently estimated and either the propensity-score model or the outcome-regression model is consistently estimated. Here, “consistently estimated” means that the corresponding nuisance estimator converges to its true population target; it does not require exact finite-sample equality between the fitted model and the true nuisance function.

Therefore, the final-stage CAIPW estimator is doubly robust with respect to the propensity-score model and the outcome-regression model , conditional on consistent estimation of the censoring survival model . The first term corrects for treatment assignment and censoring through inverse weighting, while the augmentation term stabilizes estimation through outcome regression. Their combined behavior under nuisance-model misspecification is formalized in Supplementary Proposition 2.

CA-TReL estimation for stages

For intermediate stages , the estimation of dynamic treatment regimes involves a recursive process that propagates information from future stages back to earlier stages. This process relies on pseudo-outcomes that capture the cumulative effects of treatment decisions.

We employ an adjusted version of pseudo-outcomes to mitigate the cumulative bias arising from the conditional mean models, as described by Huang et al⁴³. For complete data, rather than relying solely on model-derived values under optimal future treatments , we incorporate observed outcomes along with the anticipated loss resulting from suboptimal treatments. In the case of censored data, the estimation is based exclusively on .

The pseudo-outcome at stage  k is defined as:

6

where:

• : Censoring indicator for stage  k+1, where =1 if uncensored, 0 otherwise,

• : Pseudo-outcome at the stage  k+1,

• : Expected outcome under the optimal treatment rule at the stage  k+1,

• : Expected outcome under the observed treatment at the stage  k+1.

The stage- ( k+1) switching term in Eq. (6) has a Bellman-valid population counterpart. Specifically, under the conditional independent assumption stated in Supplementary assumption S1, the conditional expectation of the population-level switching term given equals the continuation value under the optimal stage- ( k+1) rule. This justifies the use of the switching term in the backward-induction recursion. A formal statement and proof are provided in Supplementary Proposition S1.

The estimator for the counterfactual mean outcome at the stage is given by

7

Purity measure and split-selection objective

We now define the purity criterion used to grow the CA-TReL tree explicitly. Let  denote a node at stage  k, and let denote a candidate binary split of , where is the set of allowable splits (typically feature-threshold rules based on the covariates). The split  s partitions into two disjoint child regions, denoted by and , such that

For example, if the splitting variable is Age with a threshold of 60, then one candidate split is

Define the stage-specific counterfactual-value estimator

that is, is the CAIPW estimator of the stage- k counterfactual mean outcome under treatment .

For a given split  s, define the stage-k treatment rule.

which assigns treatment to the left child and to the right child.

The split purity at node under the split  s is

where is the empirical expectation operator,  n is the number of stage- k training observations, indexes subjects, and  k is the stage index.

The corresponding no-split purity is

The gain from split  s is

The optimal split is then chosen as

8

where is the number of observations in node  B.

A node is split only if , the maximum depth has not been reached, and

Otherwise, becomes a terminal node and is assigned the treatment

9

A worked numerical example illustrating the computation of the node purity, split gain, and optimal split selection is provided in Supplementary Section “Example of the node-purity and split-selection criteria”.

Algorithm 1.

CA-TReL training by backward induction.

Implementation of CA-TReL

The implementation of CA-TReL begins with the stage  K, where the observed outcome is directly used as the pseudo-outcome (), and proceeds via backward induction through earlier stages. The propensity scores () are estimated using multinomial logistic regression, while the censoring survival probability and the conditional mean outcome () are estimated using random survival forests. The pseudo-outcomes () are updated recursively, incorporating adjustments for censoring and the expected loss from suboptimal treatments, ensuring the robustness of the recursive estimation process and enabling effective handling of incomplete data.

To optimize the performance of the proposed CA-TReL method, hyperparameter tuning was performed by “internal cross-validation”, which was within each external (outer) cross-validation training fold. (That outer cross-validation was exclusively for final evaluation, but NOT for estimating hyperparameters). The censoring and outcome-regression random survival forest components were tuned by inner 5-fold CV over.

• number of trees: ntree ∈ {1, 2, 5, 10, 50, 100, 500, 1000},

• the minimum terminal node size: nodesize ∈ {1, 20, 50, 100, 500, 1000},

• number of candidate predictors considered at each split mtry ∈ – truncated to the interval [1,], and.

• number of candidate predictors considered at each split: nsplit ∈ {0, 1, 10, 50, 100, 500},

where  p denotes the number of predictor variables in the corresponding stage-specific input matrix. For the censoring model, the best setting was selected by minimizing the mean out-of-bag (OOB) error from random survival forests fitted on the inner training splits. For the outcome-regression model, the best setting was selected by minimizing the mean inverse-probability-of-censoring-weighted (IPCW) truncated-RMST prediction error on the inner validation folds. Each CA-TReL policy-tree parameters were selected from.

• the maximum depth of the decision tree,: treedepth ∈ {2, 3, 4, 5, 10},

• the strength of regularization: ∈ {0.01, 0.02, 0.03}, and.

• the minimum number of observations required for a node to be split: minsplit ∈ { max(0.01*N, 20), max(0.05*N, 20), max(0.10*N, 20), max(0.15*N, 20) }

using an internal validation split within the outer training fold. The final CA-TReL policy-tree settings were chosen based on the highest validation restricted mean survival time (RMST). This approach ensured robust and efficient tree-based estimation tailored to the dataset characteristics.

Synthetic data generation

Given the challenges of directly validating learned policies due to the absence of true reward functions and corresponding optimal policies in real-world survival datasets, and the need to reason about counterfactual outcomes, we created a synthetic dataset. This dataset preserved the structure and characteristics of real-world data, including realistic covariate distributions and censoring percentages, while incorporating simulated outcomes and treatment assignments to ensure a controlled yet practical evaluation environment.

For this study, we utilized the MIMIC-IV v2.0 dataset as the primary source³⁹, and selected Coronary artery disease (CAD) as the focus due to its prevalence and significance as a leading cause of hospitalization and mortality worldwide. The dataset provided comprehensive clinical information, including demographics, laboratory measurements, and outcomes from hospital admissions. To refine the analysis, we focused on non-ICU admissions, treating each hospital admission as an independent patient record to capture diverse clinical trajectories.

We consider a 48-hour observation window, divided into two 24-hour intervals. For the experimental purpose, we considered the first interval as the baseline covariates for Stage 1 (renal and anemia management), while the second interval as time-dependent changes for Stage 2 (glycemic management). The features included in the analysis were Age, Creatinine, Hemoglobin, Potassium, Sodium, Glucose, Platelet Count, Hematocrit, and White Blood Cell Count. Survival time was measured from admission to death or discharge, with approximately 62.18% of CAD patients exhibiting censored survival times. This censoring rate was replicated in the synthetic dataset to maintain realistic evaluation conditions. For the initial preprocessing, we employed the recently published data processing pipeline for MIMIC-IV by Mehak et al⁴⁴. After preprocessing, a total of 65,486 patient records were available for analysis.

Data generation

The processed data were used to simulate a two-stage DTR problem. Stage 1 focused on renal and anemia management using baseline covariates, while Stage 2 addressed glycemic control. The primary goal was to identify treatment policies that maximize survival outcomes. While we aimed to align the reward function, true policy, and other simulation-defined functions with clinical reasoning, we acknowledge that these may not fully represent the complexities of individual CAD treatment. Hence, the following functions should not be considered as ground truth for clinical decision-making.

First, we describe the binary ( M=2) treatment setting. At Stage 1, the treatment indicator was assigned using a multinomial distribution, with the following treatment probability vector: the (unnormalized) reference treatment scorewas set to 1, and the alternative treatment score was modeled as . These values were then normalized to ensure they sum to 1 and used to simulate treatment assignments for each subject. Importantly, these represent treatment propensity models and do not correspond to censoring probabilities, which are separately modeled and incorporated into the data generation process. The optimal treatment rule at Stage 1, , was determined using clinically relevant covariates. The optimal treatment rule was

10

For our simulation, we generated the Stage 1 outcome (which was the time until the next decision point, which could be either another treatment stage or terminating the process — perhaps due to an event like kidney failure, or death) using the Eq. 

11

where ~ Exp(λ) represents a noise term added to capture variability in survival times.

In Stage 2, treatment assignments () were generated based on updated covariates, including the first-stage outcome (). was assigned using a multinomial distribution. The reference treatment (unnormalized) probability remained fixed at 1, while the alternative treatment unnormalized probability was . The optimal treatment rule at the second stage () was defined similarly. The optimal treatment rule was

12

The second-stage outcome () was generated using the function

13

The methodology for defining treatment probabilities, optimal treatment assignment, and outcome functions follows recent works on synthetic data generation for evaluating DTR models^13,27, ensuring consistency with state-of-the-art approaches. For the optimal treatment rule and outcome (aka “reward”) function, we selected key features based on their well-established clinical relevance. For Stage 1, Creatinine (> 1.5 mg/dL) was chosen because elevated levels indicate significant renal impairment and are strongly linked to increased cardiovascular risk⁴⁵, while low Hemoglobin is a marker for anemia, which adversely affects cardiac outcomes^46–48. For Stage 2, we use a Glucose threshold of > 140 mg/dL to identify hyperglycemia that necessitates more aggressive glycemic management⁴⁹. In addition, routinely measured features such as Age, Potassium, Sodium, Platelet Count, Hematocrit, and White Blood Cell Count are included to capture a comprehensive cardiovascular risk profile^50–55.

In our survival outcome function, the exponential transformation is employed to ensure that the predicted survival times remain strictly positive — a fundamental requirement in time-to-event analyses. This transformation is standard practice in survival analysis, as seen in both the Cox proportional hazards model and accelerated failure time models, where exponential forms are used to capture the multiplicative effect of covariates on the baseline hazard^56,57. Moreover, the quadratic penalty term is incorporated to measure the deviation of the actual treatment  A from the optimal treatment . The penalty term in our survival outcome functions explicitly incorporates patient covariates because the optimal treatment policy is determined by patient-specific clinical characteristics. For example, patients with Creatinine levels > 1.5 mg/dL are optimally treated with more aggressive renal management strategies due to increased cardiovascular risk⁴⁵, while optimal glycemic management depends on glucose levels (> 140 mg/dL)⁴⁹. Thus, deviations from such covariate-dependent optimal policies naturally reflect increased clinical risk and reduced survival benefit, capturing realistic clinical scenarios where inappropriate treatments incur disproportionately higher risks. This approach of defining the outcome function is well supported by literature in DTRs for generating the synthetic data^13,27.

Censoring mechanisms were incorporated into the data generation process to mimic real-world survival datasets, where complete observation of survival times is often unavailable. Here, censoring distributions (exponential, conditional exponential, uniform) were selected to emulate diverse censoring mechanisms commonly encountered in real-world clinical studies⁵⁶. Specifically, exponential censoring assumes a constant hazard rate of censoring, modeling scenarios of random dropout unrelated to patient covariates; conditional exponential censoring introduces covariate-dependent censoring, realistically modeling situations where dropout or loss to follow-up is influenced by patient health status or disease severity; uniform censoring represents administrative censoring occurring due to predetermined study timelines, irrespective of patient characteristics or outcomes. The inclusion of these distinct censoring mechanisms enhances robustness evaluation of the proposed method across realistic clinical scenarios. The censoring times ( C) were generated using these three distinct distributions to reflect varying levels of censoring:

1. Exponential Censoring: Censoring times were drawn from an exponential distribution, , where was chosen to achieve the desired censoring rate.

2. Conditional Exponential Censoring: To introduce covariate dependency, censoring times followed .

3. Uniform Censoring: Censoring times were generated from a uniform distribution, , where  a and  b define the lower and upper bounds, respectively, and were chosen appropriately to achieve the desired censoring rate.

These distributions were calibrated to achieve a censoring rate of around 62%, consistent with the MIMIC CAD population. Here, the exponential and uniform censoring mechanisms represent non-informative censoring in the simulation because censoring is generated independently of patient covariates. In contrast, the conditional exponential mechanism represents covariate-depend (that is, marginally informative) censoring because the censoring distribution depends on observed covariates.

To determine whether a subject proceeds to the second stage, the binary indicator variable was used, where indicates that the first-stage outcome () is less than the censoring time (C), allowing progression to the second stage, and signifies censoring after the first stage. The total survival time () was defined as:

14

where represents the second-stage outcome.

The observed total survival time is the minimum of the total survival time () and the censoring time (C). The observed outcomes for each stage were defined as follows:

1. For the second stage, the observed outcome () was calculated as:

15

Here if the patient does not enter stage 2, is set as undefined (NaN).

• 2.For the first stage, the observed outcome is defined as:

16

To facilitate understanding of the notations used in results, we summarize key notations in Table 3. This framework for censoring and outcome generation ensures that the simulated data aligns with the complexities of survival analysis in multi-stage decision-making contexts. By integrating realistic censoring percentages and structured definitions of observed outcomes, this approach facilitates robust evaluation of the proposed method under various censoring scenarios. Note the learning algorithms will see only the observed information – here and – but not . As we have a large number of instances, we split the data into 5 folds, using one fold for training and the other 4 for evaluation each time to ensure reliable performance estimation. This choice of using a single fold for training allows us to simulate a more challenging scenario, where the model is trained on limited data. Because many real-world clinical DTR applications involve limited sample sizes^58–60, we also considered a deliberately low-training-data setting in the synthetic experiments. Specifically, we trained the policy using one fold and evaluated it on the remaining four folds. This serves as a stress test to explore whether the learned policy remains robust and generalizes to unseen data when training information is limited. Additionally, reserving a larger portion of the data for evaluation provides a more precise estimate of the model’s performance by minimizing the variability in test results across folds. For each experimental split, all nuisance-model fitting, tree construction, and hyperparameter tuning were performed using only the training portion of the data; the held-out instances were used exclusively for final performance assessment.

Table 3.

Key notations for synthetic data generation process.

Notation	Description	Used in equations
	Stage 1 treatment was assigned using a multinomial distribution	(11)
	Stage 2 treatment was assigned using a multinomial distribution	(13)
	Stage 1 optimal treatment	(10,11,17)
	Stage 2 optimal treatment	(12,13,18)
	Time spent in Stage 1 if no censoring happens	(11,12,13,14,15,16, 18)
	Time at Stage 2 if no censoring happens	(13,14)
T or	Total survival time if no censoring happens ()	(14)
For Censored Data
C	Censoring time: The time at which a patient’s survival is no longer observed
	: Equals 1 if the patient enters Stage 2, otherwise 0	(14,15)
Y	Observed total survival time due to censoring:	(15,16)
	Survival time at Stage 1 after adjusting for censoring	(16)
	Survival time at Stage 2 after adjusting for censoring	(15)

In addition to the binary treatment setting, we also consider a trinary  (M=3) treatment setting, which necessitates revised treatment probabilities and optimal decision rules. The reference treatment probability remained 1, while the additional treatment probabilities and were modeled as and , respectively. The optimal treatment rules were

17

18

Supplementary Information

Below is the link to the electronic supplementary material.

Author contributions

A.K.P. performed data curation, formal analysis, methodology, software, validation, visualisation, writing—original draft, and writing—review and editing. R.G. contributed to supervision, funding acquisition, resources, and writing—review and editing.

Data availability

The synthetic datasets utilized in this study were generated based on the publicly available MIMIC-IV database (version 2.0), accessible through PhysioNet at https://physionet.org/content/mimiciv/. Access to MIMIC-IV requires completion of a data use agreement and relevant training. The SANAD dataset employed in this research is publicly available via the JSM R package, which can be accessed at https://cran.r-project.org/web/packages/JSM/index.html.

Code availability

The implementation of baseline and comparison methods was based on publicly available resources. The CA-TReL algorithm was developed by extending the base code from the supplementary materials of the paper Tree-Based Reinforcement Learning for Estimating Optimal Dynamic Treatment Regimes (Tao et al., 2018), available at https://pmc.ncbi.nlm.nih.gov/articles/instance/6457899/bin/NIHMS988847-supplement-Supplemental_B.zip. The CA-TReL repository is publicly available at https://github.com/animeshkumarpaul/CA-TReL/tree/main. OWL and ASCL were obtained from the SurvDTR GitHub repository (https://github.com/taehwa015/SurvDTR), and DWSurv from its authors’ GitHub repository (https://github.com/gabriellesimoneau/Rcode-JASA2019). Interested parties can contact the corresponding author for further details on the implementation.

Declarations

Competing interests

The authors declare no competing interests.

Ethics

This study was approved by the University of Alberta Research Ethics Board (Pro0014785), which included the use of the MIMIC-IV dataset under a data use agreement and waived the need for individual patient consent due to the de-identified nature of the data. The SANAD dataset, available through the publicly accessible JSM R package, was used in accordance with its open-access terms and does not require ethics approval or informed consent.