Wasserstein Policy Learning for Distributional Outcomes

A new study establishes statistical guarantees for offline policy learning when outcomes are entire probability distributions rather than single numbers, using a framework built on Wasserstein barycenters and quantile-function representations [1]. The paper, posted to arXiv on 17 June 2026, addresses a gap in causal inference where standard policy learning assumes scalar-valued potential outcomes [1]. In many modern applications, however, the outcome of interest is naturally a probability measure — for example, the full wealth-distribution shape that a policymaker wants to shift, or a patient’s glucose-level profile over time viewed as a distribution [3]. The authors formulate the reward through a utility functional applied to the Wasserstein barycenter of the policy-induced outcome distributions [2]. A key methodological step is the use of the quantile isometry between the Wasserstein space (𝒫₂(ℝ),𝒲₂) and the L₂ space of quantile functions. This transformation converts the non-linear barycenter problem into a tractable estimation of policy-indexed quantile curves without sacrificing geometric fidelity [2]. The authors construct both Inverse Probability Weighting (IPW) and cross-fitted Doubly Robust (DR) estimators to operationalize the approach [2]. The central theoretical question is whether moving from scalar to distribution-valued outcomes imposes an additional statistical price. By controlling the uniform deviation over the product of the combinatorial policy class Π and the infinite-dimensional quantile domain [0,1], the paper proves a finite-sample regret bound with leading dependence Õ(√(N-dim(Π)/N)) [2]. In the one-dimensional Wasserstein setting and under stated regularity conditions, the leading regret rate remains governed by the policy-class complexity [1]. A minimax lower bound is also provided, confirming that the dependence on the sample size N and the policy-class complexity N-dim(Π) is sharp [2]. While this work focuses on offline policy learning from observational data, the broader Wasserstein-policy literature includes online reinforcement-learning methods. A separate study introduced Wasserstein Policy Optimization (WPO), an actor-critic algorithm for continuous action spaces derived as an approximation to Wasserstein gradient flow over the space of policies projected into a finite-dimensional parameter space [4]. That algorithm combines properties of deterministic and classic policy-gradient methods and has been tested on the DeepMind Control Suite and a magnetic-confinement fusion task [4]. The new offline framework differs in its distributional-outcome focus and its reliance on IPW and DR estimators rather than online gradient-flow approximations. The paper’s guarantees apply to settings where the potential outcome for each treatment is a probability measure on ℝ, and the utility functional is Wasserstein-Lipschitz [2]. By reducing the barycenter problem to quantile-curve estimation, the framework avoids the geometric distortions that can arise when measure-valued outcomes are treated as densities in a linear space and averaged pointwise — a practice that can produce “barycenters” that do not represent any individual realization in the population [3].