📊 statistics

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

This paper proposes a group-robust Markov Decision Process framework with feature-wise uncertainty sets and a pessimistic offline algorithm to learn robust sequential decision policies from heterogeneous multi-site data, achieving suboptimality guarantees without relying on strong state-action rectangularity assumptions.

Original authors: Mingyuan Xu, Zongqi Xia, Tianxi Cai, Doudou Zhou, Nian Si

Published 2026-02-03

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Mingyuan Xu, Zongqi Xia, Tianxi Cai, Doudou Zhou, Nian Si

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot how to navigate a complex city to deliver packages. You don't have time to let the robot drive around and crash into things (that's "online" learning, which is dangerous and expensive). Instead, you give the robot a massive library of driving logs from three different cities: New York, Chicago, and Miami.

Here is the problem:

New York has lots of traffic lights and narrow streets.
Chicago has huge open avenues but icy winters.
Miami has heavy rain and different traffic rules.

If you just mash all these logs together into one giant pile, the robot might learn a "middle-of-the-road" strategy that works okay in the average city but fails miserably in the worst-case scenario (like getting stuck in a Chicago blizzard). This is called distributional shift.

If you teach the robot separately for each city, it might become an expert in New York but clueless in Miami, or it might get confused because there isn't enough data in any single city to be sure of the rules.

This paper proposes a smart way to teach the robot using all three cities' data while preparing it for the worst possible version of any city.

The Core Idea: "The Worst-Case Weatherman"

The authors treat the learning process like a game between two characters:

The Robot (The Agent): Wants to find the best route to deliver packages.
The Adversary (The Weatherman): Wants to make the robot's life as hard as possible by picking the worst possible traffic conditions or road rules from the data it has seen.

Usually, in these games, the Weatherman can change the rules for every single street corner independently. This makes the math impossible to solve (it's like trying to predict the weather for every atom in the atmosphere at once).

The Paper's Trick:
The authors introduce a clever shortcut called "Feature-wise Rectangularity."
Instead of letting the Weatherman change every single rule independently, they say: "Okay, Weatherman, you can change the rules for 'Traffic Lights,' 'Road Width,' and 'Weather' independently, but you have to apply the same 'worst-case' logic to all of them together."

Think of it like a menu of ingredients.

Old Way: The chef (Weatherman) can swap out the salt in the soup, the sugar in the cake, and the spice in the stew independently for every single dish. This is chaotic and hard to plan for.
New Way (This Paper): The chef can swap the salt, sugar, and spice, but they must do it in a structured way that respects the "flavor profile" of the dish. This keeps the math solvable while still being very cautious.

How the Algorithm Works: "The Cautious Chef"

The paper's algorithm (Algorithm 1) works in three steps, like a cautious chef preparing a meal for a picky guest:

Learn from Each City Separately (Ridge Regression):
First, the robot looks at the logs from New York, Chicago, and Miami separately. It tries to guess the rules for each city. But because the data might be messy or incomplete, it adds a "safety buffer" (called pessimism) to its guesses. It assumes the data might be slightly wrong.
The "Worst-Case" Mix (Row-wise Minimization):
Now, the robot combines these guesses. Instead of averaging them (which would hide the bad parts), it looks at every single rule and asks: "What is the worst version of this rule across all three cities?"
- If New York says "Speed limit is 30," Chicago says "25," and Miami says "35," the robot assumes the speed limit is 25.
- It builds a policy based on the lowest (safest) estimate for every single feature. This ensures that no matter which city's hidden "worst-case" reality shows up, the robot won't crash.
The Safety Penalty:
If the robot hasn't seen a specific situation enough times in the logs (e.g., it only saw 5 rainy days in Miami), the algorithm adds a huge "penalty" to that guess. It tells the robot: "Don't trust this number; you don't have enough data. Assume the worst." This prevents the robot from getting overconfident based on lucky, small samples.

The "Group" Strategy: Clustering Similar Cities

The paper also suggests a second trick. What if you have 50 cities, but 10 of them are very similar (e.g., all coastal towns)?
Instead of treating them as 10 separate problems, you can pool them together into one "Super-Coastal" group.

Why? It gives you more data to learn the rules of "Coastal Driving."
The Catch: You have to make sure the cities are actually similar. If you pool a desert city with a coastal city, your "Super-Group" rules will be nonsense. The paper provides math to prove that as long as the cities in the group are similar enough, pooling them makes the robot learn faster and more accurately.

The Results: Why It Matters

The authors tested this on computer simulations:

Naive Pooling: Just mixing all data together. Result: The robot failed in the worst scenarios because it ignored the unique dangers of specific cities.
Separate Learning: Learning for each city alone. Result: The robot was unstable and made mistakes because it didn't have enough data for any single city.
This Paper's Method: Result: The robot learned a policy that was consistently safe and efficient, even in the worst-case scenarios. It found the "sweet spot" between being too cautious and being too reckless.

In a Nutshell

This paper gives us a mathematical recipe for learning from multiple different sources (like hospitals, cities, or factories) without needing to trust that they are all exactly the same. It builds a decision-making system that is robust: it prepares for the worst possible version of the data it has seen, ensuring that the final plan works safely even when things go wrong or data is missing.

It's like training a pilot not just on "average weather," but by simulating the worst combination of wind, rain, and turbulence found in any of the training logs, ensuring they can land safely no matter what happens.

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

Problem Formulation

This paper addresses the challenge of learning robust sequential decision-making policies from offline datasets collected across multiple sites (e.g., hospitals) that share a common application context but exhibit heterogeneity. In standard offline Reinforcement Learning (RL), policies are learned from pre-collected data without further environment interaction. However, real-world multi-site data often suffers from two critical issues:

Distributional Shift: Offline trajectories are generated by behavior policies that may not cover the state-action distribution of the optimal policy.
Site-wise Heterogeneity: Variations in data acquisition protocols, population composition, and environmental conditions mean that trajectories are not generated from a single, fixed Markov Decision Process (MDP). Instead, each site $k$ has its own transition kernel $P_k$ and reward function $r_k$ .

The authors model this setting as a Group-Linear Distributionally Robust MDP (DRMDP). The underlying assumption is that all sites share a common, known feature map $\phi: \mathcal{S} \times \mathcal{A} \to \mathbb{R}^d$ , but the transition kernels and reward functions are linear in these features with site-specific parameters $(\theta^k_h, \mu^k_h)$ . The goal is to learn a policy that performs well under the worst-case environmental shift within a plausible uncertainty set defined by the convex hull of site-specific distributions.

Methodology

1. Feature-wise (d-rectangular) Uncertainty Sets

To handle the uncertainty arising from multiple sites while maintaining computational tractability, the authors propose a feature-wise (d-rectangular) uncertainty set.

Structure: Unlike a standard cross-site convex hull which couples uncertainty across all features (leading to non-rectangular sets and NP-hard planning), this approach allows the adversary to select mixing weights independently for each feature dimension $i \in [d]$ .
Definition: For each step $h$ and feature $i$ , the adversary chooses a mixing vector $\alpha_{h,i} \in \Delta^{K-1}$ to form a worst-case mixture of the site-specific parameters. The uncertainty set $\mathcal{U}_h$ is defined such that the reward and transition for any $(s, a)$ are linear combinations of site parameters weighted by these feature-specific mixtures.
Benefit: This relaxation preserves the key cross-site heterogeneity while restoring the robust Bellman recursions, making the problem tractable.

2. Pessimistic Value Iteration Algorithm

The authors develop an offline algorithm (Algorithm 1) based on pessimistic value iteration, which includes three core components:

Per-Site Ridge Regression: For each site $k$ , the algorithm estimates the Bellman target coefficients $\hat{\nu}^k_h$ using ridge regression on the local dataset $D_k$ . This yields site-specific estimates of the linear parameters.
Feature-wise Worst-Case Aggregation: Instead of averaging, the algorithm performs a row-wise minimum across the estimated coefficients of all sites. Specifically, for each feature dimension $i$ , the robust coefficient is set to $\min_{k \in [K]} (\hat{\nu}^k_h)_i$ . This captures the worst-case scenario across sites for each feature.
Data-Dependent Pessimism Penalty: To account for estimation uncertainty and distributional shift, a penalty term is subtracted from the Q-value. This penalty is computed using the diagonals of the inverse design matrices (covariance matrices) from the ridge regression. The penalty scales with the feature vector $\phi(s, a)$ and the maximum uncertainty observed across sites for each feature dimension.

3. Cluster-Level Extension

To improve sample efficiency when sites exhibit similarity, the authors propose a cluster-level extension. Sites are partitioned into clusters based on prior knowledge of similarity. Within each cluster, datasets are pooled to form a "super-site," and the robust Bellman operator is applied to the cluster-level representatives. This approach leverages larger pooled sample sizes while bounding the error introduced by within-cluster heterogeneity.

Theoretical Results

The paper provides finite-sample suboptimality bounds for the proposed algorithm under a robust partial coverage assumption.

Suboptimality Decomposition: The suboptimality is decomposed into spurious correlation, intrinsic uncertainty, and optimization error.
Convergence Rate: Under the assumption that the site-specific empirical covariance matrices dominate the second moments of the robust optimal policy's visitation distribution, the suboptimality of the learned policy $\hat{\pi}$ is bounded by:
$\text{SubOpt}(\hat{\pi}; s) \leq O\left( d^2 H^2 \sqrt{\frac{\log(d K H N_{\max}/\xi)}{N_{\min}}} \right)$
where $N_{\min}$ is the minimum sample size across sites, $d$ is the feature dimension, $H$ is the horizon, and $K$ is the number of sites.
Cluster-Level Bound: For the cluster-level algorithm, the bound includes an additional bias term proportional to the within-cluster parameter discrepancy ( $\Delta_h$ ), demonstrating that pooling is beneficial as long as the heterogeneity within clusters is controlled.

Experimental Results

The authors validate their approach through simulations on episodic linear MDPs with $K=3$ to $4$ sites.

Comparison: The proposed "Sitewise" algorithm is compared against:
- Naive Pooling: Aggregating all data and applying standard Pessimistic Value Iteration (PEVI).
- Per-Site PEVI (Average-Q): Averaging Q-functions from independent site-specific PEVI runs.
- Per-Site PEVI (Min-Q): Taking the pointwise minimum of site-specific Q-functions.
Performance: The proposed algorithm consistently achieves significantly lower robust suboptimality compared to all benchmarks.
- Naive Pooling fails to be robust against site-specific distributional shifts, optimizing for average performance rather than worst-case.
- Per-Site methods suffer from high variance due to limited local sample sizes and lack of information sharing.
Convergence: Empirical scaling laws confirm the theoretical $O(N_{\min}^{-1/2})$ convergence rate.
Necessity of Pessimism: Experiments show that removing the pessimism penalty (setting the scaling constant $c=0$ ) leads to a failure to converge efficiently, as the greedy estimator overfits to finite-sample noise, particularly in scenarios with imbalanced data coverage.

Significance and Claims

The paper claims to provide a principled approach to robust planning in multi-site offline RL settings without relying on strong state-action rectangularity assumptions.

Tractability vs. Robustness: By introducing the $d$ -rectangular uncertainty set, the authors bridge the gap between the intractability of general cross-site convex hulls and the excessive conservatism of standard $(s, a)$ -rectangular robust MDPs.
Privacy and Efficiency: The algorithm communicates only aggregated site-level summary statistics (ridge regression coefficients and covariance diagonals) rather than raw trajectories, which reduces data communication costs and enhances privacy.
Handling Heterogeneity: The framework explicitly models and leverages the heterogeneity across sites to ensure the learned policy is robust to the worst-case site configuration, addressing a fundamental limitation of standard offline RL when applied to multi-source data.

The authors acknowledge that their method relies on a linear structure within sites and that the $d$ -rectangular approximation may lead to conservative estimates. They suggest that future work could explore tighter approximations and the integration of representation learning to handle non-linear feature mappings.