Wasserstein Policy Learning for Distributional Outcomes

Extracted main text — title through conclusion, appendix excluded. This is what our citation measures are computed over, published so the extraction can be checked by eye.

46,821 characters · 12 sections · 24 citation commands

Rendered from LaTeX for readability, not typeset faithfully. Citation keys are highlighted; maths is left as source; figures, tables and equation environments are summarised rather than reproduced; unrecognised commands are greyed out so nothing is silently dropped. Email addresses are removed.

\footnotetext[1]{Authors are in alphabetical order.} \footnotetext[2]{Corresponding author.}

abstractOffline 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}(\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 $\mathrm{N\text{-}dim}(\Pi)$.

keywordsCausal inference, policy learning, distributional outcome

Introduction

Offline policy learning, aiming to derive individualized treatment rules from observational data to maximize population-level welfare, is an important approach for personalized decision-making in causal inference zhao2012estimating, swaminathan2015batch, zhou2017residual, kitagawa2018should, kallus2018balanced, kallus2021minimax, athey2021policy. In the classical regime, a policy $\pi: \mathcal{X} \to \mathcal{A}$ is evaluated based on the expectation of a scalar potential outcome $Y \in \mathbb{R}$ (e.g., the Average Treatment Effect). This paradigm has been extensively studied, with empirical welfare maximization (EWM) approaches achieving minimax optimal rates by leveraging plug-in estimators such as Inverse Probability Weighting (IPW) kitagawa2018should or Doubly Robust (DR) athey2021policy methods to construct unbiased surrogates of the policy value.

However, reducing welfare to a scalar expectation fails to capture distributional nuances such as risk, inequality, or tail behavior. This deficiency has motivated the study of distributional policy learning, where the objective targets functionals of the outcome distribution (e.g., quantiles or CVaR) wang2018quantile, lin2023causal, cui2025policy. Crucially, while these methods optimize distributional criteria, the underlying potential outcome $Y$ remains a scalar random variable. The complexity arises solely from the non-linearity of the objective function, not from the structure of the outcome space itself.

In contrast, many modern applications feature outcomes that are inherently stochastic processes or probability measures. For instance, policymakers may aim to optimize the entire wealth distribution shape to mitigate inequality RePEc:fip:fedcwq:191800, or healthcare systems may target physiological dynamics modeled as measures on path space zhou2025dynamic. In these settings, the outcome $\mathcal{Y}$ is naturally an element of the Wasserstein space of probability measures, $\mathcal P_2(\mathbb{R})$ equipped with $\mathcal W_2$. Treating such measure-valued outcomes as densities in a linear space (e.g., $L_2$) and applying standard functional averages can destroy the intrinsic geometry of the data. As noted in panaretos2019statistical, linear averaging fails to account for mass displacement, often yielding “barycenters" that do not represent any individual realization in the population (e.g., the linear average of two unimodal distributions may be bimodal).

This necessitates a shift to policy learning with distributional outcomes. Our goal is to learn a policy that maximizes a utility defined on the Wasserstein barycenter, which is the Fréchet mean of the induced outcome distributions. This introduces a fundamental theoretical hurdle: the Wasserstein barycenter is defined as the solution to an optimization problem over a metric space, lacking the closed-form linearity of expectations kurisu2024geodesic. Consequently, the learning problem for distributional outcome cannot be trivially reduced to standard scalar-outcome empirical welfare maximization like kitagawa2018should, athey2021policy. The policy in this case induces an entire quantile curve in $L_2([0,1])$, and our objective relies on the Wasserstein barycenter, which preserves the geometry of optimal transport and differs fundamentally from the linear averaging of densities or cumulative distribution functions. Statistically, addressing these issues requires establishing uniform convergence guarantees for policy-dependent objects in an infinite-dimensional space, which necessitates controlling the complex interplay between the combinatorial complexity of the policy class and the metric entropy of the outcome space.

To see why this poses a technical barrier, consider a concrete scenario in precision medicine: using continuous glucose monitoring data to learn an insulin dosing policy. The outcome is not a single value, but a distributional profile of glucose levels over time (a probability measure). This setting involves the following challenges:

itemize• The functional complexity: The glucose profile is an infinite-dimensional object with complex shape features---it may have multiple peaks (post-prandial spikes) or heavy tails (hypoglycemia risks). Capturing this full distributional shape requires covering the massive functional space of potential quantile curves, indexed by $t \in [0,1]$. • The policy complexity: Simultaneously, the policy searches through a vast combinatorial space of decision rules (e.g., deep decision trees based on genomic data) to find the optimal subgroup assignment. • The complex interplay: The true technical barrier arises from the exploitation of functional flexibility by the combinatorial search. A policy learner, in its pursuit of empirical welfare, may “cheat” by overfitting to statistical noise at specific quantile levels—for instance, artificially overfitting the median quantile level ($t=0.5$) to maximize utility while unknowingly destabilizing the tails ($t \to 0$ or $1$), potentially leading to clinically undesirable tail behavior.

This example demonstrates that while policy learning with distributional outcomes is essential, it presents unique theoretical hurdles. Unlike scalar policy learning, where the risk of misestimation is limited to a single expected value, distribution-valued outcomes introduce policy-indexed objects over an infinite-dimensional quantile domain. The main statistical question is whether replacing scalar outcomes by distribution-valued outcomes introduces an additional nonparametric price in policy learning. Our analysis shows that, in the one-dimensional Wasserstein setting, after the quantile-isometry reduction, the leading regret remains governed by the policy-class complexity. Technically, this requires a product-index uniform deviation bound over the policy class $\Pi$ and the quantile index $t\in[0,1]$.

In this paper, we address these challenges by developing a rigorous framework for offline policy learning with distribution-valued outcomes. We formulate the problem by combining a policy-induced $W_2$-barycenter target with a Wasserstein-Lipschitz utility functional. Our key methodological insight leverages the quantile isometry between $(\mathcal P_2(\mathbb{R}),\mathcal W_2)$ and the $L_2$ space of quantile functions. This transformation maps the non-linear barycenter problem into a tractable estimation of policy-indexed quantile curves without sacrificing geometric fidelity. To solve this, we construct Inverse Propensity Weighting (IPW) and cross-fitted Doubly Robust (DR) estimators. Theoretically, the central question is to identify the statistical price of moving from scalar outcomes to distribution-valued outcomes. \paragraph{Organization.} Our analysis proceeds in four steps. First, we define a population-level policy value through the Wasserstein barycenter of policy-induced distribution-valued outcomes. Second, in one dimension, we use the quantile isometry to convert barycenter learning into estimation of a policy-indexed mean quantile curve. Third, we construct IPW and cross-fitted DR estimators and enforce validity by monotone rearrangement. Fourth, we prove finite-sample regret upper bounds and a minimax lower bound matching the leading dependence on the sample size and policy-class complexity.

\paragraph{Contributions.} Our contributions are summarized as follows:

itemize• We formalize the offline policy learning problem where outcomes are probability measures in $\mathcal P_2(\mathbb R)$. By defining the policy value via the Wasserstein barycenter, we generalize standard welfare maximization to respect the intrinsic geometry of distributional data, distinct from functional definitions in Hilbert spaces. • Leveraging the 2-Wasserstein quantile isometry, we reduce barycenter estimation to learning policy-indexed quantile curves. We propose IPW and cross-fitted DR estimators that operate directly on the space of quantile functions. Crucially, we incorporate a monotone rearrangement step that enforces the validity of the estimated quantiles without worsening the $L_2$ estimation error. • We establish finite-sample regret bounds that identify whether distribution-valued outcomes create an additional leading-order price for policy learning. By controlling the product-index uniform deviation over $\Pi\times[0,1]$, we prove a $\widetilde{\mathcal{O}}(\sqrt{\mathrm{N\text{-}dim}(\Pi)/N})$ regret rate. Thus, in the one-dimensional Wasserstein setting, the leading statistical complexity is still governed by the policy class. • We derive minimax lower bounds showing that the leading dependence on sample size $N$ and policy complexity $\mathrm{N\text{-}dim}(\Pi)$ is rate-sharp. Hence the leading rate is sharp, and the quantile-isometry reduction does not hide an additional leading-order nonparametric penalty.

Related Works

\paragraph{Policy learning in causal inference.} A central theoretical goal for offline policy learning is the establishment of minimax regret bounds that scale with the complexity of the policy class. Empirical Welfare Maximization (EWM), relying on IPW or DR estimators, has been the dominant paradigm for scalar outcomes manski2004statistical, kitagawa2018should, athey2021policy. For instance, kitagawa2018should derived $\mathcal{O}(N^{-1/2})$ regret bounds dependent on the Vapnik-Chervonenkis (VC) dimension of the policy class, while athey2021policy extended these guarantees to observational settings using cross-fitting and orthogonal scores. This literature has expanded to cover continuous treatments kallus2018policy, chernozhukov2019semi, ai2026data, policy learning without overlap khan2023off, zhao2024positivity, jin2025policy, unmeasured confounding kallus2018confounding, kallus2021minimax, distribution shift mo2021learning, kallus2022doubly, mu2022factored, adjaho2022externally, kido2022distributionally, si2023distributionally, shen2024wasserstein, and adaptive settings bibaut2021risk, zhan2024policy. However, the theoretical machinery in these works critically relies on the outcome residing in a low-dimensional Euclidean space, leaving the regime of distribution-valued outcomes unexplored.

\paragraph{Policy learning with distributional objectives.} A growing body of work has moved beyond average welfare to optimize distributional functionals, such as quantiles, CVaR, or Gini coefficients wang2018quantile, kock2024regularizing, cui2025policy, manski2023statistical. While these methods capture risk and inequality, they represent distributional objectives on scalar outcomes, not distributional outcomes. In these settings, the potential outcome $Y$ remains a scalar random variable, and the challenge lies in the non-linearity of the utility function. Consequently, these approaches typically rely on sorting or rank-based statistics of scalar variables. In contrast, our work addresses outcomes that are themselves probability measures in the 2-Wasserstein space $(\mathcal P_2(\mathbb R),\mathcal W_2)$. This shifts the problem from scalar ranking to metric space optimization. Unlike cui2025policy which targets functionals of a scalar $Y$, we integrate Wasserstein barycenters into offline policy learning, necessitating tools that respect the mass-displacement geometry of the outcome itself, and more importantly, derive finite-sample regret bounds for empirical utility maximization, together with minimax lower bounds matching the leading dependence on the sample size and policy-class complexity.

\paragraph{Causal inference with non-scalar outcomes.} Analysis of complex outcomes broadly falls into two paradigms: Hilbert-space approaches and metric-space approaches. The former, including Functional ATE (FATE) methods ecker2024causal, testa2025doubly, treat outcomes as elements of $L_2$ and apply linear averaging. While computationally tractable, linear averaging is inadequate for distributions with phase variation, often yielding unrepresentative barycenters kurisu2024geodesic. The latter paradigm, Geodesic ATE (GATE), addresses this by estimating treatment effects via Fréchet means in metric spaces like Wasserstein space lin2023causal, kurisu2024geodesic, bhattacharjee2025doubly, raykov2025kernel. Crucially, our work differs from the GATE literature in its goal and theoretical scope. Existing GATE works focus on estimation and inference to establish the consistency or asymptotic normality of a treatment effect estimator. In contrast, we address the policy learning problem, which requires selecting the optimal policy from a policy class $\Pi$. This transition shifts the theoretical challenge to deriving a regret bound with uniform convergence rates over the combinatorial complexity of $\Pi$ (e.g., Natarajan dimension). To our knowledge, we provide the first finite-sample regret bounds for decision-making with distributional outcomes in Wasserstein space.

Policy Learning with Distributional Outcome

Notations and Assumptions

Let $\mathcal A=\{a_1,\dots,a_d\}$ be a finite action set and $\mathcal X\subseteq\mathbb R^{K}$ be a compact context space. The outcome is distribution-valued: $\mathcal Y$ is a random probability measure supported on a compact interval $\mathcal I=[\mathcal I_{\mathrm{low}},\mathcal I_{\mathrm{up}}]\subset\mathbb R$, with finite second moment. We write $\mathcal Y\in\mathcal P_2(\mathbb R)$ and assume $\mathrm{supp}(\mathcal Y)\subseteq \mathcal I$ almost surely. Equivalently, $\mathcal Y$ is represented by its (left-continuous) quantile function $\mathcal Y^{-1}:[0,1]\to\mathcal I$. We observe i.i.d.\ logged data $\{(X_i,A_i,\mathcal Y_i)\}_{i=1}^N$, where $A_i$ is drawn from a behavior policy with propensity $f_0(a|x):=\mathbb P(A=a\mid X=x)$. We adopt the standard potential outcome notation: for each $a\in\mathcal A$ there exists a potential distribution $\mathcal Y[a]\in\mathcal P_2(\mathbb R)$, with quantile function $\mathcal Y[a]^{-1}$, and the observed outcome satisfies $\mathcal Y=\mathcal Y[A]$ almost surely. We state necessary assumptions as follows.

\paragraph{Assumptions.} We impose standard causal identification assumptions, adapted to the policy learning framework with distribution-valued outcomes.

assumption[Consistency] If $A=a$, then $\mathcal Y=\mathcal Y[a]$ almost surely.

assumption[Unconfoundedness] For every $a\in\mathcal A$, $\mathcal Y[a]\perp\!\!\!\perp A \mid X$.

assumption[Overlap] There exists $\underline f>0$ such that for all $x\in\mathcal X$ and $a\in\mathcal A$, $f_0(a|x)\ge \underline f$.

assumption[Boundedness] There exist continuous functions $\underline q,\overline q:[0,1]\to\mathbb R$ such that, for all $a\in\mathcal A$ and $t\in[0,1]$, almost surely, \[ \underline q(t)\le \mathcal Y[a]^{-1}(t)\le \overline q(t). \] By consistency, the same bounds hold for the observed quantile $\mathcal Y^{-1}$. We denote a uniform bound by $M:=\sup_{t\in[0,1]} \max\{|\underline q(t)|,|\overline q(t)|\}<\infty$.

assumption[Quantile regularity in $t$] There exists $\mathcal U<\infty$ such that for all $a\in\mathcal A$ and all $s,t\in[0,1]$, \[ \big|\mathcal Y[a]^{-1}(t)-\mathcal Y[a]^{-1}(s)\big|\le \mathcal U |t-s| \qquad \text{almost surely.} \] By consistency, the same holds for the observed quantile $\mathcal Y^{-1}$.

Assumptions (ref)--(ref) are standard for off-policy evaluation and policy learning. Assumptions (ref)--(ref) control the range and the $t$-regularity of the quantile curves. They enable a discretization argument that yields uniform control over the continuum index $t\in[0,1]$ in our theoretical analysis.

The Wasserstein Barycenter Objective

Let $\mathcal P_2(\mathbb R)$ denote the space of probability measures on $\mathbb R$ with finite second moments. Let $\Pi$ be a class of deterministic policies $\pi:\mathcal X\to\mathcal A$. For a given policy $\pi$, the counterfactual distributional outcome is a random measure $\mathcal Y[\pi(X)]$ taking values in $\mathcal P_2(\mathbb R)$. We define the policy-induced target as the $W_2$-barycenter of these counterfactual measures:

equation[equation omitted — 194 chars of source]

This target is the Fr\'echet mean in $(\mathcal P_2(\mathbb R),\mathcal{W}_2)$: it summarizes the policy’s effect in a way that respects the geometry of distributions under optimal transport.

The barycenter $\mu(\pi)$ provides a geometry-aware notion of the “average” distribution induced by $\pi$. To encode preferences over distributions, we optimize a distributional utility $U:\mathcal P_2(\mathbb R)\to\mathbb R$ as welfare for distribution-valued outcome and define \[ \pi^\star\in\underset{\pi\in\Pi}{\arg\max}\;U\big(\mu(\pi)\big). \] This aggregate criterion should be distinguished from the average individually scalarized objective $\mathbb E[U(\mathcal Y[\pi(X)])]$. The latter first applies $U$ to each unit-level counterfactual distribution and then averages the resulting scalars, whereas $U(\mu(\pi))$ first aggregates the policy-induced distribution-valued outcomes through their Wasserstein barycenter and then evaluates the resulting population-level distributional profile. These two criteria generally differ when $U$ is nonlinear. For example, if half of the population has the degenerate distribution $\delta_0$ and the other half has $\delta_2$, their one-dimensional $W_2$-barycenter is $\delta_1$. For the bounded-support Wasserstein-Lipschitz utility $U(\nu)=-(\int z\,d\nu(z)-1)^2$, one has $U(\delta_1)=0$ but $\frac12U(\delta_0)+\frac12U(\delta_2)=-1$. Thus our objective evaluates the aggregate distributional profile induced by a policy. Therefore, we can measure the policy performance by the regret \[ \mathcal R(\pi):=U(\mu(\pi^\star))-U(\mu(\pi)). \]

In one dimension, Wasserstein geometry admits a convenient representation through quantile functions. For a measure $\nu\in\mathcal P_2(\mathbb R)$, let $\nu^{-1}:[0,1]\to\mathbb R$ denote its (left-continuous) quantile function. The following proposition highlights that the $\mathcal{W}_2$ metric is isometric to the $L_2([0,1])$ distance between quantile functions.

proposition[Quantile representation of $\mathcal{W}_{p}$ santambrogio2015optimal] We define $\mu_{1},\mu_{2}$ as two probability measures on $\mathbb R$ with quantile functions $\mu_1^{-1},\mu_2^{-1}$. Then for $p\ge 1$, \[ \mathcal{W}_{p}(\mu_{1},\mu_{2}) = \left(\int_{0}^{1}\left|\mu_{1}^{-1}(t)-\mu_{2}^{-1}(t)\right|^{p}\,dt\right)^{\frac{1}{p}}. \]

A direct consequence of the isometry in Proposition (ref) is that the Fr\'echet mean in the Wasserstein space corresponds to the standard Euclidean mean in the quantile space. This yields an explicit characterization of the objective in (ref), as demonstrated in the following Proposition.

proposition[Characterization of the Barycenter] Under Assumption (ref), the barycenter $\mu(\pi)$ is unique and satisfies \[ \mu(\pi)^{-1}(t)=\mathbb E\big[\mathcal Y[\pi(X)]^{-1}(t)\big],\qquad \forall t\in[0,1]. \] The proof is given in Appendix (ref).

Proposition (ref) shows that in the one-dimensional case, the barycenter geometry reduces to a mean quantile curve. Importantly, this does not collapse the learning problem to scalar policy learning: the policy-induced object is the entire function \[ q_\pi(t):=\mu(\pi)^{-1}(t)=\mathbb E[\mathcal Y[\pi(X)]^{-1}(t)] \in L_2([0,1]), \] and our regret analysis requires uniform learning over the product class $\Pi\times[0,1]$. Moreover, we assume the utility function $U$ is $L_U$-Lipschitz with respect to $\mathcal W_2$:

assumption[Lipschitz utility] There exists $0<L_U<\infty$ such that for all $\nu_1,\nu_2\in\mathcal P_2(\mathbb R)$, \[ |U(\nu_1)-U(\nu_2)|\le L_U\,\mathcal{W}_2(\nu_1,\nu_2). \]

This includes several common distributional objectives and ensures stability: small Wasserstein errors in $\mu(\pi)$ translate into small utility errors. By Assumption (ref) and Proposition (ref), we have

equation[equation omitted — 188 chars of source]

Policy-class Complexity: Natarajan Dimension

Since $\Pi$ is multi-class, we quantify its complexity via the Natarajan dimension.

definition[Natarajan dimension] A set $\{x_1,\dots,x_r\}\subseteq\mathcal X$ is Natarajan-shattered by $\Pi$ if there exist $f_1,f_2:\{x_1,\dots,x_r\}\to\mathcal A$ such that (i) $f_1(x_j)\neq f_2(x_j)$ for all $j$ and (ii) for every $S\subseteq[r]$ there exists $\pi\in\Pi$ with $\pi(x_j)=f_1(x_j)$ for $j\in S$ and $\pi(x_j)=f_2(x_j)$ for $j\notin S$. The Natarajan dimension $V:=\mathrm{N\text{-}dim}(\Pi)$ is the maximum $r$ such that some set of size $r$ is shattered.

Main Results

Identification. We study statistical guarantees for learning $\pi$ by maximizing an empirical utility $U(\hat\mu(\pi))$. By Proposition (ref), the population target is the mean quantile curve $q_\pi(t)=\mu(\pi)^{-1}(t)=\mathbb E[\mathcal Y[\pi(X)]^{-1}(t)]$. Under Assumptions (ref)--(ref), $q_\pi(t)$ admits standard identification formulas. Define the outcome regression model \[ m_0(a,x)(t):=\mathbb E[\mathcal Y^{-1}(t)\mid A=a,X=x], \] which equals $\mathbb E[\mathcal Y[a]^{-1}(t)\mid X=x]$ by unconfoundedness. Then for each fixed $(\pi,t)$, \[ q_\pi(t)=\mathbb E\!\left[\frac{\mathbf 1_{\{A=\pi(X)\}}\,\mathcal Y^{-1}(t)}{f_0(A|X)}\right] =\mathbb E\big[m_0(\pi(X),X)(t)\big]. \] This identification is useful for constructing IPW and DR estimators. \paragraph{Statistical barrier: bound over $\Pi \times [0,1]$.} The central difficulty lies in establishing concentration inequalities that hold simultaneously over both the complex policy class $\pi \in \Pi$ and the functional index $t \in [0,1]$. Unlike scalar policy learning where the target is a single value, here the index set is the product of a multi-class hypothesis space and a continuum domain.

\paragraph{Monotone rearrangement for $\hat\mu(\pi)$.} In practice, a raw estimator $\hat{q}_\pi$ (e.g., via IPW or DR) is not guaranteed to be non-decreasing, and thus may fail to be a valid quantile function. To ensure that the induced measure $\hat{\mu}(\pi)$ is well-defined, we employ the standard monotone rearrangement operator. Specifically, let $\mathcal{Q}$ denote the set of left-continuous, non-decreasing functions mapping $[0, 1]$ to $\mathbb{R}$. We define the calibrated quantile function as: $\hat{q}_\pi^\uparrow := \Pi_{\mathcal{Q}}(\hat{q}_\pi)$, where $\Pi_{\mathcal{Q}}$ is the projection that finds the nearest element in $\mathcal{Q}$ with respect to a suitable norm. The estimated measure $\hat{\mu}(\pi)$ is then uniquely defined as the distribution whose quantile function is $\hat{q}_\pi^\uparrow$. Crucially, such rearrangement operators are known to be non-expansive contractions and thus do not increase the estimation error and can be treated as a standard technicality in the subsequent regret analysis.

Regret Upper Bound for IPW Formulation

We now formally define the Inverse Propensity Weighting (IPW) estimator for the distributional setting. For a fixed policy $\pi$ and any quantile level $t\in[0,1]$, Proposition (ref) establishes that the target barycenter satisfies $q_\pi(t)=\mathbb E[\mathcal Y[\pi(X)]^{-1}(t)]$. Under the assumption of unconfoundedness and strict overlap, we can identify this quantity from observational data via the inverse-propensity representation: \[ q_\pi(t)=\mathbb E\Big[\frac{\mathbf 1\{A=\pi(X)\}}{f_0(A|X)}\ \mathcal Y^{-1}(t)\Big]. \] The empirical IPW estimator for the quantile curve is defined as \[ \hat q^{\mathrm{IPW}}_\pi(t) := \mathbb P_N\!\left[\frac{\mathbf 1_{\{A=\pi(X)\}}\mathcal Y^{-1}(t)}{f_{0}(A|X)}\right]. \] While $\hat q^{\mathrm{IPW}}_\pi(t)$ is unbiased pointwise, it is not guaranteed to be monotonically increasing with respect to $t$, which violates the definition of a valid quantile function. To enforce validity, we apply monotone rearrangement, denoted by the projection operator $\Pi_{\mathcal Q}$. Let $\hat q^{\mathrm{IPW},\uparrow}_\pi := \Pi_{\mathcal Q}(\hat q^{\mathrm{IPW}}_\pi)$ be the projected estimator. We then define the estimated policy-induced barycenter $\hat\mu^{\mathrm{IPW}}(\pi)$ as the measure whose quantile function corresponds to this valid curve, i.e., $(\hat\mu^{\mathrm{IPW}}(\pi))^{-1} = \hat q^{\mathrm{IPW},\uparrow}_\pi$. Finally, the optimal policy is estimated by maximizing the empirical utility over the policy class $\Pi$: \[ \hat\pi^{\mathrm{IPW}}\in\underset{\pi\in\Pi}{\arg\max}\;U\big(\hat{\mu}^{\mathrm{IPW}}(\pi)\big). \]

theoremFix a confidence level $\delta\in(0,1)$ and a uniform grid $\mathcal T_{\mathrm{par}}=\big\{0=t_{0}<t_{1}<\cdots<t_{\mathcal{J}}=1:t_{j}=\frac{j}{\mathcal{J}},\;0\leq j\leq\mathcal{J}\big\}$ with mesh size $\eta:=1/\mathcal J$. Suppose Assumptions (ref)--(ref) hold, and let the policy class complexity be bounded by $N\ge \mathrm{N\text{-}dim}(\Pi) \ge 1$. Furthermore, let $d=|\mathcal A|$ denote the number of actions, $\underline{f}$ be the overlap lower bound, $M$ be the uniform bound on the outcome quantiles, $\mathcal{U}$ be the Lipschitz constant of the quantile curves with respect to $t$, and $L_U$ be the Lipschitz constant of the utility functional $U$. Then, with probability at least $1-\delta$, the regret of the IPW learned policy satisfies: \[ \mathcal{R}(\hat{\pi}^{\mathrm{IPW}}) \le \frac{2L_{U}M}{\underline f}\left( \sqrt{\frac{2\mathrm{N\text{-}dim}(\Pi)\log(e\cdot N\cdot d)}{N}} + \sqrt{\frac{2\log\left(\frac{2(\mathcal{J}+1)}{\delta}\right)}{N}} \right) + \left(\frac{4L_U \cdot \eta \cdot \mathcal U}{\underline f}\right). \]

\paragraph{Proof sketch.} The proof is detailed in Appendix (ref). The proof idea proceeds in two main steps.

(1) Reduction to uniform deviation via geometry. We first relate the regret $\mathcal{R}(\hat{\pi}^{\mathrm{IPW}})$ to the estimation error of the barycenter. Using the Lipschitz property of the utility $U$ in Assumption (ref) and the argmax definition of the estimator, we bound the regret by the worst-case Wasserstein distance $\sup_{\pi}\mathcal{W}_2(\hat\mu^{\mathrm{IPW}}(\pi),\mu(\pi))$. Crucially, utilizing the Wasserstein-quantile isometry and the non-expansive property of the monotone projection operator, we reduce this geometric error to the uniform deviation of the raw quantile curves in the $L_\infty$ norm: $\sup_{\pi}\sup_t|\hat q_\pi^{\mathrm{IPW}}(t)-q_\pi(t)|$.

(2) Combinatorial complexity. Since the index $t$ is continuous, a direct union bound is infeasible. We employ a covering argument by discretizing $t$ onto a grid $\mathcal T_{\mathrm{par}}$. By the Lipschitz continuity of the outcome quantiles (Assumption (ref)), we control the approximation error between grid points. On the grid, the complexity of the function class is driven by the policy search. We control this complexity using the Natarajan dimension of $\Pi$, applying a multi-class Sauer's lemma to bound the growth function and Hoeffding’s inequality to bound the point-wise deviations.

Finally, we combine the statistical estimation error (governed by the sample size $N$ and policy complexity $\mathrm{N\text{-}dim}(\Pi)$) with the deterministic discretization error (governed by the grid mesh $\eta$). The resulting bound captures the trade-off between grid resolution and statistical variance.

\paragraph{Interpretation.} The upper bound in Theorem (ref) reveals the statistical nature of learning with distributional outcomes. The first term represents the stochastic estimation error, which scales as $\widetilde{\mathcal O}(N^{-1/2})$ up to logarithmic factors. This term depends on the difficulty of the policy search (measured by $\mathrm{N\text{-}dim}(\Pi)$), the overlap ($\underline{f}^{-1}$), and the number of grid points ($\log \mathcal{J}$). Notably, the dependence on the functional dimensionality (represented by the grid size $\mathcal{J}$) is only logarithmic, indicating that the infinite-dimensional nature of the outcome does not incur a polynomial penalty in sample complexity. The second term represents the approximation bias $\mathcal{O}(\eta)$, arising from discretizing the quantile curves. This highlights a bias-variance trade-off: a finer grid (smaller $\eta$) reduces bias but increases the logarithmic stochastic error. However, by choosing $\mathcal{J} \asymp \sqrt{N}$ and given the definition $\eta = 1/\mathcal{J}$, the bias term becomes negligible ($\mathcal{O}(N^{-1/2})$), allowing the estimator to achieve the parametric rate $\widetilde{\mathcal O}(\sqrt{V/N})$ typical of scalar policy learning, without introducing an additional leading-order nonparametric penalty under the one-dimensional quantile-isometry reduction.

Regret Upper Bound for DR Formulation

While the IPW estimator is unbiased, it can be inefficient and may suffer from high variance, especially when propensity scores are close to zero. To mitigate this and reduce sensitivity to nuisance estimation errors, we adopt the cross-fitted Doubly Robust (DR) estimator chernozhukov2018double. This estimator augments the IPW objective with a regression-based objective, improving stability through Neyman orthogonality.

We employ $L$-fold cross-fitting to decouple the nuisance estimation from the policy evaluation. The sample is randomly partitioned into $L$ disjoint folds $\mathcal I_1,\dots,\mathcal I_L$ of size $n:=N/L$. For each fold $\ell$, we construct nuisance estimators $\hat f_0^\ell$ (propensity score) and $\hat m_0^\ell$ (conditional outcome quantile) using data from the complement folds $\mathcal I_{-\ell}:=\cup_{j\neq \ell}\mathcal I_j$. The cross-fitted DR estimator for the quantile curve $q_\pi(t)$ is defined as: \[ \hat q^{\mathrm{DR}}_\pi(t) := \frac{1}{L}\sum_{\ell=1}^L \mathbb P_{n,\ell}\!\left[ \hat m_0^\ell(\pi(X),X)(t) +\frac{\mathbf 1_{\{A=\pi(X)\}}}{\hat f_0^\ell(A|X)} \big(\mathcal Y^{-1}(t)-\hat m_0^\ell(\pi(X),X)(t)\big) \right], \] where $\mathbb P_{n,\ell}$ denotes the empirical average over fold $\mathcal I_\ell$. Similar to the IPW case, we apply the monotone rearrangement operator $\Pi_{\mathcal Q}$ to obtain a valid quantile curve $\hat q^{\mathrm{DR},\uparrow}_\pi := \Pi_{\mathcal Q}(\hat q^{\mathrm{DR}}_\pi)$, and define the induced barycenter $\hat\mu^{\mathrm{DR}}(\pi)$ via $(\hat\mu^{\mathrm{DR}}(\pi))^{-1} = \hat q^{\mathrm{DR},\uparrow}_\pi$. The optimal policy is then learned by maximizing the empirical utility: \[ \hat{\pi}^{\mathrm{DR}}\in\underset{\pi\in\Pi}{\arg\max}\; U(\hat{\mu}^{\mathrm{DR}}(\pi)). \] Moreover, for each fold $\ell \in \{1,\dots,L\}$, define the uniform estimation errors:

equation*[equation* omitted — 291 chars of source]

Now we provide the main theorem for the regret bound of the cross-fitted DR formulation.

theoremFix a confidence level $\delta \in (0,1/4)$ and a uniform grid $\mathcal{T}_{\mathrm{par}}$ with mesh size $\eta = 1/\mathcal{J}$. Let $d=|\mathcal A|$, $L$ be the number of folds, and $N \ge \mathrm{N\text{-}dim}(\Pi) \ge 1$. Assume the outcome is bounded by $M$ and the overlap is bounded by $\underline{f}$. Suppose Assumptions (ref)--(ref) and Assumption (ref) hold. Further, for each fold $\ell$ and every $\gamma\in(0,1)$, assume that with probability at least $1-\gamma$, \begin{equation*} \begin{aligned} \|\hat{m}_{0}^{\ell}-m_{0}\|_{\infty,[0,1]} \leq\operatorname{Rate}_{m_{0}}(N,\gamma), \qquad \|\hat{f}_{0}^{\ell}-f_{0}\|_{\infty} \leq \operatorname{Rate}_{f_{0}}(N,\gamma). \end{aligned} \end{equation*} When $\hat f_0^\ell$ appears in an inverse-propensity weight, we maintain the bounded-away-from-zero convention $\hat f_0^\ell(a|x)\ge \underline f$. Assume additionally that the fitted outcome-quantile curves are uniformly Lipschitz in $t$, so that, for all folds $\ell$, actions $a$, contexts $x$, and $s,t\in[0,1]$, \[ \left|\left(m_0(a,x)-\hat m_0^\ell(a,x)\right)(t)-\left(m_0(a,x)-\hat m_0^\ell(a,x)\right)(s)\right| \le 2\mathcal U |t-s|, \] with the fitted Lipschitz constant absorbed into $\mathcal U$. Define \[ r_f:=\operatorname{Rate}_{f_{0}}\left(N,\frac{\delta}{3L}\right), \qquad r_m:=\operatorname{Rate}_{m_{0}}\left(N,\frac{\delta}{3L}\right). \] Without loss of generality, the rate functions are taken to be nonincreasing in the confidence parameter after replacing them by their monotone envelopes. Also define \[ \mathcal{V}_N(\Pi,\delta) := \sqrt{\frac{2L\operatorname{N\text{-}dim}(\Pi)\log\left(e\frac{N}{L}d\right)}{N}} + \sqrt{\frac{2L\log\left(\frac{8\mathcal{J}L}{\delta}\right)}{N}}. \] Then, with probability at least $1-4\delta$, the regret satisfies \begin{equation} \mathcal{R}(\hat{\pi}^{\mathrm{DR}}) \leq C_{\mathrm{or}}\mathcal{V}_N(\Pi,\delta) + C_{\mathrm{grid}}\eta + \mathrm{Rem}_{\mathrm{nuis}}(N,\delta,\eta), \end{equation} where \[ C_{\mathrm{or}} := 4L_U M\left(1+\frac{2}{\underline f}\right), \qquad C_{\mathrm{grid}} := 2L_U\mathcal U\left(4+\frac{6}{\underline f}\right), \] and the nuisance-induced remainder is \begin{align*} \mathrm{Rem}_{\mathrm{nuis}}(N,\delta,\eta) :=\;& 2L_U \Bigg[ \left( \frac{2r_fr_m}{\underline f^2} + \frac{4Mr_f}{\underline f^2} + 2\left(1+\frac{1}{\underline f}\right)r_m \right) \mathcal{V}_N(\Pi,\delta) \\ &\qquad\qquad + \frac{5r_fr_m}{\underline f^2} + \frac{4\mathcal U r_f}{\underline f^2}\eta \Bigg]. \end{align*} Equivalently, there exists a constant $C>0$, depending only on $L_U,M,\underline f,\mathcal U,L$, and $d$, such that \[ \mathcal R(\hat\pi^{\mathrm{DR}}) \le C\left[ \mathcal V_N(\Pi,\delta) + \eta + r_fr_m + r_f\mathcal V_N(\Pi,\delta) + r_m\mathcal V_N(\Pi,\delta) + r_f\eta \right]. \]

\paragraph{Proof sketch.} The proof is detailed in Appendix (ref). The regret is first reduced to the uniform $L_2$ deviation of the estimated policy-indexed quantile curve: \[ \mathcal R(\hat\pi^{\mathrm{DR}}) \le 2L_U \left( \int_0^1 \left[ \sup_{\pi\in\Pi} \left| \hat q_\pi^{\mathrm{DR}}(t)-q_\pi(t) \right| \right]^2dt \right)^{1/2}. \] The DR score is then decomposed into an oracle empirical-process term, a second-order product bias term, and nuisance-dependent centered empirical-process terms. The key point is Neyman orthogonality: the first-order conditional mean terms in the propensity and outcome-regression errors vanish after conditioning on the training folds used to estimate the nuisance functions. Consequently, the standalone nuisance bias is of product order $r_fr_m$; the first-order terms $r_f$ and $r_m$ appear only multiplied by the policy-complexity factor $\mathcal V_N(\Pi,\delta)$ or the grid mesh $\eta$.

\paragraph{Rate interpretation and Neyman orthogonality.} The bound separates the oracle policy-learning error from the nuisance-induced error. Up to constants depending only on $L_U,M,\underline f,\mathcal U,L$, and $d$, Theorem (ref) implies \[ \mathcal R(\hat\pi^{\mathrm{DR}}) \lesssim \mathcal V_N(\Pi,\delta) + \eta + r_fr_m + r_f\mathcal V_N(\Pi,\delta) + r_m\mathcal V_N(\Pi,\delta) + r_f\eta. \] The key point is that Neyman orthogonality removes the standalone first-order nuisance bias. The only non-centered nuisance bias is of product order $r_fr_m$. The terms involving $r_f$ or $r_m$ alone are centered empirical-process or discretization remainders and are multiplied by $\mathcal V_N(\Pi,\delta)$ or $\eta$. Consequently, if \[ r_f=\mathcal{O}(N^{-\alpha_f}), \qquad r_m=\mathcal{O}(N^{-\alpha_m}), \qquad \mathcal V_N(\Pi,\delta)=\widetilde{\mathcal O}(N^{-1/2}), \qquad \eta=\mathcal{O}(N^{-1/2}), \] then \[ \mathcal R(\hat\pi^{\mathrm{DR}}) = \widetilde{\mathcal O}\!\left( N^{-1/2} + N^{-(\alpha_f+\alpha_m)} + N^{-(\alpha_f+1/2)} + N^{-(\alpha_m+1/2)} \right). \] In particular, if both nuisance estimators converge at the $N^{-1/4}$ rate, then \[ \mathcal R(\hat\pi^{\mathrm{DR}}) = \widetilde{\mathcal O}(N^{-1/2}). \]

Minimax Lower Bound

To complement the upper bounds derived in the previous sections, we now establish the fundamental information-theoretic limits of policy learning with distributional outcomes. Our goal is to show that the dependence on the sample size $N$ and the policy complexity $V:=\mathrm{N\text{-}dim}(\Pi)$ in our upper bounds is essentially tight.

We focus on a specific class of distributional utilities, the integrated quantile utility, defined as: \[ U_\alpha(\nu):=\int_0^\alpha \nu^{-1}(t)\,dt,\qquad \alpha\in(0,1]. \] This utility fits within our framework as it satisfies the Lipschitz property (Assumption (ref)) via the Cauchy-Schwarz inequality: $|U_\alpha(\nu_1)-U_\alpha(\nu_2)|\le \sqrt{\alpha}\,\mathcal{W}_2(\nu_1,\nu_2)$. By constructing a hard instance for this specific utility, we demonstrate the hardness of the general problem.

theorem[Minimax Lower Bound] Assume the policy class satisfies $\mathrm{N\text{-}dim}(\Pi)\ge 1$, let $d=|\mathcal A|\ge2$, fix $\alpha\in(0,1]$, and take $0<\underline f\le 1/d$. Let $q_-,q_+:[0,1]\to\mathcal I$ be two left-continuous, non-decreasing, $\mathcal U$-Lipschitz quantile curves satisfying $q_-(t)\le q_+(t)$ and $\sup_{t\in[0,1]}\max\{|q_-(t)|,|q_+(t)|\}\le M$. Define $\Delta_Q:=\int_0^\alpha (q_+(t)-q_-(t))\,dt$, and let $\mathcal{P}_{\text{lower}}(q_-,q_+)$ be the subclass of distributions satisfying Assumptions (ref)--(ref) whose potential-outcome quantile curves take values in $\{q_-,q_+\}$ and whose behavior policy has overlap at least $\underline f$. There exists a universal constant $c_0>0$ such that for any learning algorithm that maps a dataset $\mathcal{D}_N$ to a policy $\hat\pi$, the worst-case regret is bounded from below by: \[ \inf_{\hat\pi}\ \sup_{\mathbb P\in\mathcal P_{\text{lower}}(q_-,q_+)}\ \mathbb E_{\mathcal{D}_N \sim \mathbb{P}^{\otimes N}}\!\left[\mathcal{R}(\hat{\pi})\right] \ \ge\ c_0\,\Delta_Q \min\left\{1,\sqrt{\frac{\mathrm{N\text{-}dim}(\Pi)}{\underline f\,N}}\right\}, \] where $\mathcal{R}(\hat{\pi}) = U_\alpha(\mu_{\mathbb{P}}(\pi_{\mathbb{P}}^{\ast}))-U_\alpha(\mu_{\mathbb{P}}(\hat{\pi}))$ is the regret.

\paragraph{Proof sketch.} The detailed proof is provided in Appendix (ref). We employ the method of Assouad's Lemma by constructing a hard instance family parameterized by the vertices of a hypercube $\mathcal V=\{\pm1\}^{\mathrm{N\text{-}dim}(\Pi)}$. By designing the potential outcomes as distribution-valued random mixtures of two valid base quantile curves $q_-$ and $q_+$, we reduce the utility maximization problem to a multiple hypothesis testing problem of identifying the optimal action configuration. We then bound the error probability via the Kullback-Leibler (KL) divergence between distributions induced by adjacent vertices. Since the difficulty of distinguishing the optimal actions under noisy outcomes scales inversely with the sample size $N$ and overlap $\underline{f}$, balancing the regret magnitude with this detection probability yields the leading minimax dependence $\sqrt{\mathrm{N\text{-}dim}(\Pi)/(\underline f N)}$ in the large-sample regime.

\paragraph{Interpretation.} Theorem (ref) supports the rate sharpness of the proposed framework with respect to $N$ and policy-class complexity. In the large-sample regime, the lower bound scales as $\sqrt{\mathrm{N\text{-}dim}(\Pi)/(\underline f N)}$, which matches the leading $\mathcal{O}(N^{-1/2})$ dependence of our DR upper bound (Theorem (ref)) up to logarithmic factors. The bound clarifies the cost of the distributional structure: while it depends on the integrated gap $\Delta_Q$ (analogous to the scalar outcome range), it does not introduce a leading dependence on the non-parametric metric entropy of the outcome space in this one-dimensional construction.

Conclusion

We formulated offline policy learning when each potential outcome is a probability measure and policy performance is evaluated through a utility of the policy-induced Wasserstein barycenter. Focusing on one-dimensional outcomes, we leveraged the quantile representation of $W_2$ to express the barycenter as a mean quantile curve, enabling empirical welfare maximization while preserving the optimal-transport geometry. We analyzed two estimators of the mean quantile curve, an IPW estimator with known propensities and a cross-fitted doubly robust estimator with estimated nuisance functions. Our main theoretical results provide finite-sample regret guarantees with leading dependence $\widetilde{\mathcal O}(\sqrt{\mathrm{N\text{-}dim}(\Pi)/N})$ on the Natarajan dimension of the policy class and the sample size. We also established a minimax lower bound showing sharpness of the leading dependence on $(\mathrm{N\text{-}dim}(\Pi),N)$ for a representative one-dimensional construction.

Several directions remain open. This work mainly focuses on one-dimensional distributional outcomes, where the quantile isometry provides an exact representation of Wasserstein geometry and makes finite-sample policy regret analysis tractable. Extending the analysis to multivariate Wasserstein spaces is substantially more challenging, since there is no canonical quantile ordering and Wasserstein barycenters generally lack the explicit structure used in our proofs. Another important direction is to relax the quantile regularity conditions imposed here, for example by allowing weaker smoothness or tail behavior. Finally, stochastic policies, continuous treatments, and sequential decision problems involve different policy classes and complexity measures, and therefore require separate regret analyses. We leave these extensions for future work.

Acknowledgments

Yiyan Huang was supported by the Startup Funds of Great Bay University (No. YJKY250111) and the Innovative Team Program for Regular Universities in Guangdong Province (No. 2025KCXTD031). Qi Wu was supported by the CityU-JD Digits Joint Laboratory in Financial Technology and Engineering, the Hong Kong Research Grants Council General Research Fund (Nos. 11219420/9043008 and 11200219/9042900), the HK Institute of Data Science, the InnoHK initiative of the Government of the HKSAR, and the Laboratory for AI-Powered Financial Technologies. Zhiheng Zhang was supported by the Fundamental Research Funds for the Central Universities (No. 2025110602), the Independent Research Project funded by the School of Statistics and Data Science (No. 2026110081), and the Shanghai Engineering Research Center of Finance Intelligence (No. 19DZ2254600).