Distributionally-Robust Policy Learning with Observational Data
cs.LG
Submitted: 2026-10-18
Updated: 2026-10-18
Comments: To appear at NeurIPS 2026. Generative AI was used only to polish writing and implement the main algorithm
License: http://creativecommons.org/licenses/by/4.0/
The gist: 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
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, χ squared, and CVaR ambiguity sets. We establish two convergence guarantees. First, the estimated policy attains the standard (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 (n-(1+α)/(2+α)), interpolating between (1/sqrt n) and (1/n) as the margin parameter α ranges over [0, infinity). 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.
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks