Learning Sequential Decisions from Multiple Sources via Group-Robust Markov Decision Processes
Organizations: Department of Statistics and Data Science, National University of Singapore · Department of Neurology, University of Pittsburgh · Department of Biostatistics, Harvard T.H. Chan School of Public Health · Department of Biomedical Informatics, Harvard Medical School ·
Abstract
Learning a common policy from multiple source environments requires accounting for cross-source heterogeneity, and the offline setting introduces an additional challenge of limited data coverage. We formulate this problem as a distributionally robust Markov decision process under a shared linear MDP representation with source-specific transition and reward components. The uncertainty set is constructed through feature-wise mixtures of these components, yielding a -rectangular structure that captures cross-source heterogeneity and admits a robust Bellman recursion. Building on this formulation, we develop Group-Robust Pessimistic Value Iteration (GR-PEVI), an offline policy-learning algorithm that combines feature-wise worst-case aggregation with a data-dependent pessimism penalty. Under a robust partial coverage condition, we establish finite-sample suboptimality guarantees for GR-PEVI and extend the analysis to clustered sources, characterizing the trade-off between the statistical gain from pooling similar sources and the approximation error induced by within-cluster heterogeneity. Simulation studies demonstrate the effectiveness of GR-PEVI, with empirical behavior consistent with the theoretical analysis. We apply GR-PEVI to a multi-source electronic health record study of multiple sclerosis treatment decisions to illustrate its use in a clinical setting.
Figures & tables
| Characteristic | Overall | MGB-young | MGB-old | UPMC |
|---|---|---|---|---|
| Patients, | 5,700 | 2,370 | 1,925 | 1,405 |
| Female sex, (%) | 4,244 (74.5%) | 1,750 (73.8%) | 1,398 (72.6%) | 1,096 (78.0%) |
| Age, mean (SD) | 40.9 (11.4) | 32.3 (6.5) | 51.3 (6.7) | 41.2 (11.0) |
| White race, (%) | 4,943 (86.9%) | 1,941 (81.9%) | 1,748 (90.8%) | 1,254 (90.0%) |
| Standard-efficacy DMT, (%) | 4,357 (76.4%) | 1,851 (78.1%) | 1,459 (75.8%) | 1,047 (74.5%) |
| High-efficacy DMT, (%) | 1,343 (23.6%) | 519 (21.9%) | 466 (24.2%) | 358 (25.5%) |
| Method | Mean | SD | Median | Win count | GR-PEVI win rate |
|---|---|---|---|---|---|
| GR-PEVI | 3.2042 | 0.0760 | 3.2107 | 30/50 | – |
| Operator-Avg PEVI | 3.1535 | 0.0803 | 3.1594 | 7/50 | 72% |
| Avg-PEVI | 3.1463 | 0.0814 | 3.1364 | 6/50 | 82% |
| Majority-PEVI | 3.1319 | 0.1050 | 3.1287 | 7/50 | 76% |
| Pooled PEVI | 3.0151 | 0.0803 | 3.0276 | 0/50 | 98% |
| Method | MGB-young | MGB-old | UPMC |
|---|---|---|---|
| GR-PEVI | (1.85) | (2.51) | (4.34) |
| Operator-Avg PEVI | (2.08) | (2.72) | (4.76) |
| Avg-PEVI | (2.10) | (2.98) | (5.58) |
| Majority-PEVI | (2.57) | (3.69) | (6.11) |
| Pooled PEVI | (2.18) | (2.25) | (3.30) |
| Method | Main setting | Milder transition heterogeneity |
|---|---|---|
| GR-PEVI | 0.017 (0.007) | 0.020 (0.006) |
| Pooled PEVI | 0.329 (0.042) | 0.274 (0.034) |
| Avg-PEVI | 0.951 (0.021) | 0.872 (0.020) |
| Majority-PEVI | 0.895 (0.022) | 0.778 (0.025) |
| Operator-Avg PEVI | 0.242 (0.021) | 0.196 (0.015) |
| Generic name | Brand name(s) | RxNorm ID |
| High efficacy | ||
| Natalizumab | Tysabri | 354770 |
| Ocrelizumab | Ocrevus | 1876366 |
| Ofatumumab | Arzerra; Kesimpta | 712566 |
| Mitoxantrone | Novantrone | 7005 |
| Alemtuzumab | Campath; Lemtrada | 117055 |
| Method | Mean | SD | Median | Win count | GR-PEVI win rate |
|---|---|---|---|---|---|
| GR-PEVI | 4.6629 | 0.0247 | 4.6576 | 38/50 | – |
| Operator-Avg PEVI | 4.6345 | 0.0128 | 4.6369 | 6/50 | 86% |
| Majority-PEVI | 4.6269 | 0.0237 | 4.6277 | 4/50 | 84% |
| Avg-PEVI | 4.6095 | 0.0348 | 4.6187 | 2/50 | 92% |
| Pooled PEVI | 4.3767 | 0.0577 | 4.3824 | 0/50 | 100% |
| Method | MICU/SICU | TSICU | SICU | MICU | CCU | CVICU |
|---|---|---|---|---|---|---|
| GR-PEVI | (0.65) | (0.89) | (0.71) | (0.55) | (0.59) | (0.62) |
| Operator-Avg PEVI | (0.25) | (0.45) | (0.39) | (0.24) | (0.79) | (0.30) |
| Avg-PEVI | (0.34) | (0.96) | (0.73) | (0.45) | (1.15) | (0.44) |
| Majority-PEVI | (0.28) | (0.68) | (0.51) | (0.31) | (0.83) | (0.31) |
| Pooled PEVI | (0.43) | (1.24) | (0.86) | (0.46) | (0.82) | (0.54) |
Robust Markov Decision Processes
Robust Markov decision processes optimize decisions under uncertain dynamics, seeking reliable performance against adverse transition models.
See all 27 papersReinforcement Learning for Individual Optimal Policy from Heterogeneous Data
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.