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

TL;DR

This paper develops a robust offline reinforcement learning framework for multi-site data with heterogeneity, introducing group-structured distributionally robust MDPs and an algorithm that leverages shared features and site clustering.

Contribution

It proposes a novel group-structured distributionally robust MDP model with tractable Bellman recursions and an offline algorithm incorporating site-specific ridge regression, worst-case aggregation, and clustering extensions.

Findings

Provides a suboptimality bound under partial coverage assumptions.

Introduces feature-wise uncertainty sets preserving tractability.

Demonstrates improved sample efficiency with site clustering.

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.