# Wasserstein Policy Learning for Distributional Outcomes

arXiv: https://arxiv.org/abs/2606.19117 · HTML: https://mlsanity.com/paper/2606.19117/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 22 of its 41 references that are in the map.
- First posted on arXiv 17 Jun 2026.

## Builds on (most cited first, 22 of 22)

- [Statistical Aspects of Wasserstein Distances](https://mlsanity.com/paper/1806.05500/) (2018)
- [Confounding-Robust Policy Improvement](https://mlsanity.com/paper/1805.08593/) (2018)
- [Policy Evaluation and Optimization with Continuous Treatments](https://mlsanity.com/paper/1802.06037/) (2018)
- [Balanced Policy Evaluation and Learning](https://mlsanity.com/paper/1705.07384/) (2017)
- [Residual Weighted Learning for Estimating Individualized Treatment Rules](https://mlsanity.com/paper/1508.03179/) (2015)
- [Doubly Robust Distributionally Robust Off-Policy Evaluation and Learning](https://mlsanity.com/paper/2202.09667/) (2022)
- [Policy learning "without" overlap: Pessimism and generalized empirical Bernstein's inequality](https://mlsanity.com/paper/2212.09900/) (2022)
- [Learning Optimal Distributionally Robust Individualized Treatment Rules](https://mlsanity.com/paper/2006.15121/) (2020)
- [Semi-Parametric Efficient Policy Learning with Continuous Actions](https://mlsanity.com/paper/1905.10116/) (2019)
- [Policy Learning with Adaptively Collected Data](https://mlsanity.com/paper/2105.02344/) (2021)
- [Distributionally Robust Batch Contextual Bandits](https://mlsanity.com/paper/2006.05630/) (2020)
- [Causal Inference on Distribution Functions](https://mlsanity.com/paper/2101.01599/) (2021)
- [Risk Minimization from Adaptively Collected Data: Guarantees for Supervised and Policy Learning](https://mlsanity.com/paper/2106.01723/) (2021)
- [Wasserstein Distributionally Robust Policy Evaluation and Learning for Contextual Bandits](https://mlsanity.com/paper/2309.08748/) (2023)
- [Causal inference with a functional outcome](https://mlsanity.com/paper/2304.07113/) (2023)
- [Geodesic Causal Inference](https://mlsanity.com/paper/2406.19604/) (2024)
- [Positivity-free Policy Learning with Observational Data](https://mlsanity.com/paper/2310.06969/) (2023)
- [Policy Learning with Distributional Welfare](https://mlsanity.com/paper/2311.15878/) (2023)
- [Doubly robust estimation of causal effects for random object outcomes with continuous treatments](https://mlsanity.com/paper/2506.22754/) (2025)
- [Off-policy evaluation beyond overlap: partial identification through smoothness](https://mlsanity.com/paper/2305.11812/) (2023)
- [Doubly-Robust Functional Average Treatment Effect Estimation](https://mlsanity.com/paper/2501.06024/) (2025)
- [Kernel-based estimators for functional causal effects](https://mlsanity.com/paper/2503.05024/) (2025)

## Abstract (arXiv)

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 ℝ 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 widetildeO(√mathrmN-dim(Pi)/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 mathrmN-dim(Pi).