Robust Reinforcement Learning for Dynamic Treatment Regimes under Model and Measurement Uncertainty
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 | Gain | Harm (%) | Benefit (%) | PF at (%) | 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 |
| Kernel-robust RL | ||||||
| Condition | Nominal RL | 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 | ||||||
| 45.74 (8.9) | 53.45 (0.9) | 53.31 (1.0) | 53.07 (0.9) | 19.04 (44.4) | 16.11 | |
| 44.20 (8.9) | 51.20 (0.9) | 51.06 (1.2) | 50.88 (1.0) | 19.13 (43.8) | 16.04 | |
| 39.12 (12.2) | 46.49 (0.9) | 46.57 (1.3) | 46.50 (1.1) | 17.92 (45.9) | 15.68 | |
| Trained on | Trained on | |||||
| Method | Nominal model | Nominal model | ||||
| 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, | 46.20 | 45.58 | 53.34 | 44.32 | 20.03 | 32.77 |
| Kernel-robust RL, | 46.46 | 46.03 | 52.84 | 43.31 | 20.79 | 31.90 |
| Kernel-robust RL, | 46.11 | 45.77 | 53.31 | 44.59 | 21.18 | 32.94 |
Appendix figures & tables6 assets
Supplementary material from the paper’s appendix.
Appendix
| Method | Nominal model | |
| 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) |
| Shifts in the dynamics | Errors in PSA | ||||
| Method | Nominal model | ||||
| 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, | 48.77 | 53.31 | 46.57 | 23.26 | 38.65 |
| 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 |
| Method | Seed | Nominal model | ||||
|---|---|---|---|---|---|---|
| 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 |
| Kernel-robust RL | ||||||
|---|---|---|---|---|---|---|
| Condition | Nominal RL | 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 | ||||||
| 46.14 (7.1) | 53.34 (0.7) | 52.84 (0.9) | 53.31 (0.9) | 17.82 (45.4) | 14.97 | |
| 42.08 (10.3) | 50.15 (0.9) | 50.03 (1.1) | 50.19 (1.4) | 17.72 (46.1) | 14.83 | |
| 33.99 (15.3) | 44.32 (1.3) | 43.31 (2.9) | 44.59 (2.0) | 16.91 (48.3) | 14.37 | |
| 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 |
Robust RL
Robust reinforcement learning develops agents that maintain performance and safety under uncertainty, disturbances, attacks, and distribution shifts.
See all 230 papersAdaptive Chemotherapy Control under Tumor Heterogeneity via Reinforcement Learning
Distributionally-Robust Policy Learning with Observational Data
KL, , and CVaR ambiguity sets. We establish two convergence guarantees. First, the estimated policy attains the standard rate on its DRO regret, and inherits the doubly-robust property -- the rate continues to hold whenever either the outcome model or the propensity model is consistently estimated, even at slow nonparametric rates. We then propose a sample-splitting bias-correction algorithm derived from the Lagrangian dual, and show that under a Tsybakov-style margin condition on the treatment-effect, the estimator attains a fast rate of , interpolating between and as the margin parameter ranges over . To our knowledge, this is the first work to incorporate a margin condition into the analysis of distributionally-robust policy learning, and the first to obtain fast rates in this setting. Empirically, on synthetic and semi-synthetic experiments calibrated to a welfare-to-work program, distributional robustness yields meaningful improvements over empirical welfare maximization when training and test distributions differ.