Distributionally-Robust Policy Learning with Observational Data

Debmalya Mandal

cs.LGOct 18, 2026

Distributionally-Robust Policy Learning with Observational Data

Authors: Debmalya Mandal

Organizations: Dept. of Computer Science University of Warwick, UK

Abstract

Policy learning from observational data is fundamental to applications ranging from personalized medicine to welfare-program targeting, but standard approaches optimize empirical performance and may fail when the deployment distribution differs from the training distribution. We study distributionally-robust policy learning: given observational data, we seek a policy $π\in Π$ that maximizes the worst-case policy value over an $f$ -divergence ball of radius $ρ$ around the data-generating distribution. We propose an estimator built on cross-fitted doubly-robust scores for the conditional treatment effect, paired with a DRO solver that handles KL, $χ^2$ , and CVaR ambiguity sets. We establish two convergence guarantees. First, the estimated policy attains the standard $\widetilde{O}(n^{-1/2})$ 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 $\widetilde{O}(n^{-(1+α)/(2+α)})$ , interpolating between $\widetilde{O}(1/\sqrt{n})$ and $\widetilde{O}(1/n)$ as the margin parameter $α$ ranges over $[0,\infty)$ . 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.

Figures & tables

\widehat{m}\in\argmin_{m}\max_{Q:D_{f}(Q\lVert\hat{P}_{n})\leq\rho}\sum_{i=n+1}^{2n}q_{i}\cdot(m(X_{i},W_{i})-Y_{i})^{2}

Algorithm 1 Distributionally Robust Policy Learning from Observational Data

\widehat{m}_{g}\in\argmin_{m}\max_{Q:D_{f}(Q\lVert\hat{P}_{n_{g}})\leq\rho}\sum_{i\in\mathcal{D}_{2,g}}q_{i}\cdot(m(X_{i},W_{i})-Y_{i})^{2}

Algorithm 2 Group Robust Policy Learning from Observational Data

Figure 1: Out-of-Distribution Robustness on Semi-Synthetic Data: We observe that $\text{DR}(\textrm{CVaR})$ performs significantly better than standard empirical risk minimizer (ERM) on both training and testing counties.

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

County	Role	$n_{c}$	$p_{c}$	$\mu_{\mathrm{age}}^{c}$	$\mu_{\mathrm{edu}}^{c}$	$\mu_{\mathrm{earn}}^{c}$
Riverside	train	400	0.50	32	12.0	2.0
San Diego	train	400	0.52	34	12.5	2.4
Tulare	train	400	0.48	36	11.0	1.6
Alameda	test	1500	0.50	33	12.2	2.5
Butte	test	1500	0.55	35	11.5	1.8
Los Angeles	test	1500	0.45	34	10.6	2.1

Appendix

Table 1: County-level demographic and propensity parameters.

County	Role	$\tau_{0}^{c}$	$\tau_{\mathrm{age}}^{c}$	$\tau_{\mathrm{edu}}^{c}$	$\tau_{\mathrm{edu}^{2}}^{c}$
Riverside	train	$+0.10$	$-0.001$	$+0.12$	$+0.005$
San Diego	train	$+0.08$	$-0.001$	$+0.10$	$+0.005$
Tulare	train	$+0.07$	$-0.002$	$+0.08$	$+0.005$
Alameda	test	$+0.05$	$-0.001$	$-0.10$	$-0.005$
Butte	test	$+0.03$	$-0.002$	$-0.08$	$-0.005$
Los Angeles	test	$+0.05$	$-0.001$	$-0.12$	$-0.005$

Appendix

Table 2: County-specific treatment-effect coefficients.

Policy Learning

Policy learning develops decision rules that choose actions or interventions to maximize outcomes from data, feedback, or interaction.

See all 484 papers

Jun 17, 2026

Wasserstein Policy Learning for Distributional Outcomes

Offline policy learning has received growing attention in causal inference. The primary objective is to learn a policy (individualized treatment rule) as a mapping from covariates to treatment that maximizes the empirical welfare defined as the mean of scalar-valued potential outcomes. In this paper, we study offline policy learning with distribution-valued outcomes, where each potential outcome is a probability measure on

\mathbb{R}

and the reward is defined through a utility functional applied to the Wasserstein barycenter of induced outcome distributions. We establish statistical guarantees for the policy learning framework based on both Inverse Probability Weighting (IPW) and Doubly Robust (DR) estimators. By handling the challenging uniform deviation over the product of the combinatorial policy class and the infinite-dimensional quantile domain, we prove that the finite-sample regret has leading dependence

\widetilde{\mathcal{O}}(\sqrt{\mathrm{N\text{-}dim}(Π)/N})

. In the one-dimensional Wasserstein setting and under the stated regularity conditions, the leading regret rate is still governed by the policy-class complexity. Moreover, we provide a minimax lower bound establishing the sharpness of the leading dependence on

N

and

\mathrm{N\text{-}dim}(Π)

Yiyan Huang, Cheuk Hang Leung, Qi Wu +1

School of Statistics and Data Science & Institute of Big Data Research, Shanghai University of Finance and Economics, Shanghai, China

Offline RL Wasserstein Barycenters Causal Inference

Ask PDF

Sep 26, 2026

Distributionally Robust Average-Reward Reinforcement Learning: Finite-Sample Guarantees under Weak Communication

We study distributionally robust reinforcement learning (DR-RL) in the average-reward setting under weak communication. Our main result provides finite-sample guarantees for estimating the robust optimal average reward and learning a near-optimal policy, covering both SA-rectangular and S-rectangular structures with divergence-based and distance-based uncertainty sets. Specifically, for Kullback--Leibler and

f_k

-divergence balls, we establish explicit radius conditions under which the robust average-reward Bellman equation admits a constant-gain solution, while for total variation and Wasserstein balls, any positive radius suffices without requiring the nominal MDP to be weakly communicating. Our algorithm is prior-knowledge-free and achieves sample complexities of

\widetilde O(|\mathcal{S}||\mathcal{A}|p_{\wedge}^{-1}\operatorname{Span}^{2}(u_δ^{\ast})ε^{-2})

for estimating the robust optimal average reward and

\widetilde O(|\mathcal{S}||\mathcal{A}|p_{\wedge}^{-2}\operatorname{Span}^{2}(u_δ^{\ast})ε^{-2})

for learning an

ε

-optimal policy. Here,

p_{\wedge}

is the smallest positive nominal transition probability and

u_δ^{\ast}

is a robust optimal bias function. We further provide an almost-tight explicit upper bound on

\operatorname{Span}(u_δ^{\ast})

. Finally, we validate the predicted

n^{-1/2}

convergence rate through numerical experiments.

Chenyu Lu, Zijun Chen, Nian Si

The Hong Kong University of Science and Technology

Distributionally Robust RL Average-Reward RL Sample-Efficient RL

Ask PDF

Jul 15, 2026

Learning Who to Treat When Treatment is Missing

Policy learning methods are increasingly used to inform treatment allocation under budget constraints. Most proposed methods assume complete treatment data, yet applications frequently suffer from missingness that can bias estimates and lead to suboptimal policies. We address this gap by extending efficient estimators for average treatment effect (ATE) estimation to policy value and conditional average treatment effect (CATE) estimation under missing at random (MAR) and missing completely conditionally at random (MCCAR) treatment data. Through asymptotic efficiency analysis, we prove that the MAR estimator, which leverages partially-observed units, is both valid and more efficient than the MCCAR estimator when MCCAR assumptions hold. This result provides formal justification for preferring MAR-based estimation in policy learning under both missing data settings. Our comprehensive experiments using synthetic and semi-synthetic datasets confirm that correctly specifying the missingness mechanism is crucial: misspecified estimators remain biased regardless of sample size, while our estimators achieve near-oracle performance when assumptions are satisfied. Our work provides practitioners with theoretically grounded, empirically validated tools for robust policy learning in the presence of missing treatment data.

Johnna Sundberg, Rayid Ghani, Eli Ben-Michael +1

Carnegie Mellon University, Pittsburgh, Pennsylvania

Causal Effect Estimation Conditional Average Treatment Effect Estimation Learning with Missing Data Policy Optimization

Ask PDF

Abstract

Policy Learning

Explore similar work

Wasserstein Policy Learning for Distributional Outcomes

Distributionally Robust Average-Reward Reinforcement Learning: Finite-Sample Guarantees under Weak Communication

Learning Who to Treat When Treatment is Missing