Robust Reinforcement Learning for Dynamic Treatment Regimes under Model and Measurement Uncertainty

Chengze Xie; Xuan Wang; Ying Ding; Hairong Wang; Yan Li

cs.LGOct 15, 2026

Robust Reinforcement Learning for Dynamic Treatment Regimes under Model and Measurement Uncertainty

Authors: Chengze Xie, Xuan Wang, Ying Ding, Hairong Wang, Yan Li

Organizations: Department of Industrial and Systems Engineering, Texas A&M University. · Department of Mechanical Engineering, The University of Texas at Austin. · School of Information, The University of Texas at Austin.

Abstract

We develop a robust reinforcement learning (RL) framework for learning adaptive treatment policies that generalize across patients and remain reliable under errors in the measured biomarker and shifts in the disease dynamics. By modeling the patient as a hidden context, we address generalization and robustness in a unified manner through robust contextual Markov decision processes. Applying our framework to adaptive therapy for metastatic prostate cancer in a simulator calibrated to clinical data, we obtain generalizable and robust policies that substantially outperform the existing adaptive therapy method. In particular, we attain robustness with almost no loss of performance when the biomarker is accurate and the disease evolves as anticipated, and find that the resulting policy Pareto-dominates the existing method across a diverse array of scenarios. Our findings suggest that robust RL holds great promise as a principled methodology for learning reliable adaptive treatment policies.

Figures & tables

Method	$\bar{G}$	Gain	Harm (%)	Benefit (%)	PF at $T$ (%)	Treated (%)	Courses
Virtual test population (1,000 patients)
Continuous treatment	15.81	–	–	–	2.7	100.0	1.0
AT50	29.58	+13.77	28.4	54.0	32.5	28.7	6.2
Nominal RL	50.35	+34.54	0.2	91.9	61.5	30.3	12.9
Fitted patients (32 patients)
Continuous treatment	14.66	–	–	–	0.0	100.0	1.0

Table 1: Nominal RL and the clinical rules in the nominal model, in terms of the mean progression-free time

\bar{G}

(months) and its gain over continuous treatment, the harm and benefit rates, the fraction of patients without progression after

T=65

months (PF at

T

), the fraction of treated months, and the mean number of treatment courses, i.e., maximal runs of consecutive treated months.

Figure 1: Nominal RL, AT50, and continuous treatment on the virtual test population, split into tertiles of the initial resistant fraction

R_{0}/B_{0}

. (a) Mean progression-free time. (b) Harm rate.

Figure 2: Policies on the virtual test population in the nominal model and under random perturbations of PSA with

\varepsilon=0.2

. (a) Mean progression-free time. (b) Harm rate. See Table 4 of Appendix C for the values.

		Kernel-robust RL
Condition	Nominal RL	$\rho_{\mathrm{tr}}=0.1$	$\rho_{\mathrm{tr}}=0.25$	$\rho_{\mathrm{tr}}=0.5$	AT50	Continuous treatment
Nominal model	50.35 (0.2)	48.13 (0.1)	48.77 (0.0)	48.45 (0.0)	29.58 (28.4)	15.81
Training kernels
$\rho=0$	45.74 (8.9)	53.45 (0.9)	53.31 (1.0)	53.07 (0.9)	19.04 (44.4)	16.11
$\rho=0.1$	44.20 (8.9)	51.20 (0.9)	51.06 (1.2)	50.88 (1.0)	19.13 (43.8)	16.04
$\rho=0.25$	39.12 (12.2)	46.49 (0.9)	46.57 (1.3)	46.50 (1.1)	17.92 (45.9)	15.68

Table 2: Mean progression-free time (months) under shifts in the patient dynamics on the virtual test population for policies trained on the calibrated population, with harm rates (%) in parentheses.

Figure 3: Mean progression-free time (top) and harm rate (bottom) on the virtual test population. The left panels show the nominal model, and the middle and right panels show the performance as a function of the evaluation radius

\rho

, (a) on the training kernels and (b) on the held-out kernels. Kernel-robust RL is shown for

\rho_{\mathrm{tr}}=0.25

Figure 4: Cross-evaluation on the virtual test population. Each axis shows the mean progression-free time divided by the highest value attained in the corresponding condition by the policies in this figure. Kernel shift denotes the evaluation radius

\rho

on the training kernels, and random perturbation denotes errors in PSA with

\varepsilon=0.2

. Kernel-robust RL is trained with

\rho_{\mathrm{tr}}=0.25

. See Table 5 of Appendix C for the values.

Figure 5: Mean progression-free time on the 32 fitted patients in the nominal model for policies trained on the calibrated population and on the broad population.

	Trained on $\mathcal{D}_{\mathrm{cal}}$	Trained on $\mathcal{D}_{\mathrm{wide}}$
Method	Nominal model	Nominal model	$\rho=0$	$\rho=0.25$	$\rho=1$	$\varepsilon=0.2$
Nominal RL	48.43	48.59	46.14	33.99	13.64	26.85
Random-perturbation RL	45.17	43.54	46.64	38.37	22.04	42.08
Kernel-robust RL, $\rho_{\mathrm{tr}}=0.1$	46.20	45.58	53.34	44.32	20.03	32.77
Kernel-robust RL, $\rho_{\mathrm{tr}}=0.25$	46.46	46.03	52.84	43.31	20.79	31.90
Kernel-robust RL, $\rho_{\mathrm{tr}}=0.5$	46.11	45.77	53.31	44.59	21.18	32.94

Table 3: Mean progression-free time (months) on the 32 fitted patients for policies trained on

\mathcal{D}_{\mathrm{cal}}

and on

\mathcal{D}_{\mathrm{wide}}

under the conditions of Table 5 .

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Method	Nominal model	$\varepsilon=0.2$
Continuous treatment	15.81	15.81
AT50	29.58 (28.4)	16.46 (43.8)
Nominal RL	50.35 (0.2)	32.96 (14.4)
Random-perturbation RL	45.94 (0.1)	44.49 (2.0)

Appendix

Table 4: Mean progression-free time (months) on the virtual test population in the nominal model and under random perturbations of PSA with

\varepsilon=0.2

, with harm rates (%) in parentheses.

		Shifts in the dynamics			Errors in PSA
Method	Nominal model	$\rho=0$	$\rho=0.25$	$\rho=1$	$\varepsilon=0.2$
Continuous treatment	15.81	16.11	15.68	12.72	15.81
AT50	29.58	19.04	17.92	11.65	16.46
Nominal RL	50.35	45.74	39.12	17.46	32.96
Random-perturbation RL	45.94	47.15	40.91	18.80	44.49
Kernel-robust RL, $\rho_{\mathrm{tr}}=0.25$	48.77	53.31	46.57	23.26	38.65

Appendix

Table 5: Cross-evaluation on the virtual test population (mean progression-free time in months). The columns

\rho\in\{0,0.25,1\}

denote evaluation radii on the training kernels, and the column

\varepsilon=0.2

denotes random perturbations of PSA.

Initial resistant fraction	Nominal RL	AT50	Continuous treatment
Virtual test population
Low	60.81 (0.4)	37.28 (39.9)	23.92
Middle	54.70 (0.1)	31.56 (34.5)	15.89
High	35.57 (0.0)	19.93 (10.8)	7.63
Fitted patients
Low	62.30 (0.0)	31.40 (50.0)	22.50

Appendix

Table 6: Mean progression-free time (months) by tertile of the initial resistant fraction

R_{0}/B_{0}

in the nominal model, with harm rates (%) in parentheses. Tertiles are computed separately for each test population.

Method	Seed	Nominal model	$\rho=0$	$\rho=0.25$	$\rho=1$	$\varepsilon=0.2$
Nominal RL	11	50.33	46.00	39.91	14.87	34.66
	23	50.32	46.25	39.56	15.97	31.85
	37	50.40	44.96	37.89	21.53	32.38
Random-perturbation RL	11	45.76	47.58	41.47	17.11	44.50
	23	46.12	47.35	41.51	18.99	44.46
	37	45.95	46.52	39.75	20.29	44.51

Appendix

Table 7: Per-seed mean progression-free time (months) on the virtual test population for the conditions of Table 5 .

		Kernel-robust RL
Condition	Nominal RL	$\rho_{\mathrm{tr}}=0.1$	$\rho_{\mathrm{tr}}=0.25$	$\rho_{\mathrm{tr}}=0.5$	AT50	Continuous treatment
Nominal model	48.59 (0.0)	45.58 (0.0)	46.03 (0.0)	45.77 (0.0)	25.34 (28.1)	14.66
Training kernels
$\rho=0$	46.14 (7.1)	53.34 (0.7)	52.84 (0.9)	53.31 (0.9)	17.82 (45.4)	14.97
$\rho=0.1$	42.08 (10.3)	50.15 (0.9)	50.03 (1.1)	50.19 (1.4)	17.72 (46.1)	14.83
$\rho=0.25$	33.99 (15.3)	44.32 (1.3)	43.31 (2.9)	44.59 (2.0)	16.91 (48.3)	14.37

Appendix

Table 8: Shifts in the patient dynamics on the 32 fitted patients, policies trained on the broad population. Notation as in Table 2 .

Population	Continuous	AT50	Forced only	Nominal RL	Seeds 11 / 23 / 37
Virtual	15.81	52.60	53.79	53.79	53.79 / 53.79 / 53.79
Fitted patients	14.66	53.06	53.75	53.75	53.75 / 53.75 / 53.75

Appendix

Table 9: Mean progression-free time (months) under the alternative definition of progression in the nominal model. Forced only denotes the rule that treats only when forced.

Robust RL

Robust reinforcement learning develops agents that maintain performance and safety under uncertainty, disturbances, attacks, and distribution shifts.

See all 231 papers

Robust Adversarial Policy Optimization Under Dynamics Uncertainty
Apr 13, 2026Mintae Kim, Koushil SreenathAdversarial MARL Reinforcement Learning
Hybrid Robotics, UC Berkeley Berkeley, CA 94720, USA
Adaptive Chemotherapy Control under Tumor Heterogeneity via Reinforcement Learning
Sep 10, 2026Bereket Sitotaw Kidane, Md Samiul Haque Motayed, Shuo WangRL Control Reinforcement Learning
Department of Mechanical and Aerospace Engineering, University of Texas at Arlington, Arlington, TX 76019, USA
Distributionally-Robust Policy Learning with Observational Data
Oct 18, 2026Debmalya MandalPolicy Learning Distribution Shift Robustness
Dept. of Computer Science University of Warwick, UK

Apr 13, 2026J/K move · Enter open · S save

Robust Adversarial Policy Optimization Under Dynamics Uncertainty

Mintae Kim, Koushil Sreenath

Hybrid Robotics, UC Berkeley Berkeley, CA 94720, USA

Adversarial MARL Reinforcement Learning Robust RL Distributionally Robust RL Generalization in RL

Open paper AskPDF arXiv

Reinforcement learning (RL) policies often fail under dynamics that differ from training, a gap not fully addressed by domain randomization or existing adversarial RL methods. Distributionally robust RL provides a formal remedy but still relies on surrogate adversaries to approximate intractable primal problems, leaving blind spots that potentially cause instability and over-conservatism. We propose a dual formulation that directly exposes the robustness-performance trade-off. At the trajectory level, a temperature parameter from the dual problem is approximated with an adversarial network, yielding efficient and stable worst-case rollouts within a divergence bound. At the model level, we employ Boltzmann reweighting over dynamics ensembles, focusing on more adverse environments to the current policy rather than uniform sampling. The two components act independently and complement each other: trajectory-level steering ensures robust rollouts, while model-level sampling provides policy-sensitive coverage of adverse dynamics. The resulting framework, robust adversarial policy optimization (RAPO) outperforms robust RL baselines, improving resilience to uncertainty and generalization to out-of-distribution dynamics while maintaining dual tractability.

Abstract

Robust RL

Explore similar work

Robust Adversarial Policy Optimization Under Dynamics Uncertainty

Adaptive Chemotherapy Control under Tumor Heterogeneity via Reinforcement Learning

Distributionally-Robust Policy Learning with Observational Data