arXiv CSopen access

Wasserstein Policy Learning for Distributional Outcomes

Proceedings of Machine Learning Research vol 336:1–33, 2026

39th Annual Conference on Learning Theory

Wasserstein Policy Learning for Distributional Outcomes Yiyan Huang* †

HUANGYIYAN @ GBU . EDU . CN School of Computing and Information Technology, Great Bay University, Guangdong, China

Cheuk Hang Leung*

CHLEUNG 87@ CITYU . EDU . HK

Department of Data Science, City University of Hong Kong, Hong Kong, China

Qi Wu*

QIWU 55@ CITYU . EDU . HK

arXiv:2606.19117v1 [stat.ME] 17 Jun 2026

Department of Data Science, City University of Hong Kong, Hong Kong, China

Zhiheng Zhang*

ZHANGZHIHENG @ MAIL . SHUFE . EDU . CN

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

Editors: Steve Hanneke and Tor Lattimore

Abstract 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 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 finitep e N-dim(Π)/N ). In the one-dimensional Wasserstein sample regret has leading dependence O( 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 N-dim(Π). Keywords: Causal inference, policy learning, distributional outcome

1. 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 (Zhao et al., 2012; Swaminathan and Joachims, 2015; Zhou et al., 2017; Kitagawa and Tetenov, 2018; Kallus, 2018; Kallus and Zhou, 2021; Athey and Wager, 2021). In the classical regime, a policy π : X → A is evaluated based on the expectation of a scalar potential outcome Y ∈ 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) (Kitagawa and Tetenov, 2018) or Doubly Robust (DR) (Athey and Wager, 2021) methods to construct unbiased surrogates of the policy value. *. Authors are in alphabetical order. †. Corresponding author. © 2026 Y. Huang* † , C.H. Leung* , Q. Wu* & Z. Zhang* .

H UANG* † L EUNG* W U* Z HANG*

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) (Wang et al., 2018; Lin et al., 2023; Cui and Han, 2025). 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 (Aliprantis et al., 2022), or healthcare systems may target physiological dynamics modeled as measures on path space (Zhou and Mueller, 2025). In these settings, the outcome Y is naturally an element of the Wasserstein space of probability measures, P2 (R) equipped with W2 . Treating such measure-valued outcomes as densities in a linear space (e.g., L2 ) and applying standard functional averages can destroy the intrinsic geometry of the data. As noted in Panaretos and Zemel (2019), 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 (Kurisu et al., 2024). Consequently, the learning problem for distributional outcome cannot be trivially reduced to standard scalar-outcome empirical welfare maximization like Kitagawa and Tetenov (2018); Athey and Wager (2021). The policy in this case induces an entire quantile curve in L2 ([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 infinitedimensional 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: • 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 ∈ [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, 2

WASSERSTEIN P OLICY L EARNING FOR D ISTRIBUTIONAL O UTCOMES

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 → 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 Π and the quantile index t ∈ [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 policyinduced W2 -barycenter target with a Wasserstein-Lipschitz utility functional. Our key methodological insight leverages the quantile isometry between (P2 (R), W2 ) and the L2 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. 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. Contributions.

Our contributions are summarized as follows:

• We formalize the offline policy learning problem where outcomes are probability measures in P2 (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 L2 estimation error. • We establish finite-sample regret bounds that identify whether distribution-valued outcomes create an additional leading-order price for policy p learning. By controlling the product-index e N-dim(Π)/N ) regret rate. Thus, in the uniform deviation over Π × [0, 1], we prove a O( 3

H UANG* † L EUNG* W U* Z HANG*

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 N-dim(Π) is rate-sharp. Hence the leading rate is sharp, and the quantileisometry reduction does not hide an additional leading-order nonparametric penalty.

2. Related Works 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 (Manski, 2004; Kitagawa and Tetenov, 2018; Athey and Wager, 2021). For instance, Kitagawa and Tetenov (2018) derived O(N −1/2 ) regret bounds dependent on the VapnikChervonenkis (VC) dimension of the policy class, while Athey and Wager (2021) extended these guarantees to observational settings using cross-fitting and orthogonal scores. This literature has expanded to cover continuous treatments (Kallus and Zhou, 2018b; Chernozhukov et al., 2019; Ai et al., 2026), policy learning without overlap (Khan et al., 2023; Zhao et al., 2024; Jin et al., 2025), unmeasured confounding (Kallus and Zhou, 2018a, 2021), distribution shift (Mo et al., 2021; Kallus et al., 2022; Mu et al., 2022; Adjaho and Christensen, 2022; Kido, 2022; Si et al., 2023; Shen et al., 2024), and adaptive settings (Bibaut et al., 2021; Zhan et al., 2024). 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. 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 (Wang et al., 2018; Kock and Preinerstorfer, 2024; Cui and Han, 2025; Manski and Tetenov, 2023). 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 (P2 (R), W2 ). This shifts the problem from scalar ranking to metric space optimization. Unlike Cui and Han (2025) 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. 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 (Ecker et al., 2024; Testa et al., 2025), treat outcomes as elements of L2 and apply linear averaging. While computationally tractable, linear averaging is inadequate for distributions with phase variation, often yielding unrepresentative barycenters (Kurisu et al., 2024). The latter paradigm, Geodesic ATE (GATE), addresses this by estimating treatment effects via Fréchet means in metric spaces like Wasserstein space (Lin et al., 2023; Kurisu et al., 2024; 4

WASSERSTEIN P OLICY L EARNING FOR D ISTRIBUTIONAL O UTCOMES

Bhattacharjee et al., 2025; Raykov et al., 2025). 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 Π. This transition shifts the theoretical challenge to deriving a regret bound with uniform convergence rates over the combinatorial complexity of Π (e.g., Natarajan dimension). To our knowledge, we provide the first finite-sample regret bounds for decision-making with distributional outcomes in Wasserstein space.

3. Policy Learning with Distributional Outcome 3.1. Notations and Assumptions Let A = {a1 , . . . , ad } be a finite action set and X ⊆ RK be a compact context space. The outcome is distribution-valued: Y is a random probability measure supported on a compact interval I = [Ilow , Iup ] ⊂ R, with finite second moment. We write Y ∈ P2 (R) and assume supp(Y) ⊆ I almost surely. Equivalently, Y is represented by its (left-continuous) quantile function Y −1 : [0, 1] → I. We observe i.i.d. logged data {(Xi , Ai , Yi )}N i=1 , where Ai is drawn from a behavior policy with propensity f0 (a|x) := P(A = a | X = x). We adopt the standard potential outcome notation: for each a ∈ A there exists a potential distribution Y[a] ∈ P2 (R), with quantile function Y[a]−1 , and the observed outcome satisfies Y = Y[A] almost surely. We state necessary assumptions as follows. Assumptions. We impose standard causal identification assumptions, adapted to the policy learning framework with distribution-valued outcomes. Assumption 1 (Consistency) If A = a, then Y = Y[a] almost surely. Assumption 2 (Unconfoundedness) For every a ∈ A, Y[a] ⊥⊥ A | X. Assumption 3 (Overlap) There exists f > 0 such that for all x ∈ X and a ∈ A, f0 (a|x) ≥ f . Assumption 4 (Boundedness) There exist continuous functions q, q : [0, 1] → R such that, for all a ∈ A and t ∈ [0, 1], almost surely, q(t) ≤ Y[a]−1 (t) ≤ q(t). By consistency, the same bounds hold for the observed quantile Y −1 . We denote a uniform bound by M := supt∈[0,1] max{|q(t)|, |q(t)|} < ∞. Assumption 5 (Quantile regularity in t) There exists U < ∞ such that for all a ∈ A and all s, t ∈ [0, 1], Y[a]−1 (t) − Y[a]−1 (s) ≤ U|t − s| almost surely. By consistency, the same holds for the observed quantile Y −1 . Assumptions 1–3 are standard for off-policy evaluation and policy learning. Assumptions 4–5 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 ∈ [0, 1] in our theoretical analysis. 5

H UANG* † L EUNG* W U* Z HANG*

3.2. The Wasserstein Barycenter Objective Let P2 (R) denote the space of probability measures on R with finite second moments. Let Π be a class of deterministic policies π : X → A. For a given policy π, the counterfactual distributional outcome is a random measure Y[π(X)] taking values in P2 (R). We define the policy-induced target as the W2 -barycenter of these counterfactual measures: h
i µ(π) ∈ arg min E W22 µ, Y[π(X)] . (1) µ∈P2 (R)

This target is the Fréchet mean in (P2 (R), W2 ): it summarizes the policy’s effect in a way that respects the geometry of distributions under optimal transport. The barycenter µ(π) provides a geometry-aware notion of the “average” distribution induced by π. To encode preferences over distributions, we optimize a distributional utility U : P2 (R) → R as welfare for distribution-valued outcome and define
π ⋆ ∈ arg max U µ(π) . π∈Π

This aggregate criterion should be distinguished from the average individually scalarized objective E[U (Y[π(X)])]. The latter first applies U to each unit-level counterfactual distribution and then averages the resulting scalars, whereas U (µ(π)) first aggregates the policy-induced distributionvalued outcomes through their Wasserstein barycenter and then evaluates the resulting populationlevel distributional profile. These two criteria generally differ when U is nonlinear. For example, if half of the population has the degenerate distribution δ0 and the other half has δ2 , their onedimensional W2 -barycenter is δ1 . For the bounded-support Wasserstein-Lipschitz utility U (ν) = R −( z dν(z) − 1)2 , one has U (δ1 ) = 0 but 21 U (δ0 ) + 12 U (δ2 ) = −1. Thus our objective evaluates the aggregate distributional profile induced by a policy. Therefore, we can measure the policy performance by the regret R(π) := U (µ(π ⋆ )) − U (µ(π)). In one dimension, Wasserstein geometry admits a convenient representation through quantile functions. For a measure ν ∈ P2 (R), let ν −1 : [0, 1] → R denote its (left-continuous) quantile function. The following proposition highlights that the W2 metric is isometric to the L2 ([0, 1]) distance between quantile functions. Proposition 6 (Quantile representation of Wp (Santambrogio, 2015)) We define µ1 , µ2 as two −1 probability measures on R with quantile functions µ−1 1 , µ2 . Then for p ≥ 1, Z 1 Wp (µ1 , µ2 ) = 0

p −1 dt µ−1 1 (t) − µ2 (t)

 p1 .

A direct consequence of the isometry in Proposition 6 is that the Fréchet mean in the Wasserstein space corresponds to the standard Euclidean mean in the quantile space. This yields an explicit characterization of the objective in (1), as demonstrated in the following Proposition. Proposition 7 (Characterization of the Barycenter) Under Assumption 4, the barycenter µ(π) is unique and satisfies   µ(π)−1 (t) = E Y[π(X)]−1 (t) , ∀t ∈ [0, 1]. The proof is given in Appendix A. 6

WASSERSTEIN P OLICY L EARNING FOR D ISTRIBUTIONAL O UTCOMES

Proposition 7 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π (t) := µ(π)−1 (t) = E[Y[π(X)]−1 (t)] ∈ L2 ([0, 1]), and our regret analysis requires uniform learning over the product class Π × [0, 1]. Moreover, we assume the utility function U is LU -Lipschitz with respect to W2 : Assumption 8 (Lipschitz utility) There exists 0 < LU < ∞ such that for all ν1 , ν2 ∈ P2 (R), |U (ν1 ) − U (ν2 )| ≤ LU W2 (ν1 , ν2 ). This includes several common distributional objectives and ensures stability: small Wasserstein errors in µ(π) translate into small utility errors. By Assumption 8 and Proposition 7, we have Z 1

⋆

R(π) ≤ LU W2 (µ(π ), µ(π)) = LU

2

1/2

|q (t) − qπ (t)| dt π⋆

.

(2)

0

3.3. Policy-class Complexity: Natarajan Dimension Since Π is multi-class, we quantify its complexity via the Natarajan dimension. Definition 9 (Natarajan dimension) A set {x1 , . . . , xr } ⊆ X is Natarajan-shattered by Π if there exist f1 , f2 : {x1 , . . . , xr } → A such that (i) f1 (xj ) ̸= f2 (xj ) for all j and (ii) for every S ⊆ [r] there exists π ∈ Π with π(xj ) = f1 (xj ) for j ∈ S and π(xj ) = f2 (xj ) for j ∈ / S. The Natarajan dimension V := N-dim(Π) is the maximum r such that some set of size r is shattered.

4. Main Results Identification. We study statistical guarantees for learning π by maximizing an empirical utility U (µ̂(π)). By Proposition 7, the population target is the mean quantile curve qπ (t) = µ(π)−1 (t) = E[Y[π(X)]−1 (t)]. Under Assumptions 1–3, qπ (t) admits standard identification formulas. Define the outcome regression model m0 (a, x)(t) := E[Y −1 (t) | A = a, X = x], which equals E[Y[a]−1 (t) | X = x] by unconfoundedness. Then for each fixed (π, t), " #   1{A=π(X)} Y −1 (t) qπ (t) = E = E m0 (π(X), X)(t) . f0 (A|X) This identification is useful for constructing IPW and DR estimators. Statistical barrier: bound over Π×[0, 1]. The central difficulty lies in establishing concentration inequalities that hold simultaneously over both the complex policy class π ∈ Π and the functional index t ∈ [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. 7

H UANG* † L EUNG* W U* Z HANG*

Monotone rearrangement for µ̂(π). In practice, a raw estimator q̂π (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 µ̂(π) is well-defined, we employ the standard monotone rearrangement operator. Specifically, let Q denote the set of left-continuous, non-decreasing functions mapping [0, 1] to R. We define the calibrated quantile function as: q̂π↑ := ΠQ (q̂π ), where ΠQ is the projection that finds the nearest element in Q with respect to a suitable norm. The estimated measure µ̂(π) is then uniquely defined as the distribution whose quantile function is q̂π↑ . 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. 4.1. Regret Upper Bound for IPW Formulation We now formally define the Inverse Propensity Weighting (IPW) estimator for the distributional setting. For a fixed policy π and any quantile level t ∈ [0, 1], Proposition 7 establishes that the target barycenter satisfies qπ (t) = E[Y[π(X)]−1 (t)]. Under the assumption of unconfoundedness and strict overlap, we can identify this quantity from observational data via the inverse-propensity representation: h 1{A = π(X)} i qπ (t) = E Y −1 (t) . f0 (A|X) The empirical IPW estimator for the quantile curve is defined as " # −1 (t) 1 Y {A=π(X)} q̂πIPW (t) := PN . f0 (A|X) While q̂πIPW (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 ΠQ . Let q̂πIPW,↑ := ΠQ (q̂πIPW ) be the projected estimator. We then define the estimated policy-induced barycenter µ̂IPW (π) as the measure whose quantile function corresponds to this valid curve, i.e., (µ̂IPW (π))−1 = q̂πIPW,↑ . Finally, the optimal policy is estimated by maximizing the empirical utility over the policy class Π:
π̂ IPW ∈ arg max U µ̂IPW (π) . π∈Π

 Theorem 10 Fix a confidence level δ ∈ (0, 1) and a uniform grid Tpar = 0 = t0 < t1 < · · · < tJ = 1 : tj = Jj , 0 ≤ j ≤ J with mesh size η := 1/J . Suppose Assumptions 1–5 hold, and let the policy class complexity be bounded by N ≥ N-dim(Π) ≥ 1. Furthermore, let d = |A| denote the number of actions, f be the overlap lower bound, M be the uniform bound on the outcome quantiles, U be the Lipschitz constant of the quantile curves with respect to t, and LU be the Lipschitz constant of the utility functional U . Then, with probability at least 1 − δ, the regret of the IPW learned policy satisfies:  v   u r u 2 log 2(J +1)    δ 2LU M  2N-dim(Π) log(e · N · d) t 4LU · η · U IPW   R(π̂ )≤ + .  + f N N f

8

WASSERSTEIN P OLICY L EARNING FOR D ISTRIBUTIONAL O UTCOMES

Proof sketch. The proof is detailed in Appendix B. The proof idea proceeds in two main steps. (1) Reduction to uniform deviation via geometry. We first relate the regret R(π̂ IPW ) to the estimation error of the barycenter. Using the Lipschitz property of the utility U in Assumption 8 and the argmax definition of the estimator, we bound the regret by the worst-case Wasserstein distance supπ W2 (µ̂IPW (π), µ(π)). 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∞ norm: supπ supt |q̂πIPW (t) − qπ (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 Tpar . By the Lipschitz continuity of the outcome quantiles (Assumption 5), 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 Π, 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 N-dim(Π)) with the deterministic discretization error (governed by the grid mesh η). The resulting bound captures the trade-off between grid resolution and statistical variance. Interpretation. The upper bound in Theorem 10 reveals the statistical nature of learning with distributional outcomes. The first term represents the stochastic estimation error, which scales e −1/2 ) up to logarithmic factors. This term depends on the difficulty of the policy search as O(N (measured by N-dim(Π)), the overlap (f −1 ), and the number of grid points (log J ). Notably, the dependence on the functional dimensionality (represented by the grid size 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 O(η), arising from discretizing the quantile curves. This highlights a bias-variance trade-off: a finer grid (smaller η) √ reduces bias but increases the logarithmic stochastic error. However, by choosing J ≍ N and given the definition η = 1/J , the p bias term becomes negligible (O(N −1/2 )), allowing the estimator e to achieve the parametric rate O( V /N ) typical of scalar policy learning, without introducing an additional leading-order nonparametric penalty under the one-dimensional quantile-isometry reduction. 4.2. 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 (Chernozhukov et al., 2018). 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 I1 , . . . , IL of size n := N/L. For each fold ℓ, we construct nuisance estimators fˆ0ℓ (propensity score) and m̂ℓ0 (conditional outcome quantile) using data from the complement folds I−ℓ := ∪j̸=ℓ Ij . The cross-fitted DR estimator for the quantile curve qπ (t) is defined as: " # L X
1 1 {A=π(X)} Pn,ℓ m̂ℓ0 (π(X), X)(t) + q̂πDR (t) := Y −1 (t) − m̂ℓ0 (π(X), X)(t) , L fˆ0ℓ (A|X) ℓ=1

9

H UANG* † L EUNG* W U* Z HANG*

where Pn,ℓ denotes the empirical average over fold Iℓ . Similar to the IPW case, we apply the monotone rearrangement operator ΠQ to obtain a valid quantile curve q̂πDR,↑ := ΠQ (q̂πDR ), and define the induced barycenter µ̂DR (π) via (µ̂DR (π))−1 = q̂πDR,↑ . The optimal policy is then learned by maximizing the empirical utility: π̂ DR ∈ arg max U (µ̂DR (π)). π∈Π

Moreover, for each fold ℓ ∈ {1, . . . , L}, define the uniform estimation errors: ∥fˆ0ℓ − f0 ∥∞ :=

sup x∈X ,a∈A

fˆ0ℓ (a|x) − f0 (a|x) ,

∥m̂ℓ0 − m0 ∥∞,[0,1] := sup m̂ℓ0 (a, x)(t) − m0 (a, x)(t) . x,a t∈[0,1]

Now we provide the main theorem for the regret bound of the cross-fitted DR formulation. Theorem 11 Fix a confidence level δ ∈ (0, 1/4) and a uniform grid Tpar with mesh size η = 1/J . Let d = |A|, L be the number of folds, and N ≥ N-dim(Π) ≥ 1. Assume the outcome is bounded by M and the overlap is bounded by f . Suppose Assumptions 1–5 and Assumption 8 hold. Further, for each fold ℓ and every γ ∈ (0, 1), assume that with probability at least 1 − γ, ∥m̂ℓ0 − m0 ∥∞,[0,1] ≤ Ratem0 (N, γ),

∥fˆ0ℓ − f0 ∥∞ ≤ Ratef0 (N, γ).

When fˆ0ℓ appears in an inverse-propensity weight, we maintain the bounded-away-from-zero convention fˆ0ℓ (a|x) ≥ f . Assume additionally that the fitted outcome-quantile curves are uniformly Lipschitz in t, so that, for all folds ℓ, actions a, contexts x, and s, t ∈ [0, 1],     m0 (a, x) − m̂ℓ0 (a, x) (t) − m0 (a, x) − m̂ℓ0 (a, x) (s) ≤ 2U|t − s|, with the fitted Lipschitz constant absorbed into U. Define     δ δ rf := Ratef0 N, , rm := Ratem0 N, . 3L 3L 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 s
s
N 2L N-dim(Π) log e L d 2L log 8Jδ L + . VN (Π, δ) := N N Then, with probability at least 1 − 4δ, the regret satisfies R(π̂ DR ) ≤ Cor VN (Π, δ) + Cgrid η + Remnuis (N, δ, η), where

 Cor := 4LU M

2 1+ f



 Cgrid := 2LU U

,

10

6 4+ f

(3)

 ,

WASSERSTEIN P OLICY L EARNING FOR D ISTRIBUTIONAL O UTCOMES

and the nuisance-induced remainder is " Remnuis (N, δ, η) := 2LU

 !  2rf rm 4M rf 1 rm VN (Π, δ) + +2 1+ f f2 f2 # 5rf rm 4Urf + η . + f2 f2

Equivalently, there exists a constant C > 0, depending only on LU , M, f , U, L, and d, such that R(π̂ DR ) ≤ C [VN (Π, δ) + η + rf rm + rf VN (Π, δ) + rm VN (Π, δ) + rf η] . Proof sketch. The proof is detailed in Appendix C. The regret is first reduced to the uniform L2 deviation of the estimated policy-indexed quantile curve:

R(π̂

DR

Z 1 ) ≤ 2LU

sup 0

π∈Π

q̂πDR (t) − qπ (t)

!1/2

2 dt

.

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 rf rm ; the first-order terms rf and rm appear only multiplied by the policy-complexity factor VN (Π, δ) or the grid mesh η. Rate interpretation and Neyman orthogonality. The bound separates the oracle policy-learning error from the nuisance-induced error. Up to constants depending only on LU , M, f , U, L, and d, Theorem 11 implies R(π̂ DR ) ≲ VN (Π, δ) + η + rf rm + rf VN (Π, δ) + rm VN (Π, δ) + rf η. The key point is that Neyman orthogonality removes the standalone first-order nuisance bias. The only non-centered nuisance bias is of product order rf rm . The terms involving rf or rm alone are centered empirical-process or discretization remainders and are multiplied by VN (Π, δ) or η. Consequently, if rf = O(N −αf ),

rm = O(N −αm ),

e −1/2 ), VN (Π, δ) = O(N

η = O(N −1/2 ),

then   e N −1/2 + N −(αf +αm ) + N −(αf +1/2) + N −(αm +1/2) . R(π̂ DR ) = O In particular, if both nuisance estimators converge at the N −1/4 rate, then e −1/2 ). R(π̂ DR ) = O(N 11

H UANG* † L EUNG* W U* Z HANG*

5. 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 := N-dim(Π) in our upper bounds is essentially tight. We focus on a specific class of distributional utilities, the integrated quantile utility, defined as: Z α ν −1 (t) dt, α ∈ (0, 1]. Uα (ν) := 0

This utility fits within our framework as it satisfies the Lipschitz property (Assumption 8) via the √ Cauchy-Schwarz inequality: |Uα (ν1 ) − Uα (ν2 )| ≤ α W2 (ν1 , ν2 ). By constructing a hard instance for this specific utility, we demonstrate the hardness of the general problem. Theorem 12 (Minimax Lower Bound) Assume the policy class satisfies N-dim(Π) ≥ 1, let d = |A| ≥ 2, fix α ∈ (0, 1], and take 0 < f ≤ 1/d. Let q− , q+ : [0, 1] → I be two left-continuous, nondecreasing, U-Lipschitz quantile curves satisfying q− (t) ≤ q+ (t) and supt∈[0,1] max{|q− (t)|, |q+ (t)|} ≤ Rα M . Define ∆Q := 0 (q+ (t) − q− (t)) dt, and let Plower (q− , q+ ) be the subclass of distributions satisfying Assumptions 1–5 whose potential-outcome quantile curves take values in {q− , q+ } and whose behavior policy has overlap at least f . There exists a universal constant c0 > 0 such that for any learning algorithm that maps a dataset DN to a policy π̂, the worst-case regret is bounded from below by: ( s ) N-dim(Π) inf sup EDN ∼P⊗N [R(π̂)] ≥ c0 ∆Q min 1, , π̂ P∈Plower (q− ,q+ ) fN where R(π̂) = Uα (µP (πP∗ )) − Uα (µP (π̂)) is the regret. Proof sketch. The detailed proof is provided in Appendix D. We employ the method of Assouad’s Lemma by constructing a hard instance family parameterized by the vertices of a hypercube V = {±1}N-dim(Π) . 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 f , balancing the regret magnitude with this detection q probability yields the leading minimax dependence N-dim(Π)/(f N ) in the large-sample regime. Interpretation. Theorem 12 supports the rate sharpness of the proposed framework with respect q to N and policy-class complexity. In the large-sample regime, the lower bound scales as N-dim(Π)/(f N ), which matches the leading O(N −1/2 ) dependence of our DR upper bound (Theorem 11) up to logarithmic factors. The bound clarifies the cost of the distributional structure: while it depends on the integrated gap ∆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. 12

WASSERSTEIN P OLICY L EARNING FOR D ISTRIBUTIONAL O UTCOMES

6. 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 W2 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.pOur main theoretical results provide finite-sample regret guarantees with leading e N-dim(Π)/N ) on the Natarajan dimension of the policy class and the sample dependence O( size. We also established a minimax lower bound showing sharpness of the leading dependence on (N-dim(Π), 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).

References Christopher Adjaho and Timothy Christensen. Externally valid treatment choice. arXiv preprint arXiv:2205.05561, 1(1), 2022. Chunrong Ai, Yue Fang, and Haitian Xie. Data-driven policy learning for continuous treatments. Journal of Econometrics, 253:106170, 2026. Dionissi Aliprantis, Daniel Carroll, and Eric Young. The dynamics of the racial wealth gap. Working Papers 19-18R, Federal Reserve Bank of Cleveland, 2022. URL https://EconPapers. repec.org/RePEc:fip:fedcwq:191800. 13

H UANG* † L EUNG* W U* Z HANG*

Susan Athey and Stefan Wager. Policy learning with observational data. Econometrica, 89(1): 133–161, 2021. Satarupa Bhattacharjee, Bing Li, Xiao Wu, and Lingzhou Xue. Doubly robust estimation of causal effects for random object outcomes with continuous treatments. arXiv preprint arXiv:2506.22754, 2025. Aurélien Bibaut, Nathan Kallus, Maria Dimakopoulou, Antoine Chambaz, and Mark van Der Laan. Risk minimization from adaptively collected data: Guarantees for supervised and policy learning. Advances in neural information processing systems, 34:19261–19273, 2021. Victor Chernozhukov, Denis Chetverikov, Mert Demirer, Esther Duflo, Christian Hansen, Whitney Newey, and James Robins. Double/debiased machine learning for treatment and structural parameters: Double/debiased machine learning. The Econometrics Journal, 21(1), 2018. Victor Chernozhukov, Mert Demirer, Greg Lewis, and Vasilis Syrgkanis. Semi-parametric efficient policy learning with continuous actions. Advances in Neural Information Processing Systems, 32, 2019. Yifan Cui and Sukjin Han. Policy learning with distributional welfare. Journal of the American Statistical Association, pages 1–12, 2025. Kreske Ecker, Xavier de Luna, and Lina Schelin. Causal inference with a functional outcome. Journal of the Royal Statistical Society Series C: Applied Statistics, 73(1):221–240, 2024. David Haussler and Philip M Long. A generalization of sauer’s lemma. Journal of Combinatorial Theory, Series A, 71(2):219–240, 1995. Ying Jin, Zhimei Ren, Zhuoran Yang, and Zhaoran Wang. Policy learning “without” overlap: Pessimism and generalized empirical bernstein’s inequality. The Annals of Statistics, 53(4):1483– 1512, 2025. Nathan Kallus. Balanced policy evaluation and learning. Advances in neural information processing systems, 31, 2018. Nathan Kallus and Angela Zhou. Confounding-robust policy improvement. Advances in neural information processing systems, 31, 2018a. Nathan Kallus and Angela Zhou. Policy evaluation and optimization with continuous treatments. In International conference on artificial intelligence and statistics, pages 1243–1251. PMLR, 2018b. Nathan Kallus and Angela Zhou. Minimax-optimal policy learning under unobserved confounding. Management Science, 67(5):2870–2890, 2021. Nathan Kallus, Xiaojie Mao, Kaiwen Wang, and Zhengyuan Zhou. Doubly robust distributionally robust off-policy evaluation and learning. In International Conference on Machine Learning, pages 10598–10632. PMLR, 2022. Samir Khan, Martin Saveski, and Johan Ugander. Off-policy evaluation beyond overlap: partial identification through smoothness. arXiv preprint arXiv:2305.11812, 2023. 14

WASSERSTEIN P OLICY L EARNING FOR D ISTRIBUTIONAL O UTCOMES

Daido Kido. Distributionally robust policy learning with wasserstein distance. arXiv preprint arXiv:2205.04637, 2022. Toru Kitagawa and Aleksey Tetenov. Who should be treated? empirical welfare maximization methods for treatment choice. Econometrica, 86(2):591–616, 2018. Anders Bredahl Kock and David Preinerstorfer. Regularizing discrimination in optimal policy learning with distributional targets. arXiv preprint arXiv:2401.17909, 2024. Daisuke Kurisu, Yidong Zhou, Taisuke Otsu, and Hans-Georg Müller. Geodesic causal inference. arXiv preprint arXiv:2406.19604, 2024. Zhenhua Lin, Dehan Kong, and Linbo Wang. Causal inference on distribution functions. Journal of the Royal Statistical Society Series B: Statistical Methodology, 85(2):378–398, 2023. Charles F Manski. Statistical treatment rules for heterogeneous populations. Econometrica, 72(4): 1221–1246, 2004. Charles F Manski and Aleksey Tetenov. Statistical decision theory respecting stochastic dominance. The Japanese Economic Review, 74(4):447–469, 2023. Weibin Mo, Zhengling Qi, and Yufeng Liu. Learning optimal distributionally robust individualized treatment rules. Journal of the American Statistical Association, 116(534):659–674, 2021. Tong Mu, Yash Chandak, Tatsunori B Hashimoto, and Emma Brunskill. Factored dro: Factored distributionally robust policies for contextual bandits. Advances in Neural Information Processing Systems, 35:8318–8331, 2022. Victor M Panaretos and Yoav Zemel. Statistical aspects of wasserstein distances. Annual review of statistics and its application, 6(1):405–431, 2019. Yordan P Raykov, Hengrui Luo, Justin D Strait, and Wasiur R KhudaBukhsh. Kernel-based estimators for functional causal effects. arXiv preprint arXiv:2503.05024, 2025. Filippo Santambrogio. Optimal transport for applied mathematicians: Calculus of variations, pdes, and modeling, volume 87 of progress in nonlinear differential equations and their applications, 2015. Yi Shen, Pan Xu, and Michael Zavlanos. Wasserstein distributionally robust policy evaluation and learning for contextual bandits. Transactions on Machine Learning Research, 2024. ISSN 2835-8856. URL https://openreview.net/forum?id=NmpjDHWIvg. Featured Certification. Nian Si, Fan Zhang, Zhengyuan Zhou, and Jose Blanchet. Distributionally robust batch contextual bandits. Management Science, 69(10):5772–5793, 2023. Adith Swaminathan and Thorsten Joachims. Batch learning from logged bandit feedback through counterfactual risk minimization. The Journal of Machine Learning Research, 16(1):1731–1755, 2015. 15

H UANG* † L EUNG* W U* Z HANG*

Lorenzo Testa, Tobia Boschi, Francesca Chiaromonte, Edward H Kennedy, and Matthew Reimherr. Doubly-robust functional average treatment effect estimation. arXiv preprint arXiv:2501.06024, 2025. Alexandre B Tsybakov. Nonparametric estimators. In Introduction to Nonparametric Estimation, pages 1–76. Springer, 2008. Lan Wang, Yu Zhou, Rui Song, and Ben Sherwood. Quantile-optimal treatment regimes. Journal of the American Statistical Association, 113(523):1243–1254, 2018. Ruohan Zhan, Zhimei Ren, Susan Athey, and Zhengyuan Zhou. Policy learning with adaptively collected data. Management Science, 70(8):5270–5297, 2024. Pan Zhao, Antoine Chambaz, Julie Josse, and Shu Yang. Positivity-free policy learning with observational data. In International Conference on Artificial Intelligence and Statistics, pages 1918– 1926. PMLR, 2024. Yingqi Zhao, Donglin Zeng, A John Rush, and Michael R Kosorok. Estimating individualized treatment rules using outcome weighted learning. Journal of the American Statistical Association, 107(499):1106–1118, 2012. Xin Zhou, Nicole Mayer-Hamblett, Umer Khan, and Michael R Kosorok. Residual weighted learning for estimating individualized treatment rules. Journal of the American Statistical Association, 112(517):169–187, 2017. Yidong Zhou and Hans-Georg Mueller. Dynamic modelling of sparse longitudinal data and functional snippets with stochastic differential equations. Journal of the Royal Statistical Society Series B: Statistical Methodology, 87(3):833–849, 2025.

16

WASSERSTEIN P OLICY L EARNING FOR D ISTRIBUTIONAL O UTCOMES

Appendix A. Wasserstein barycenter representation Proof [Proof of Proposition 7] Let Q be the set of all left-continuous, non-decreasing functions q : [0, 1] → R. In one dimension, the map ν 7→ ν −1 is an isometry between (P2 (R), W2 ) and (Q, ∥ · ∥L2 ) (see Proposition 6). Fix a policy π and denote the random quantile curve as Q(t) := Y[π(X)]−1 (t). By Proposition 6, for any measure ν ∈ P2 (R) with quantile function qν := ν −1 ∈ Q, we have: Z 1  Z 1  2    2 E W2 (Y[π(X)], ν) = E |Q(t) − qν (t)| dt = E |Q(t) − qν (t)|2 dt. 0

0

For each fixed t ∈ [0, 1], we expand the integrand:     E |Q(t) − qν (t)|2 = E |Q(t) − E[Q(t)]|2 + |E[Q(t)] − qν (t)|2 , where the cross term vanishes since E[Q(t) − E[Q(t)]] = 0. Consequently, Z 1  2  |E[Q(t)] − qν (t)|2 dt, E W2 (Y[π(X)], ν) = Cπ + 0

R1

where Cπ := 0 E[|Q(t) − E[Q(t)]|2 ] dt is a constant independent of ν. It follows that any minimizer must satisfy qν (t) = E[Q(t)] for almost every t. By Assumption 4, Q is uniformly bounded by M ; hence, the map t 7→ E[Q(t)] is well-defined and belongs to L2 ([0, 1]). Furthermore, since Q(·) is non-decreasing and left-continuous almost surely, the pointwise expectation E[Q(·)] is also non-decreasing. Left-continuity of E[Q(·)] follows from the dominated convergence: for any t ∈ (0, 1], letting s ↑ t, we have E[Q(s)] → E[Q(t)] because |Q(s)| ≤ M and Q(s) → Q(t) a.s. Thus, E[Q] ∈ Q, ensuring thatRa minimizer exists within Q. Uniqueness follows from the strict 1 convexity of the functional q 7→ 0 |E[Q(t)] − q(t)|2 dt on Q. Therefore, the unique barycenter µ(π) satisfies: (µ(π))−1 (t) = E[Y[π(X)]−1 (t)], ∀t ∈ [0, 1].

Appendix B. Proof of Theorem 10 Before presenting the proof of Theorem 10, we first give two lemmas that are useful for our proofs. Lemma 13 Given δ > 0, M > 0, ϵ ≥ 0, L ≥ 0, β ∈ (0, 1], and a probability space (Ω, F, P). Suppose that f (t; ω) : T × Ω → R with T = [tLow , tUpp ] and |f (t; ω) − f (s; ω)| ≤ L|t − s|β + ϵ

∀ ω ∈ Ω.

Let Tpar := {tLow = t0 < t1 < · · · < tJ = tUpp } be a partition of [tLow , tUpp ] with tj = j(t −t ) tLow + UppJ Low and ( ) tUpp − tLow β δ P f (t; ω) ≥ M − ϵ − L ≤ , J J 17

H UANG* † L EUNG* W U* Z HANG*

for any t ∈ Tpar , then we have P{sup f (t; ω) ≥ M} ≤ δ. t∈T

Proof For any η > 0, there exists t̃ such that sup f (t; ω) − η < f (t̃; ω) ≤ sup f (t; ω). t∈T

t∈T

Now, choose tj ∈ Tpar such that tj−1 < t̃ ≤ tj . Then we have sup f (t; ω) = sup f (t; ω) − f (t̃; ω) − η + η + f (t̃; ω) t∈T

t∈T

≤ f (t̃; ω) + η = f (t̃; ω) − f (tj ; ω) + f (tj ; ω) + η ≤ |f (t̃; ω) − f (tj ; ω)| + max f (t; ω) + η t∈Tpar

β

β

≤ ϵ + L|t̃ − tj | + max f (t; ω) + η ≤ ϵ + L|tj − tj−1 | + max f (t; ω) + η t∈Tpar

=ϵ+L

t∈Tpar

β

tUpp − tLow J

+ max f (t; ω) + η. t∈Tpar

Since η is arbitrary, we must have sup f (t; ω) ≤ ϵ + L t∈T

tUpp −tLow β + max J t∈T

f (t; ω). As a result,

par

we have (

tUpp − tLow β P{sup f (t; ω) ≥ M} ≤ P max f (t; ω) ≥ M − ϵ − L t∈Tpar J t∈T ( ) J X tUpp − tLow β ≤ P f (tj ; ω) ≥ M − ϵ − L ≤ δ. J

)

j=1

Lemma 14 (Multi-class Sauer bound Haussler and Long (1995)) Let Π be a class of functions mapping X to a d-element set A. Let V := N-dim(Π). For any sample x1:n with n ≥ V , the number of distinct labelings mΠ (n) := |{(π(x1 ), . . . , π(xn )) : π ∈ Π}| satisfies mΠ (n) ≤ (end)V ,

and hence

log mΠ (n) ≤ V log(end).

Proof This is a standard consequence of the multi-class Sauer-type lemma for the Natarajan dimension. We include the (weak but sufficient) form used in the main text: mΠ (n) grows at most polynomially in n with exponent V and an additional factor d for the label set. Now we will prove Theorem 10 in the following. 18

WASSERSTEIN P OLICY L EARNING FOR D ISTRIBUTIONAL O UTCOMES

Proof [Proof of Theorem 10] Recall qπ (t) = µ(π)−1 (t) = E[Y[π(X)]−1 (t)]. By unconfoundedness and overlap, for each fixed (π, t), " # 1{A=π(X)} Y −1 (t) qπ (t) = E . f0 (A|X)   Also recall q̂πIPW (t) = PN 1{A=π(X)} Y −1 (t)/f0 (A|X) . First, the regret can be reduced to uniform estimation error. Let π̂ IPW ∈ arg maxπ∈Π U (µ̂IPW (π)). By the argmax property and Lipschitzness of U , we have R(π̂ IPW ) = U (µ(π ⋆ )) − U (µ(π̂ IPW ))     = U (µ(π ⋆ )) − U (µ̂IPW (π ⋆ )) + U (µ̂IPW (π ⋆ )) − U (µ̂IPW (π̂ IPW ))   + U (µ̂IPW (π̂ IPW )) − U (µ(π̂ IPW ))     ≤ U (µ(π ⋆ )) − U (µ̂IPW (π ⋆ )) + 0 + U (µ̂IPW (π̂ IPW )) − U (µ(π̂ IPW )) ≤ sup U (µ(π)) − U (µ̂IPW (π)) + sup U (µ̂IPW (π)) − U (µ(π)) π∈Π

π∈Π IPW

= 2 sup U (µ̂

(π)) − U (µ(π))

π∈Π

≤ 2LU sup W2 (µ̂IPW (π), µ(π)), π∈Π

To handle the monotonicity constraint, we use the fact that qπ ∈ Q. Since ΠQ is a non-expansive contraction in L2 , we have: W2 (µ̂IPW (π), µ(π)) = ∥q̂πIPW,↑ − qπ ∥L2 ≤ ∥q̂πIPW − qπ ∥L2 ≤ sup q̂πIPW (t) − qπ (t) . t∈[0,1]

Therefore, it follows that R(π̂ IPW ) ≤ 2LU sup sup q̂πIPW (t) − qπ (t) . t∈[0,1] π∈Π

Define

"

1{A=π(X)} Y −1 (t) ∆(t) := sup (PN − P) f0 (A|X) π∈Π

# .

It suffices to bound supt∈[0,1] ∆(t). For any fixed t, conditional on X1:N , the set of distinct labelings {(π(X1 ), . . . , π(XN )) : π ∈ Π} has cardinality mΠ (N ). For a fixed labeling, the summands are i.i.d. and bounded by |Y −1 (t)|/f0 (A|X) ≤ M/f . Thus Hoeffding’s inequality and a union bound over labelings yield   N ε2 P{∆(t) ≥ ε} ≤ 2 mΠ (N ) exp − . 2(M/f )2 Let Tpar = {tj : j = 0, . . . , J } be the uniform grid with J + 1 points (mesh η = 1/J ). A union bound over tj yields: with probability at least 1 − δ, v   u u 2 log 2(J +1)mΠ (N ) t δ M max ∆(tj ) ≤ . 0≤j≤J f N 19

H UANG* † L EUNG* W U* Z HANG*

Now we will extend from grid to continuum by Lipschitz regularity. By Assumption 5, ∆(·) is pathwise Lipschitz with constant U/f for both P and PN terms. For any t, let tj be the nearest grid point such that |t − tj | ≤ η. Then:  −1   −1  Y (t) Y −1 (tj ) Y (t) Y −1 (tj ) 2U ∆(t) ≤ ∆(tj ) + sup − −E − ≤ ∆(tj ) + η. f f f f f π 0 0 0 0 Taking the supremum over t ∈ [0, 1], we obtain supt ∆(t) ≤ maxj ∆(tj ) + (2U /f )η. Finally, we convert mΠ (N ) to Natarajan dimension. √ Recall V := N-dim(Π). By Lemma 14, √ √ log mΠ (N ) ≤ V log(eN d). Using a + b ≤ a + b, we obtain  v   u r u 2 log 2(J +1)  t δ M  2V log(eN d) + . max ∆(tj ) ≤   j f N N Combining the above bounds yields the theorem statement.

Appendix C. Proof of Theorem 11 Proof [Proof of Theorem 11] Recall qπ (t) = µ(π)−1 (t) = E[Y[π(X)]−1 (t)]. By unconfoundedness, qπ (t) = E[m0 (π(X), X)(t)],

m0 (a, x)(t) := E[Y −1 (t) | A = a, X = x].

For notational convenience, define the reciprocal weights by g0 (a, x) :=

1 , f0 (a|x)

g(a, x) := g0 (a, x),

ĝ0ℓ (a, x) :=

1 fˆ0ℓ (a|x)

.

When fˆ0ℓ appears in a denominator, it is maintained under the bounded-away-from-zero convention fˆ0ℓ (a|x) ≥ f . Throughout the proof, we condition on the training folds I−ℓ used to construct fˆ0ℓ and m̂ℓ0 ; hence the nuisance estimates are fixed when taking expectations over the held-out fold Iℓ . As in scalar policy learning, the regret can be decomposed as R(π̂ DR ) = U (µ(π ⋆ )) − U (µ(π̂ DR ))     = U (µ(π ⋆ )) − U (µ̂DR (π ⋆ )) + U (µ̂DR (π ⋆ )) − U (µ̂DR (π̂ DR ))   + U (µ̂DR (π̂ DR )) − U (µ(π̂ DR ))     ≤ U (µ(π ⋆ )) − U (µ̂DR (π ⋆ )) + 0 + U (µ̂DR (π̂ DR )) − U (µ(π̂ DR )) ≤ sup U (µ(π)) − U (µ̂DR (π)) + sup U (µ̂DR (π)) − U (µ(π)) π∈Π

π∈Π DR

= 2 sup U (µ̂

(π)) − U (µ(π))

π∈Π

≤ 2LU sup W2 (µ̂DR (π), µ(π)). π∈Π

20

WASSERSTEIN P OLICY L EARNING FOR D ISTRIBUTIONAL O UTCOMES

Using the L2 representation of the Wasserstein distance and the non-expansiveness of the isotonic projection ΠQ , we observe that for any fixed π ∈ Π:

DR

W2 (µ̂

(π), µ(π)) = ∥q̂πDR,↑ − qπ ∥L2 ≤ ∥q̂πDR − qπ ∥L2 =

Z 1

1/2

2 q̂πDR (t) − qπ (t) dt

.

0

To obtain a uniform bound over Π, we note that for each t ∈ [0, 1] and any π ∈ Π: q̂πDR (t) − qπ (t) ≤ sup q̂πDR ′ (t) − qπ ′ (t) . π ′ ∈Π

Squaring both sides and integrating over t ∈ [0, 1], we obtain: Z 1

2 q̂πDR (t) − qπ (t) dt ≤

Z 1 sup 0

0

π ′ ∈Π

q̂πDR ′ (t) − qπ ′ (t)

2 dt.

Since the right-hand side is independent of π, we can take the supremum over π ∈ Π on the left-hand side and take the square root: Z 1 sup π∈Π

1/2 DR ≤ q̂π (t) − qπ (t) dt

Z 1

2

0

0

sup q̂πDR (t) − qπ (t)

2

!1/2 dt

.

π∈Π

Consequently, define D(t) := sup q̂πDR (t) − qπ (t) . π∈Π

Then the preceding display gives

R(π̂

DR

Z 1 ) ≤ 2LU

2

D(t) dt

1/2 .

0

Therefore, now it suffices to bound D(t), which can be split into three main steps: Step I: Decomposing D(t).

Note that

D(t) = sup qπ (t) − q̂πDR (t) π∈Π

" # L   1X 1{A=π(X)} −1 −1 ℓ ℓ = sup E Y[π(X)] (t) − Pn,ℓ m̂0 (π(X), X) + (Y − m̂0 (π(X), X)) (t) . L π∈Π fˆ0ℓ (A|X) ℓ=1

21

H UANG* † L EUNG* W U* Z HANG*

Denote ĝ0ℓ (·, ·) = ˆℓ 1 . We then have f0 (·|·)

D(t) " # L  1X 1{A=π(X)} −1 ℓ ℓ = sup E Y[π(X)] (t) − Pn,ℓ m̂0 (π(X), X) + (Y − m̂0 (π(X), X)) (t) L π∈Π fˆ0ℓ (A|X) ℓ=1   1{A=π(X)} −1 ℓ ℓ   (Y − m̂ (π(X), X)) m̂ (π(X), X) +   L 0   0 1 X fˆ0ℓ (A|X) ≤ sup Pn,ℓ (t)   1 π∈Π L  ℓ=1   − m0 (π(X), X) − {A=π(X)} (Y −1 − m0 (π(X), X)) f0 (A|X)   L 1{A=π(X)} −1 1 X Pn,ℓ m0 (π(X), X) + (Y − m0 (π(X), X)) (t) − E[Y −1 [π(X)]](t) + sup f0 (A|X) π∈Π L ℓ=1     L 1 X (4) ≤ sup Pn,ℓ 1{A=π(X)} Y −1 − m̂ℓ0 (π(X), X) (t) ĝ0ℓ (A, X) − g(A, X) π∈Π L ℓ=1   L
1 X ℓ + sup Pn,ℓ 1{A=π(X)} g(A, X) − 1 m0 (π(X), X) − m̂0 (π(X), X) (t) (5) π∈Π L ℓ=1   1{A=π(X)} −1 (6) (Y − m0 (π(X), X)) (t) − E[Y −1 [π(X)]](t) . + sup PN m0 (π(X), X) + f0 (A|X) π∈Π 

−1

The goal is finding stochastic bounds of Eqns. (4) - (6). Before bounding these terms, we record the two orthogonality identities that remove the standalone first-order nuisance bias. For any fixed π, t, and fold ℓ, conditioning on I−ℓ gives h i

E 1{A=π(X)} Y −1 (t) − m0 (π(X), X)(t) ĝ0ℓ (A, X) − g(A, X) | I−ℓ = 0, h i

E 1{A=π(X)} g(A, X) − 1 m0 (π(X), X) − m̂ℓ0 (π(X), X) (t) | I−ℓ = 0.

(7) (8)

Indeed, (7) follows from   E Y −1 (t) − m0 (A, X)(t) | A, X = 0, while (8) follows from   E 1{A=π(X)} g(A, X) − 1 | X = f0 (π(X)|X)

1 − 1 = 0. f0 (π(X)|X)

Thus, the only non-centered nuisance bias is of product order in the propensity and outcomeregression errors. Step II: Studying Eqns. (4) and (5). corrected decomposition.

We now bound the two nuisance-dependent terms in the

22

WASSERSTEIN P OLICY L EARNING FOR D ISTRIBUTIONAL O UTCOMES

We first bound Eqn. (4). Indeed, we have Eqn. (4)     L 1 X −1 ℓ ℓ = sup Pn,ℓ 1{A=π(X)} Y − m̂0 (π(X), X) (t) ĝ0 (A, X) − g(A, X) π∈Π L ℓ=1       −1 ℓ ℓ 1 Y − m̂ (π(X), X) (t) ĝ (A, X) − g(A, X) L {A=π(X)} 0 0   1 X      ≤ sup Pn,ℓ    L π∈Π ℓ=1 − E 1{A=π(X)} Y −1 − m̂ℓ0 (π(X), X) (t) ĝ0ℓ (A, X) − g(A, X)      L 1 X −1 ℓ ℓ + . sup E 1{A=π(X)} Y − m̂0 (π(X), X) (t) ĝ0 (A, X) − g(A, X) L π∈Π

(9)

(10)

ℓ=1

We bound Eqn. (9). Denote       1{A=π(X)} Y −1 − m̂ℓ0 (π(X), X) (t) ĝ0ℓ (A, X) − g(A, X)        Γℓ2 (t) := sup Pn,ℓ   .  π∈Π − E 1{A=π(X)} Y −1 − m̂ℓ0 (π(X), X) (t) ĝ0ℓ (A, X) − g(A, X)

Also, denote       −1 ℓ ℓ 1{A=ai } Y − m̂0 (ai , X) (t) ĝ0 (A, X) − g(A, X)   ℓ;i      Pn,ℓ    := Γ2 (t). − E 1{A=ai } Y −1 − m̂ℓ0 (ai , X) (t) ĝ0ℓ (A, X) − g(A, X) Define nℓ := |Iℓ | = N/L. Using Hoeffding’s inequality, we have for any ε > 0,   n o n ℓ ε2 ℓ;i P Γ2 (t) ≥ ε I−ℓ ≤ 2 exp − 2 , 2B where



  ℓ 1{A=ai } Y (t) ĝ0 (A, X) − g(A, X)     −1 ℓ ≤ 1{A=ai } Y − m0 (ai , X) (t) ĝ0 (A, X) − g(A, X)     ℓ ℓ + 1{A=ai } m0 (ai , X) − m̂0 (ai , X) (t) ĝ0 (A, X) − g(A, X)   ℓ ℓ ∥fˆ0 − f0 ∥∞ 2M + ∥m̂0 − m0 ∥∞,[0,1] ≤ := B. f2 −1

− m̂ℓ0 (ai , X)



Since there are at most mΠ (nℓ ) distinct labeling patterns realized by Π on Iℓ . Applying a union bound over these patterns yields   n o n ℓ ε2 ℓ P Γ2 (t) ≥ ε I−ℓ ≤ 2mΠ (nℓ ) exp − 2 . 2B 23

H UANG* † L EUNG* W U* Z HANG*

Choose ε so that the RHS equals δ:

  n ℓ ε2 2 mΠ (nℓ ) exp − 2 = δ 2B

⇐⇒

v   u u 2 log 2mΠ (nℓ ) t δ ε=B . nℓ

As a result, for each fixed t ∈ [0, 1], with probability at least 1 − δ, 

v   u 2M + ∥m̂ℓ0 − m0 ∥∞,[0,1] u 2 log 2mΠ (nℓ ) t δ Γℓ2 (t) ≤ nℓ f2   s ℓ ℓ ˆ
s 2∥f0 − f0 ∥∞ 2M + ∥m̂0 − m0 ∥∞,[0,1] 2L N-dim(Π) log e N 2L log d L  + ≤ 2 N N f ∥fˆ0ℓ − f0 ∥∞

2 δ

 .

In particular, for each fixed grid point tj , with probability at least 1 − Jδ ,

2∥fˆ0ℓ − f0 ∥∞



2M + ∥m̂ℓ0 − m0 ∥∞,[0,1]

Γℓ2 (tj ) ≤

f2 s ×

 Next, we consider



s
 2J 2L N-dim(Π) log e N 2L log d L δ . + N N

Y −1 − m̂ℓ0 (π(X), X)



  −1 ℓ (t) − Y − m̂0 (π(X), X) (s) . Note that for any

s, t,

    −1 ℓ −1 ℓ Y − m̂0 (π(X), X) (t) − Y − m̂0 (π(X), X) (s) ≤ Y −1 (t) − Y −1 (s) + m̂ℓ0 (π(X), X)(s) − m̂ℓ0 (π(X), X)(t) ≤ U|t − s| + m̂ℓ0 (π(X), X)(s) − m0 (π(X), X)(s) + |m0 (π(X), X)(s) − m0 (π(X), X)(t)| + m0 (π(X), X)(t) − m̂ℓ0 (π(X), X)(t) ≤ U|t − s| + ∥m̂ℓ0 − m0 ∥∞,[0,1] + U |t − s| + ∥m̂ℓ0 − m0 ∥∞,[0,1] = 2U|t − s| + 2 ∥m̂ℓ0 − m0 ∥∞,[0,1] , 24

WASSERSTEIN P OLICY L EARNING FOR D ISTRIBUTIONAL O UTCOMES

we thus have

Γℓ2 (t) − Γℓ2 (s)      E 1{A=π(X)} Y −1 − m̂ℓ0 (π(X), X) (t) ĝ0ℓ (A, X) − g(A, X)      ≤ sup π∈Π −1 ℓ ℓ − E 1{A=π(X)} Y − m̂0 (π(X), X) (s) ĝ0 (A, X) − g(A, X)     −1 ℓ ℓ Pn,ℓ 1{A=π(X)} Y − m̂0 (π(X), X) (s) ĝ0 (A, X) − g(A, X)     + sup π∈Π −1 ℓ ℓ − Pn,ℓ 1{A=π(X)} Y − m̂0 (π(X), X) (t) ĝ0 (A, X) − g(A, X)
4∥fˆ0ℓ − f0 ∥∞ U|t − s| + ∥m̂ℓ0 − m0 ∥∞,[0,1] . ≤ f2

Similarly, we also have Γℓ2 (s) − Γℓ2 (t) ≤ cases, we have Γℓ2 (t) − Γℓ2 (s)

4∥fˆ0ℓ −f0 ∥∞ (U |t−s|+∥m̂ℓ0 −m0 ∥∞,[0,1] ) . f2

Combining the two

4∥fˆ0ℓ −f0 ∥∞ (U |t−s|+∥m̂ℓ0 −m0 ∥∞,[0,1] ) ≤ . By Lemma 13, we then have, f2

with probability at least 1 − δ,

2∥fˆ0ℓ − f0 ∥∞ sup Γℓ2 (t) ≤



2M + ∥m̂ℓ0 − m0 ∥∞,[0,1]



f2

t∈[0,1]

s ×

+

s
 2J d 2L N-dim(Π) log e N 2L log L δ  + N N

4∥fˆ0ℓ − f0 ∥∞ ∥m̂ℓ0 − m0 ∥∞,[0,1] f2

+

4η U∥fˆ0ℓ − f0 ∥∞ . f2

As a result, we have, with probability at least 1 − 3δ,   2 Ratef0 (N, δ) 2M + Ratem0 (N, δ) sup Γℓ2 (t) ≤

f2

t∈[0,1]

s × +

s
 2J 2L N-dim(Π) log e N d 2L log L δ  + N N

4 Ratef0 (N, δ) Ratem0 (N, δ) 4η U Ratef0 (N, δ) + . f2 f2 25

H UANG* † L EUNG* W U* Z HANG*

Equivalently, we have, with probability at least 1 − Lδ ,

sup Γℓ2 (t) ≤

δ 2 Ratef0 N, 3L



δ 2M + Ratem0 N, 3L



f2

t∈[0,1]

s

s
 6J L 2L N-dim(Π) log e N 2L log d L δ  × + N N


δ δ δ 4 Ratef0 N, 3L Ratem0 N, 3L 4η U Ratef0 N, 3L + + , f2 f2

which implies that, with probability at least 1 − δ, L

L

1X ℓ 1X Eqn. (9) = Γ2 (t) ≤ sup Γℓ2 (t) L L t∈[0,1] ℓ=1 ℓ=1  

δ δ 2 Ratef0 N, 3L 2M + Ratem0 N, 3L ≤ f2 s
s
 6J L 2L N-dim(Π) log e N d 2L log L δ  × + N N


δ δ δ 4η U Ratef0 N, 3L 4 Ratef0 N, 3L Ratem0 N, 3L + . + f2 f2

(11)

Next, we bound Eqn. (10). By the orthogonality identity (7), conditioning on I−ℓ yields

h i

E 1{A=π(X)} Y −1 − m̂ℓ0 (π(X), X) (t) ĝ0ℓ (A, X) − g(A, X) | I−ℓ i h

= E 1{A=π(X)} m0 (π(X), X) − m̂ℓ0 (π(X), X) (t) ĝ0ℓ (A, X) − g(A, X) | I−ℓ h i
≤ E m0 (π(X), X) − m̂ℓ0 (π(X), X) (t) ĝ0ℓ (A, X) − g(A, X) | I−ℓ ≤

∥m̂ℓ0 − m0 ∥∞,[0,1] ∥fˆ0ℓ − f0 ∥∞ f2

(12)

.

Consequently, on the nuisance event in Theorem 11,

Eqn. (10) ≤

δ δ Ratef0 N, 3L Ratem0 N, 3L f2 26

.

(13)

WASSERSTEIN P OLICY L EARNING FOR D ISTRIBUTIONAL O UTCOMES

Combining Eqns. (11) and (13), we have, with probability at least 1 − δ, Eqn. (4) ≤ Eqn. (9) + Eqn. (10)  

δ δ 2 Ratef0 N, 3L 2M + Ratem0 N, 3L ≤ f2 s

s N 6J L 2L N-dim(Π) log e L d 2L log δ  + × N N


δ δ δ 5 Ratef0 N, 3L Ratem0 N, 3L 4ηU Ratef0 N, 3L + + . f2 f2

(14)

We also bound Eqn. (5) in a similar manner. Indeed, Eqn. (5)   L
1 X Pn,ℓ 1{A=π(X)} g(A, X) − 1 m0 (π(X), X) − m̂ℓ0 (π(X), X) (t) π∈Π L ℓ=1    
ℓ 1 g(A, X) − 1 m (π(X), X) − m̂ (π(X), X) (t) 0 L 0   {A=π(X)} 1 X     sup Pn,ℓ  ≤  
L π∈Π ℓ=1 − E 1{A=π(X)} g(A, X) − 1 m0 (π(X), X) − m̂ℓ0 (π(X), X) (t) = sup

L

+

1 X sup E L π∈Π

 1{A=π(X)} g(A, X) − 1

   m0 (π(X), X) − m̂ℓ0 (π(X), X) (t) .

(15)

(16)

ℓ=1

We bound Eqn. (15). Let ∆ℓ (a, x, t) := m0 (a, x)(t) − m̂ℓ0 (a, x)(t),

Wπ (X, A) := 1{A=π(X)} g(A, X) − 1.

By the fold-conditional orthogonality identity (8), for every fixed (π, t), E [Wπ (X, A)∆ℓ (π(X), X, t) | I−ℓ ] = 0. Thus the centered empirical process in Eqn. (15) can be written as Γℓ3 (t) := sup |Pn,ℓ [Wπ (X, A)∆ℓ (π(X), X, t)]| . π∈Π

For the growth-function argument below, we first condition on the nuisance-training folds and on the held-out covariates; to avoid overloading notation, this conditioning is still written as conditioning on I−ℓ . Let nℓ := |Iℓ | = N/L and define the set of labeling patterns realized by Π on the held-out fold, Vℓ := {(π(Xr ))r∈Iℓ : π ∈ Π} . For a fixed pattern v = (vr )r∈Iℓ ∈ Vℓ , define
Zr,vr (t) := 1{Ar =vr } g(vr , Xr ) − 1 ∆ℓ (vr , Xr , t). 27

H UANG* † L EUNG* W U* Z HANG*

Then

 E [Zr,vr (t) | I−ℓ ] = 0,

|Zr,vr (t)| ≤

1 1+ f



∥m̂ℓ0 − m0 ∥∞,[0,1] =: Bℓ .

Using Hoeffding’s inequality, for any fixed v ∈ Vℓ , fixed t, and any ε > 0,      1 X  n ℓ ε2 P Zr,vr (t) ≥ ε I−ℓ ≤ 2 exp − 2 .  nℓ  2Bℓ r∈Iℓ

Since |Vℓ | ≤ mΠ (nℓ ), a union bound over the policy-induced labeling patterns gives   n o n ℓ ε2 ℓ P Γ3 (t) ≥ ε | I−ℓ ≤ 2mΠ (nℓ ) exp − 2 . 2Bℓ Choosing the right-hand side to be δ yields, for each fixed t ∈ [0, 1], with probability at least 1 − δ, s   1 2 log (2mΠ (nℓ )/δ) Γℓ3 (t) ≤ 1 + . ∥m̂ℓ0 − m0 ∥∞,[0,1] f nℓ Using the multi-class Sauer bound and nℓ = N/L, set u = δ/(4LJ ) in the preceding fixed-t bound. Then, for any fixed fold-grid pair (ℓ, j), with probability at least 1 − u,  s
r   N 2L N-dim(Π) log e d 1 2L log (8J L/δ)  L Γℓ3 (tj ) ≤ 2 1 + ∥m̂ℓ0 − m0 ∥∞,[0,1]  + . f N N A union bound over all LJ fold-grid pairs implies that the same bound holds simultaneously for all ℓ = 1, . . . , L and all grid points j = 1, . . . , J , with probability at least 1 − δ/4. Next, for s, t ∈ [0, 1], the elementary inequality sup |uπ | − sup |vπ | ≤ sup |uπ − vπ | π

π

π

implies |Γℓ3 (t) − Γℓ3 (s)| ≤ sup |Pn,ℓ [Wπ (X, A){∆ℓ (π(X), X, t) − ∆ℓ (π(X), X, s)}]| . π∈Π

For the grid-to-continuum step, we use the fitted-Lipschitz condition maintained in Theorem 11, so that sup |∆ℓ (a, x, t) − ∆ℓ (a, x, s)| ≤ 2U|t − s|. a,x

In particular, the deterministic envelope bound gives   1 ℓ ℓ |Γ3 (t) − Γ3 (s)| ≤ 2 1 + U|t − s|. f On this fixed-grid event, Lemma 13 and the deterministic t-oscillation bound below imply that, for every ℓ = 1, . . . , L,  s
r     N 2L N-dim(Π) log e d 1 2L log (8J L/δ)  1 L sup Γℓ3 (t) ≤ 2 1 + ∥m̂ℓ0 −m0 ∥∞,[0,1]  + +2 1 + Uη. f N N f t∈[0,1] 28

WASSERSTEIN P OLICY L EARNING FOR D ISTRIBUTIONAL O UTCOMES

Intersecting the preceding fixed-grid event with the nuisance event   δ ℓ max ∥m̂0 − m0 ∥∞,[0,1] ≤ Ratem0 N, , 1≤ℓ≤L 3L and using a union bound over the nuisance folds, we obtain, with probability at least 1 − δ, L

L

1X ℓ 1X Eqn. (15) ≤ Γ3 (t) ≤ sup Γℓ3 (t) L L t∈[0,1] ℓ=1 ℓ=1 s 
r     d δ  2L N-dim(Π) log e N 2L log (8J L/δ) 1 L  Ratem0 N, + ≤2 1+ f 3L N N   1 Uη. +2 1+ f (17) Next, we bound Eqn. (16). By the orthogonality identity (8), conditioning on I−ℓ gives i h

E 1{A=π(X)} g(A, X) − 1 m0 (π(X), X) − m̂ℓ0 (π(X), X) (t) | I−ℓ h  i 
(18) = E E 1{A=π(X)} g(A, X) − 1 | X m0 (π(X), X) − m̂ℓ (π(X), X) (t) | I−ℓ 0

= 0. Therefore, Eqn. (16) = 0.

(19)

Combining Eqns. (17) and (19), we have, with probability at least 1 − δ, Eqn. (5) ≤ Eqn. (15) + Eqn. (16) s
s
     N 8J L 2L log δ 1 δ  2L N-dim(Π) log e L d  ≤2 1+ Ratem0 N, + f 3L N N   1 +2 1+ Uη. f Step III: Studying Eqn. (6)

Recall that   1{A=π(X)} −1 −1 E[Y [π(X)]](t) = E m0 (π(X), X) + (Y − m0 (π(X), X)) (t), f0 (A|X)

we therefore have   1{A=π(X)} −1 sup PN m0 (π(X), X) + (Y − m0 (π(X), X)) (t) − E[Y −1 [π(X)]](t) f0 (A|X) π∈Π   1{A=π(X)} −1 PN m0 (π(X), X) + (Y − m0 (π(X), X)) (t) f0 (A|X)   = sup := △(t). 1{A=π(X)} −1 π∈Π − E m0 (π(X), X) + (Y − m0 (π(X), X)) (t) f0 (A|X) 29

(20)

H UANG* † L EUNG* W U* Z HANG*

Also, denote   1{A=ai } −1 PN m0 (ai , X) + (Y − m0 (ai , X)) (t) f0 (A|X)   := △i (t). 1{A=ai } −1 − E m0 (ai , X) + (Y − m0 (ai , X)) (t) f0 (A|X) Using Hoeffding’s inequality, we have: for any ε > 0,    i N ε2 P △ (t) ≥ ε X1:N ≤ 2 exp − 2 , 2B 1

−1 − m (a , X))(t)| ≤ B such that we can choose B as (1 + 2 )M . i} where |m0 (ai , X) + f0{A=a 0 i f (A|X) (Y Since there are at most mΠ (N ) distinct labeling patterns realized by Π on X1:N . Applying a union bound over these patterns yields    N ε2 P △(t) ≥ ε X1:N ≤ 2mΠ (N ) exp − 2 . 2B

Choose ε so that the RHS equals δ:   N ε2 2 mΠ (N ) exp − 2 = δ 2B

⇐⇒

ε=B

v   u u 2 log 2mΠ (N ) t δ N

.

As a result, for each fixed t ∈ [0, 1], with probability at least 1 − δ, v   u   u 2 log 2mΠ (N ) t δ 2 △(t) ≤ 1 + M . f N In particular, for each fixed grid point tj , with probability at least 1 − Jδ , s
s
   N 2J 2L N-dim(Π) log e L d 2L log δ 2 . △(tj ) ≤ 2 1 + + M f N N Note that △ (t) − △(s)   1{A=π(X)} −1 (Y − m0 (π(X), X)) (t) PN m0 (π(X), X) + f0 (A|X)   ≤ sup 1{A=π(X)} −1 π∈Π − PN m0 (π(X), X) + (Y − m0 (π(X), X)) (s) f0 (A|X)   1{A=π(X)} −1 E m0 (π(X), X) + (Y − m0 (π(X), X)) (s) f0 (A|X)   + sup 1{A=π(X)} −1 π∈Π − E m0 (π(X), X) + (Y − m0 (π(X), X)) (t) f0 (A|X)   2 U|t − s|, ≤2 1+ f

30

WASSERSTEIN P OLICY L EARNING FOR D ISTRIBUTIONAL O UTCOMES

  implying that | △ (t) − △(s)| ≤ 2 1 + f2 U|t − s| by symmetrization. By Lemma 13, we then have, with probability at least 1 − δ, s
s
   2J 2L N-dim(Π) log e N d 2L log 2 L δ  M sup △ (t) ≤2 1 + + f N N t∈[0,1]   2 + 1+ Uη. f

Final combination.

(21)

Define   δ rm := Ratem0 N, . 3L

  δ rf := Ratef0 N, , 3L

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 s
s
2L N-dim(Π) log e N d 2L log 8Jδ L L VN (Π, δ) := + . N N Combining Eqns. (14), (20), and (21), we have, with probability at least 1 − 4δ, R(π̂

DR

      2 6 ) ≤ 2LU 2M 1 + VN (Π, δ) + 4 + Uη f f "   ! 4M rf 1 2rf rm + +2 1+ + 2LU rm VN (Π, δ) f f2 f2 # 4Urf 5rf rm + η . + f2 f2

This matches (3).

Appendix D. Proof of Theorem 12 Proof Let V := N-dim(Π) ≥ 1. By definition of the Natarajan dimension, there exists a set S = {x1 , . . . , xV } ⊂ X that is Natarajan-shattered by Π. Hence there exist two functions f1 , f2 : S → A with f1 (xi ) ̸= f2 (xi ) such that for every subset S0 ⊆ S there exists π ∈ Π with π(x) = f1 (x) for x ∈ S0 and π(x) = f2 (x) for x ∈ S \ S0 . Fix δ0 ∈ (0, 1/4) (to be chosen later) and let V = {±1}V . We construct a finite subset {Pv }v∈V ⊂ Plower (q− , q+ ) such that learning the optimal policy reduces to identifying the vertex v. First, we construct distribution-valued potential outcomes. Under the distribution Pv : • X is uniformly distributed on S. • The behavior policy satisfies, for each xi ∈ S, f0 (f2 (xi ) | xi ) = f ,

f0 (a | xi ) = f

for a ∈ / {f1 (xi )},

and f0 (f1 (xi ) | xi ) = 1 − (d − 1)f . 31

H UANG* † L EUNG* W U* Z HANG*

• Potential outcomes are distribution-valued and take only the two valid base quantile curves q− and q+ from Theorem 12. Let ∆(t) := q+ (t) − q− (t) ≥ 0. For each i and each action a ∈ A, define a Bernoulli latent variable Zi,a ∈ {0, 1}, drawn conditionally on X = xi and independently of A, and set Y[a]−1 (t) (X = xi ) := q− (t) + Zi,a ∆(t),

t ∈ [0, 1].

We choose the success probabilities: Zi,f1 (xi ) ∼ Bern(1/2),

Zi,f2 (xi ) ∼ Bern(1/2 + vi δ0 ),

Zi,a ∼ Bern(0) otherwise.

• The observed outcome satisfies Y = Y[A] (consistency). This construction satisfies Assumptions 1–5: unconfoundedness holds because the latent variables are conditionally independent of A given X, boundedness and quantile regularity follow from the properties of q− and q+ , and the overlap constant is f by design since 0 < f ≤ 1/d. Now we can evaluate Uα and identify the optimal policy. Under Pv , let µv (π) denote the policyinduced barycenter quantile (which depends on the distribution Pv ). By Proposition 7:   µv (π)−1 (t) = E(X,Y)∼Pv Y[π(X)]−1 (t) . For each xi , choosing f2 (xi ) instead of f1 (xi ) changes the mean quantile by δ0 vi ∆(t). Hence the optimal policy πv⋆ for Pv satisfies πv⋆ (xi ) =

( f2 (xi ), vi = +1, f1 (xi ), vi = −1,

and for any π ∈ Π, δ0 Uα (µv (πv⋆ )) − Uα (µv (π)) ≥ V

V X

!Z

1{π(xi ) ̸= πv⋆ (xi )}

α

∆(t) dt. 0

i=1

Let PN v denote the joint distribution of the dataset DN (containing N logged samples) under Pv . Recall that the estimator π̂ is a function of DN . sup EPN [Uα (µv (πv⋆ )) − Uα (µv (π̂))] ≥ v v∈V

1 X EPN [Uα (µv (πv⋆ )) − Uα (µv (π̂))] . v 2V v∈V

Therefore, based on the above results, we have V

1 X δ0 X X N ⋆ E Pv π̂(xi ) ̸= πv⋆ (xi ) N [Uα (µv (πv )) − Uα (µv (π̂))] ≥ P V V v 2 V ·2 v∈V i=1

v∈V

Z α ∆(t) dt. 0

We now apply Assouad’s symmetrization. For each i ∈ {1, . . . , V }, let Mi [v] be the vertex in V that differs from v only in the i-th coordinate. By swapping the order of summation and pairing 32

WASSERSTEIN P OLICY L EARNING FOR D ISTRIBUTIONAL O UTCOMES

each v with Mi [v] for vi = 1, we obtain X X 


 ⋆ N PN PN v π̂(xi ) ̸= πv (xi ) = v π̂(xi ) ̸= f2 (xi ) + PMi [v] π̂(xi ) ̸= f1 (xi ) v:vi =1

v∈V

≥

X 

 N PN v π̂(xi ) ̸= f2 (xi ) + 1 − PMi [v] π̂(xi ) ̸= f2 (xi )

v:vi =1

≥

X 

 N 1 − TV(PN v , PMi [v] ) ,

v:vi =1

where TV(P, Q) represents the total variation distance between P and Q. Applying this lower bound to the sum yields   Z V X X X
 α 1 δ0  1 ⋆ N N EPN [Uα (µv (πv )) − Uα (µv (π̂))] ≥ 1 − TV(Pv , PMi [v] )  ∆(t) dt.  V v 2V V 2 0 i=1 v∈V vi =1

v∈V

According to the relationship between TV distance and KL divergence (e.g., Lemma 2.6 in Tsybakov (2008)), we have: for any distributions P, Q, 1 − TV(P, Q) ≥

1 exp − KL(P∥Q) . 2

In our construction, Pv and PMi [v] differ only when (X = xi , A = f2 (xi )), which happens with probability V1 f . Conditioned on this event, the observed outcome reveals the Bernoulli latent variable Zi,f2 (xi ) ∼ Bern(1/2+δ0 ) versus Bern(1/2−δ0 ). A standard bound for Bernoulli KL implies that for δ0 ∈ (0, 1/4), 12f N 2 N KL(PN δ0 . v ∥PMi [v] ) ≤ V Choosing ( s ) 1 V δ0 := √ min 1, fN 24 ensures δ0 ∈ (0, 1/4) and 12f N 2 1 δ0 ≤ . V 2 1 −1/2 N Hence 1 − TV(PN . v , PMi [v] ) ≥ 2 e Plugging this bound into the preceding display yields

δ0 sup EPN [Uα (µv (πv⋆ )) − Uα (µv (π̂))] ≥ e−1/2 v 4 v∈V

Z α 0

( s )Z α e−1/2 V ∆(t) dt = √ min 1, (q+ (t)−q− (t)) dt. fN 4 24 0

Since {Pv }v∈V ⊂ Plower (q− , q+ ), this lower bound holds for supP∈Plower (q− ,q+ ) as well. This proves √ the theorem with c0 = e−1/2 /(4 24).

33