Wasserstein Policy Optimization

• Wasserstein Policy Optimization (WPO) is a geometric framework that updates policies as state-indexed probability distributions using the 2-Wasserstein metric.
• It leverages variational formulations, Wasserstein gradient flows, and proximal schemes to compute transport-driven updates for continuous-action reinforcement learning.
• Empirical studies show that WPO methods can achieve enhanced sample efficiency and stability compared to traditional KL-based policy optimization techniques.

Searching arXiv for papers on Wasserstein Policy Optimization and closely related methods. Wasserstein Policy Optimization (WPO) denotes policy optimization performed in probability space using Wasserstein geometry. In this perspective, a policy is not treated primarily as a vector of parameters but as a state-indexed family of action distributions, and successive updates are governed by optimal transport structure rather than solely by Kullback–Leibler trust regions or Euclidean parameter gradients. Across the literature, WPO appears both as a broad paradigm and as a collection of concrete algorithms, including Wasserstein gradient flows, Wasserstein trust-region methods, Wasserstein natural gradients, and Wasserstein proximal updates for entropy-regularized reinforcement learning (Pfau et al., 1 May 2025, Zhu et al., 3 Mar 2026).

1. Conceptual scope and geometric foundations

WPO rests on the view that policies are measurable maps from states into a Wasserstein space of action probabilities. In the most explicit geometric formulation, the admissible policy space is

$$\Pi := L^0(S; P(A)) = \{ \pi:S\to P(A) \mid \pi \text{ is } (B(S), B(P(A)))\text{-measurable}\},$$

where $P(A)$ denotes probability measures on the action space with finite second moment (Dus, 16 Apr 2026). This replaces the usual emphasis on parameter space by a geometry on state-conditional action distributions themselves.

The metric structure is typically induced by the $2$-Wasserstein distance. For probability measures $\mu,\nu$ on $\mathbb{R}^d$ with finite second moment,

$$W_2^2(\mu,\nu) := \inf_{\gamma\in\Gamma(\mu,\nu)} \int \|x-y\|^2\, d\gamma(x,y),$$

where $\Gamma(\mu,\nu)$ is the set of couplings with marginals $\mu,\nu$ (Zhu et al., 3 Mar 2026). In continuous-action reinforcement learning, this geometry assigns low cost to moving probability mass a short distance in action space and high cost to transporting it far away. That property differentiates WPO from KL-based methods, which constrain changes in distribution without encoding the underlying action metric.

A central mathematical object is the continuity equation associated with Wasserstein gradient flow. In the early formulation of policy optimization as Wasserstein gradient flow, policy updates are written as distributional transport on probability-measure space rather than direct parameter ascent (Zhang et al., 2018). In later formulations based on Otto calculus, the induced policy-space gradient is expressed through $\nabla_a Q^\pi(s,a)$, so the geometry of the action space enters the optimization rule explicitly (Dus, 16 Apr 2026). This suggests a unifying interpretation: many seemingly different WPO algorithms can be viewed as different discretizations, projections, or relaxations of transport-driven policy evolution.

A recurrent contrast is with KL trust-region methods such as TRPO and PPO. In those methods, successive policies are constrained in KL geometry; in WPO, the trust region is defined by an optimal-transport discrepancy. The practical implication is that WPO respects the metric structure of actions, so moving mass from $a$ to a nearby $P(A)$0 is treated as fundamentally different from moving it across distant regions of action space (Terpin et al., 2022).

2. Variational formulations and Wasserstein flows

One foundational route to WPO uses the Jordan–Kinderlehrer–Otto proximal scheme. In its generic form, a policy update is written as

$P(A)$1

so that improvement in the objective $P(A)$2 is balanced against Wasserstein proximity to the current policy (Ziesche et al., 2023). In reinforcement learning, $P(A)$3 is instantiated by a return-based free-energy functional, often including entropy regularization.

For entropy-regularized continuous-action RL, Wasserstein Proximal Policy Gradient gives a per-state proximal problem

$P(A)$4

with $P(A)$5 (Zhu et al., 3 Mar 2026). The update is a Wasserstein trust-region step with an explicit entropy term, and it decouples across states.

An operator-splitting derivation separates this proximal update into a drift step and a diffusion step. The drift step is an optimal transport update defined by a transport map $P(A)$6 satisfying

$P(A)$7

followed by a heat step

$P(A)$8

that is, Gaussian convolution with covariance $P(A)$9 (Zhu et al., 3 Mar 2026). In the small-step limit, the Monge map becomes $2$0, which yields the drift dynamics familiar from gradient-flow interpretations.

A closely related continuous-time formulation is Wasserstein policy gradient for entropy-regularized RL. There the state-conditional policy evolves according to

$2$1

equivalently via the Langevin SDE

$2$2

with Euler–Maruyama discretization

$2$3

(Zhu et al., 25 May 2026). In this form, WPO becomes a transport-plus-diffusion mechanism in action space.

3. Representative algorithmic instantiations

The name WPO encompasses several distinct algorithmic realizations rather than a single update rule. One direct actor-critic formulation derives a closed-form parameter update by projecting Wasserstein gradient flow over the space of all policies into a finite-dimensional parameter space. The resulting actor step is

$2$4

which combines deterministic-policy-gradient use of $2$5 with applicability to stochastic policies without the reparameterization trick (Pfau et al., 1 May 2025).

Another line uses trust-region duality directly. OT-TRPO formulates policy optimization with an optimal-transport discrepancy constraint,

$2$6

and proves a one-dimensional convex dual

$2$7

with an optimal policy update characterized by push-forwards and mass splitting (Terpin et al., 2022).

In deterministic systems, Truly Deterministic Policy Optimization uses Wasserstein-based quadratic regularization because KL and TV are ill-defined or maximal for deterministic policies. In that setting,

$2$8

and the resulting surrogate includes both a squared Wasserstein action-distance term and a policy-sensitivity term involving action-to-observation Jacobians (Saleh et al., 2022).

Structured policy classes have also motivated specialized WPO variants. For Gaussian mixture policies, policy adaptation is cast as a Wasserstein gradient flow on the space of GMMs, with Gaussian means and covariances updated by Bures–Wasserstein Riemannian optimization and mixture weights updated via an optimal-transport coupling over mixture components (Ziesche et al., 2023). For behavioral-distribution optimization, Wasserstein Natural Gradient defines a Wasserstein information matrix $2$9 through the local expansion

$\mu,\nu$0

and uses the preconditioned direction $\mu,\nu$1 in policy gradient or evolution-strategy updates (Moskovitz et al., 2020).

For orientation, the main instantiations can be summarized as follows.

Variant	Core update	Distinctive property
WGF/JKO WPO	Wasserstein proximal step	Policy optimization on measure space
OT-TRPO	1D dual over $\mu,\nu$2	Trust region via OT discrepancy
TDPO	Quadratic Wasserstein surrogate	Deterministic-policy regularization
WNG/WNPG	$\mu,\nu$3	Wasserstein natural gradient
Actor-critic WPO	$\mu,\nu$4	Closed-form projected flow
WPPG	OT drift + Gaussian heat	No log-density evaluation

This taxonomy suggests that “WPO” is best understood as a metric-geometric design principle, not as a single canonical algorithm.

4. Convergence theory and analytical structure

Convergence analyses of WPO differ sharply across formulations. In finite-action nonparametric settings, metric-aware trust-region methods prove monotonic performance improvement for Wasserstein policy optimization and show that, with a decaying Lagrangian multiplier for the trust-region constraint, WPO converges to global optimality; the same work also shows that Sinkhorn policy optimization converges to WPO as the entropic regularizer diminishes (Song et al., 2023).

For entropy-regularized continuous-action RL, WPPG establishes a global linear convergence rate under bounded rewards, bounded action sets, bounded initial negative entropy, and a uniform Talagrand transportation-information inequality $\mu,\nu$5. With

$\mu,\nu$6

and step size $\mu,\nu$7, the exact-evaluation theorem gives

$\mu,\nu$8

For actor–critic with inexact $\mu,\nu$9, the same quantity contracts up to an additive $\mathbb{R}^d$0 term, so approximation and stochastic errors do not accumulate (Zhu et al., 3 Mar 2026).

A parallel convergence theory for Wasserstein policy gradient exploits the Bellman structure of entropy-regularized RL rather than ordinary convexity. The key identity is the statewise KL representation of the soft Bellman residual,

$\mathbb{R}^d$1

where $\mathbb{R}^d$2 is the Gibbs policy tied to $\mathbb{R}^d$3. Together with a Bellman resolvent identity connecting value improvement to relative Fisher information and a uniform log-Sobolev inequality along the moving Gibbs family, this yields a distributional Polyak–Łojasiewicz condition and geometric contraction up to an $\mathbb{R}^d$4 discretization bias for the Euler–Maruyama scheme (Zhu et al., 25 May 2026).

A shorter convergence note reaches a similar qualitative conclusion for entropy-regularized WPO by combining monotonic energy dissipation with a local log-Sobolev inequality. Under the existence of a sufficiently regular solution to the gradient-flow equation, the value function is argued to converge linearly to the global optimum (Šiška et al., 21 May 2026). Taken together, these results indicate that the analytically tractable regime for WPO is presently centered on entropy-regularized problems, bounded or regularized action spaces, and either exact or controlled-approximation critics.

5. Empirical regimes and applications

Empirical work on WPO is concentrated in continuous control and robotics. The actor-critic WPO formulation was evaluated on the DeepMind Control Suite and a magnetic confinement fusion task, where results were reported as comparing favorably with state-of-the-art continuous-control methods (Pfau et al., 1 May 2025). In high-dimensional action spaces, WPO was described as taking off earlier than MPO, which in turn was faster than DDPG, with the trend strengthening as dimensionality increased.

WPPG reports experiments on MuJoCo continuous-control benchmarks including Hopper-v5, Walker2d-v5, HalfCheetah-v5, Reacher-v5, Swimmer-v5, and Humanoid-v5. In those experiments, WPPG with an explicit Gaussian actor was competitive with SAC, whereas WPPG-I, using an implicit actor $\mathbb{R}^d$5, consistently outperformed all baselines across nearly all tasks, including PPO, SAC, and WPO, and was notably more stable than WPO on challenging domains such as Humanoid and Swimmer, where WPO could be unstable or fail, including on Reacher (Zhu et al., 3 Mar 2026).

In robotics adaptation, Wasserstein gradient flow over Gaussian mixture policies has been evaluated on reaching motions, collision-avoidance behaviors, and multi-goal tasks. Representative results reported that WGF/WPO achieved success rate approximately $\mathbb{R}^d$6 after roughly $\mathbb{R}^d$7, $\mathbb{R}^d$8, and $\mathbb{R}^d$9 steps on reaching, collision avoidance, and multi-goal tasks respectively, while PPO required substantially more samples and SAC often failed or exhibited higher variance (Ziesche et al., 2023). This regime is especially structured: policies are initialized from demonstrations via EM, state dependence is produced by Gaussian Mixture Regression, and covariance updates exploit the Bures–Wasserstein manifold.

Deterministic Wasserstein policy optimization has been evaluated in robotic control environments with non-local rewards in the frequency domain and a long horizon of $$W_2^2(\mu,\nu) := \inf_{\gamma\in\Gamma(\mu,\nu)} \int \|x-y\|^2\, d\gamma(x,y),$$0 time steps. In those settings, TDPO significantly outperformed PPO, TRPO, DDPG, and TD3, a result attributed in the paper to low-variance deterministic gradients and the avoidance of noise that would excite resonance (Saleh et al., 2022).

The diffusion-model reinterpretation of maximum-entropy RL introduces DiffWPO, alongside DiffSAC and DiffPPO. On standard continuous-control benchmarks, the paper reports that DiffSAC, DiffPPO, and DiffWPO achieve better returns and higher sample efficiency than SAC and PPO, with performance improving as the number of diffusion steps increases (Sanokowski et al., 1 Dec 2025). This suggests that Wasserstein-based policy transport is compatible with expressive diffusion policies, although the mechanism is mediated through a reverse-KL upper bound and an augmented diffusion MDP rather than explicit OT solves.

6. Relations, misconceptions, limitations, and extensions

A common misconception is to treat WPO as synonymous with a single 2025 actor-critic algorithm. The literature instead uses the term in at least three senses: as a broad geometric paradigm for optimizing policies in probability space; as a specific continuous-action actor-critic algorithm based on projected Wasserstein gradient flow; and as a family of trust-region or proximal methods using optimal-transport discrepancies (Pfau et al., 1 May 2025, Terpin et al., 2022). This broader usage is already visible in early work that interpreted policy optimization itself as Wasserstein gradient flow (Zhang et al., 2018).

A second misconception is that WPO simply substitutes Wasserstein distance for KL while leaving the rest of policy optimization unchanged. In fact, different WPO variants change different parts of the optimization problem. OT-TRPO reformulates the trust-region constraint via a one-dimensional dual and an analytical push-forward update; deterministic WPO replaces KL by a quadratic Wasserstein surrogate that remains meaningful for Dirac policies; WPPG replaces KL projection by a Wasserstein proximal step plus a heat flow; and WNG replaces Fisher geometry by a Wasserstein information matrix on behavioral embeddings (Terpin et al., 2022, Saleh et al., 2022, Zhu et al., 3 Mar 2026, Moskovitz et al., 2020).

The relation between WPO and WPPG is particularly important. WPPG is closely related to WPO in spirit, because both optimize policies in Wasserstein geometry, but it is not a special case of the 2025 WPO method. The distinction is algorithmic and functional: Pfau-style WPO follows a Wasserstein gradient flow and then projects back to a parametric policy manifold via a KL projection step requiring tractable log densities, whereas WPPG performs a Wasserstein proximal step with operator splitting, avoids any KL projection, and never evaluates $$W_2^2(\mu,\nu) := \inf_{\gamma\in\Gamma(\mu,\nu)} \int \|x-y\|^2\, d\gamma(x,y),$$1, which makes it directly applicable to expressive implicit policies (Zhu et al., 3 Mar 2026).

The principal limitations of current WPO theory are assumption-heavy. Convergence results often require bounded rewards, bounded or regularized action sets, smoothness of $$W_2^2(\mu,\nu) := \inf_{\gamma\in\Gamma(\mu,\nu)} \int \|x-y\|^2\, d\gamma(x,y),$$2 or the transition model, uniform $$W_2^2(\mu,\nu) := \inf_{\gamma\in\Gamma(\mu,\nu)} \int \|x-y\|^2\, d\gamma(x,y),$$3 or log-Sobolev inequalities, exact or controlled-approximation critics, or finite action spaces for the strongest duality results (Zhu et al., 3 Mar 2026, Zhu et al., 25 May 2026, Song et al., 2023). High-dimensional exact Monge maps are generally intractable, so practical algorithms typically use sample-wise drifts $$W_2^2(\mu,\nu) := \inf_{\gamma\in\Gamma(\mu,\nu)} \int \|x-y\|^2\, d\gamma(x,y),$$4, entropic regularization, Sinkhorn-style approximations, or parameter-space projections rather than exact transport maps.

Beyond sequential RL, Wasserstein policy optimization has also been extended to contextual bandits and distributionally robust policy learning. In one line of work, max–min off-policy learning is carried out over Wasserstein balls around empirical context and cost distributions, with dual formulations, entropy-regularized surrogates, and finite-sample and iteration/query complexity guarantees (Shen et al., 2023). In another, individualized treatment rules are learned under source–target shift using pointwise $$W_2^2(\mu,\nu) := \inf_{\gamma\in\Gamma(\mu,\nu)} \int \|x-y\|^2\, d\gamma(x,y),$$5 ambiguity sets over conditional potential-outcome laws, yielding a tractable pessimistic translation of the source conditional mean response by $$W_2^2(\mu,\nu) := \inf_{\gamma\in\Gamma(\mu,\nu)} \int \|x-y\|^2\, d\gamma(x,y),$$6 (Kido, 2022). These are not MDP-style WPO algorithms, but they show that the Wasserstein-policy perspective extends naturally to robust decision-making under distribution shift.

In that enlarged sense, WPO is best viewed as a transport-geometric program for policy learning. Its unifying claim is not that all policies should be updated by the same numerical scheme, but that the underlying optimization object is a family of action distributions whose geometry matters. The literature to date suggests that this viewpoint is especially valuable when action-space geometry is consequential, policy densities are intractable, deterministic or near-deterministic policies are of interest, or entropy-regularized continuous control admits analysis through transport inequalities and gradient-flow methods.