Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Learning Sequential Decisions from Multiple Sources via Group-Robust Markov Decision Processes

Mingyuan Xu, Zongqi Xia, Tianxi Cai, Doudou Zhou, Nian Si

TL;DR

The paper tackles offline reinforcement learning when data come from multiple heterogeneous sites, causing distributional shift and cross-site variability. It introduces a group-linear distributionally robust MDP with a $d$-rectangular, feature-wise uncertainty set that preserves cross-site structure and yields tractable robust Bellman recursions. The authors design a pessimistic offline algorithm that performs per-site ridge regression for Bellman targets, aggregates via rowwise site-wise minima, and applies a data-dependent pessimism penalty, with a cluster-level pooling extension to boost sample efficiency; they also prove a suboptimality bound under a robust partial coverage condition. Empirical results on multi-site simulations show the proposed method outperforms naive pooling and per-site baselines, with favorable convergence rates and robustness to distributional shifts. The framework offers principled, scalable planning across heterogeneous data sources and highlights the trade-offs between conservatism and statistical efficiency in robust multi-site RL.

Abstract

We often collect data from multiple sites (e.g., hospitals) that share common structure but also exhibit heterogeneity. This paper aims to learn robust sequential decision-making policies from such offline, multi-site datasets. To model cross-site uncertainty, we study distributionally robust MDPs with a group-linear structure: all sites share a common feature map, and both the transition kernels and expected reward functions are linear in these shared features. We introduce feature-wise (d-rectangular) uncertainty sets, which preserve tractable robust Bellman recursions while maintaining key cross-site structure. Building on this, we then develop an offline algorithm based on pessimistic value iteration that includes: (i) per-site ridge regression for Bellman targets, (ii) feature-wise worst-case (row-wise minimization) aggregation, and (iii) a data-dependent pessimism penalty computed from the diagonals of the inverse design matrices. We further propose a cluster-level extension that pools similar sites to improve sample efficiency, guided by prior knowledge of site similarity. Under a robust partial coverage assumption, we prove a suboptimality bound for the resulting policy. Overall, our framework addresses multi-site learning with heterogeneous data sources and provides a principled approach to robust planning without relying on strong state-action rectangularity assumptions.

Learning Sequential Decisions from Multiple Sources via Group-Robust Markov Decision Processes

TL;DR

The paper tackles offline reinforcement learning when data come from multiple heterogeneous sites, causing distributional shift and cross-site variability. It introduces a group-linear distributionally robust MDP with a -rectangular, feature-wise uncertainty set that preserves cross-site structure and yields tractable robust Bellman recursions. The authors design a pessimistic offline algorithm that performs per-site ridge regression for Bellman targets, aggregates via rowwise site-wise minima, and applies a data-dependent pessimism penalty, with a cluster-level pooling extension to boost sample efficiency; they also prove a suboptimality bound under a robust partial coverage condition. Empirical results on multi-site simulations show the proposed method outperforms naive pooling and per-site baselines, with favorable convergence rates and robustness to distributional shifts. The framework offers principled, scalable planning across heterogeneous data sources and highlights the trade-offs between conservatism and statistical efficiency in robust multi-site RL.

Abstract

We often collect data from multiple sites (e.g., hospitals) that share common structure but also exhibit heterogeneity. This paper aims to learn robust sequential decision-making policies from such offline, multi-site datasets. To model cross-site uncertainty, we study distributionally robust MDPs with a group-linear structure: all sites share a common feature map, and both the transition kernels and expected reward functions are linear in these shared features. We introduce feature-wise (d-rectangular) uncertainty sets, which preserve tractable robust Bellman recursions while maintaining key cross-site structure. Building on this, we then develop an offline algorithm based on pessimistic value iteration that includes: (i) per-site ridge regression for Bellman targets, (ii) feature-wise worst-case (row-wise minimization) aggregation, and (iii) a data-dependent pessimism penalty computed from the diagonals of the inverse design matrices. We further propose a cluster-level extension that pools similar sites to improve sample efficiency, guided by prior knowledge of site similarity. Under a robust partial coverage assumption, we prove a suboptimality bound for the resulting policy. Overall, our framework addresses multi-site learning with heterogeneous data sources and provides a principled approach to robust planning without relying on strong state-action rectangularity assumptions.
Paper Structure (16 sections, 5 theorems, 56 equations, 3 figures, 1 algorithm)

This paper contains 16 sections, 5 theorems, 56 equations, 3 figures, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Our Contributions
  4. Model and Method
  5. Multi-Site Linear MDP and Offline Data
  6. Robust Formulation and Learning Objective
  7. Algorithm
  8. Theory
  9. Decomposition of Suboptimality for Robust MDP
  10. Suboptimality Upper Bound
  11. Cross-Site Similarity
  12. Simulation
  13. Multi-site Performance Comparison
  14. Convergence Rate
  15. The Necessity of the Pessimism Penalty
  16. ...and 1 more sections

Key Result

Proposition 1

Under the $d$-rectangular uncertainty sets $\{\mathcal{U}_h\}_{h=1}^H$, for any policy $\pi$, the functions $V_h^\pi$ and $Q_h^\pi$ satisfy the following recursive relationships:

Figures (3)

  • Figure 1: Comparison of suboptimality between the proposed Sitewise algorithm and benchmark methods. (A) Distribution of suboptimality over 50 runs. (B) Mean and 95% confidence intervals of the average suboptimality values. The Sitewise method demonstrates significantly lower suboptimality and higher stability.
  • Figure 2: Scaling relationships between minimum sample size and performance metrics. (A) Suboptimality decreases with $N_{\min}$ following a power law (slope = $-0.45\pm0.03$, $R^2 = 0.995$). (B) Value Gap shows similar power-law decay (slope = $-0.62 \pm 0.03$, $R^2 = 0.997$). Data points represent the mean $\pm$ 95% CI across independent runs. Dashed lines show fitted power-law relationships in log-log space.
  • Figure 3: Convergence comparison between estimators without and with pessimism penalty. The estimator without penalty ($c=0$, dark blue) fails to converge efficiently, exhibiting a plateau with a flat slope ($-0.15 \pm 0.09$). Incorporating the pessimism penalty ($c>0$) restores the power-law decay in suboptimality. Note: For the strong penalty case ($c=0.001$, orange), the algorithm achieves exact optimal policy (zero suboptimality) at the largest sample size ($N_{\min} = 2 \times 10^6$). The reported slope ($-0.88 \pm 0.23$) is thus calculated using the first $5$ (pre-convergence) data points.

Theorems & Definitions (9)

  • Remark 1: Simplex features
  • Proposition 1: Robust Bellman Equations
  • Definition 1: $\xi$-uncertainty quantifier jin2021pessimism
  • Lemma 2: Decomposition of Suboptimality for Robust MDP
  • Theorem 3: Suboptimality
  • Corollary 4: Suboptimality with partial coverage
  • Remark 2
  • Corollary 5: Cluster-Level Suboptimality
  • Remark 3