Wasserstein Formulation of Reinforcement Learning. An Optimal Transport Perspective on Policy Optimization

Abstract: We present a geometric framework for Reinforcement Learning (RL) that views policies as maps into the Wasserstein space of action probabilities. First, we define a Riemannian structure induced by stationary distributions, proving its existence in a general context. We then define the tangent space of policies and characterize the geodesics, specifically addressing the measurability of vector fields mapped from the state space to the tangent space of probability measures over the action space. Next, we formulate a general RL optimization problem and construct a gradient flow using Otto's calculus. We compute the gradient and the Hessian of the energy, providing a formal second-order analysis. Finally, we illustrate the method with numerical examples for low-dimensional problems, computing the gradient directly from our theoretical formalism. For high-dimensional problems, we parameterize the policy using a neural network and optimize it based on an ergodic approximation of the cost.

• The paper introduces a novel geometric framework for policy optimization by modeling policies as measurable maps from state space to the 2-Wasserstein space.
• The methodology leverages Otto calculus to derive both gradients and Hessians of the RL objective within a rigorously defined Riemannian structure.
• Empirical results demonstrate scalability with neural parameterizations, achieving rapid convergence and robust control in nonlinear, high-dimensional tasks.

Wasserstein Formulation of Reinforcement Learning: An Optimal Transport Perspective on Policy Optimization

Overview

This paper presents a rigorous geometric framework for policy optimization in reinforcement learning (RL) based on the differential structure of optimal transport, specifically the Wasserstein space of action distributions (2604.14765). The authors formalize policies as measurable maps from the state space to the $2$-Wasserstein space of probability measures on the action space, endow this space with a Riemannian geometry induced by the stationary distribution, and derive both the gradient and Hessian of the expected long-term cost using Otto calculus. The work also bridges theoretical foundations (existence/uniqueness of invariant measures) with practical algorithms and demonstrates both analytical tractability and scalability via neural parameterizations with ergodic approximations.

Theoretical Contributions

Policy Space as a Fiber Bundle of Wasserstein Spaces

Policies are treated as Borel measurable functions $\pi: S \to P(A)$, where $S$ is the (Polish) state space and $P(A)$ the $2$-Wasserstein space over actions. The Wasserstein metric encodes not just distributional dissimilarity but also the geometry of the action space, bypassing the limitations of KL-based geometries that disregard action space topology.

Existence and Uniqueness of Invariant Measures

The development relies on the existence of a unique stationary state distribution $\mu_\pi$ induced by any policy $\pi$. The authors provide rigorous sufficient conditions:

• A contraction property on the dynamics and Lipschitz continuity of policy mappings is shown to guarantee strict contraction of the Markov operator in the $W_2$ metric, ensuring exponential convergence to a unique invariant measure.
• Alternative, more general Markov chain results (Krylov–Bogoliubov, Doeblin’s criterion, Strong/Weak Feller properties) are invoked to ensure existence and uniqueness when direct contraction does not apply.

Differential Structure and Metric

A formal Riemannian metric is introduced on the policy space $\Pi$:

$$\langle \xi, \eta \rangle_\pi = \int_{S} \langle v^\xi(s, \cdot), v^\eta(s, \cdot) \rangle_{L^2(\pi(s))} d\mu_\pi(s)$$

where $\pi: S \to P(A)$0 are tangent vector fields at each state, associated to infinitesimal transport plans in the action space.

Gradient and Hessian via Otto Calculus

The main formal machinery stems from gradient flows on Wasserstein spaces. The authors apply Otto calculus to derive:

• Gradient: The Wasserstein-gradient of the RL objective is

$\pi: S \to P(A)$1

where $\pi: S \to P(A)$2 is a soft Q-function built using the Poisson equation to combine costs and long-term value sensitivity. This gives a direct and theoretically principled generalization of policy gradients, embedding them in a Riemannian space (Fig. 1).

• Hessian: The paper provides a comprehensive second-order analysis, expressing the Hessian of the RL objective as a sum of a local action-space curvature term (the Hessian of $\pi: S \to P(A)$3) and a non-local transport/dynamics curvature term reflecting sensitivity of the stationary distribution to policy changes.

Figure 1: Convergence of the Average Cost during Policy Iteration.

A significant technical step is mapping variations of a neural network policy to infinitesimal Wasserstein transport plans (using the continuity equation), which justifies the use of “natural gradient” methods adapted to the metric tensor induced by the optimal transport geometry.

Geodesic Convexity

A sharp characterization is provided: in the special case where environment dynamics are independent of action (decoupled dynamics), the RL objective is geodesically convex under the Wasserstein metric if and only if the per-state cost is convex in the action measure. Otherwise, nontrivial transport/dynamics curvature generally introduces non-convexity.

Numerical Methods and Empirical Results

Model Classes and Algorithms

Three classes of environments were considered: a scalar nonlinear regulator, the inverted pendulum, and high-dimensional chains of coupled oscillators.

Two main numerical strategies are used:

• Exact/Particle Grid Policy Updates: In low-dimensional settings, policies are represented as mixtures of Diracs (“particles”) and the Bellman equations are solved exactly on grids for both value and invariant distribution, enabling direct implementation of the full Wasserstein gradient and Hessian.
• Scalable Neural Parameterization with Differentiable Physics or World Models: In high dimensions or with unknown dynamics, the policy is represented by a neural network. Either the gradients are backpropagated through known/differentiable dynamics (direct) or through a learned world model. Ergodic averages over long simulations are used to approximate the stationary distribution, and parameter updates are performed using Adam or an approximate natural gradient (aligned with the metric induced by Wasserstein geometry).

Visualization and Dynamics

Scalar Nonlinear Regulator

• Rapid convergence of average cost, robust to noise (Fig. 1 and Fig. 2).
• Learned value function displayed a parabolic landscape, and policy is approximately linear in the unsaturated regime.

  Figure 2: Simulated trajectory. Top: State evolution with and without noise. Bottom: Control actions applied by the agent.

Inverted Pendulum

• Grid-based policy recovers sharp switching (bang-bang) policies and correct value landscape (Fig. 3).
• Trajectory-based neural policies (with both Adam and natural gradient optimizers using differentiable physics) match grid-based performance, generating smooth policies that “pump” and stabilize the pendulum efficiently (Fig. 5, Fig. 6).
• With learned world models, no degradation in control quality was observed—the agent plans with similar efficiency once the world model fit converges (Fig. 7), as seen in aligned loss curves for policy and world model (Fig. 11).

  Figure 3: Grid-based results. Left: Value Function heatmap. Right: Policy mean action heatmap.

  

Figure 4: Direct Diff. Physics: Trajectory and Control.

Figure 5: Direct Diff. Physics: Trajectory and Control.

Figure 6: World Model: Generated Trajectory.

Figure 7: Joint Training Losses: Policy Loss (left) and World Model Prediction Loss (right).

High-Dimensional Oscillators

• For $\pi: S \to P(A)$4 oscillators, direct optimization using the Wasserstein-adapted natural gradient achieves stabilization of all masses efficiently, even in the presence of strong nonlinear couplings.
• Using learned world models, performance and stabilization are effectively maintained, and learning curves indicate rapid convergence for both the policy and the model.

  Figure 8: Direct Differentiable Physics: State trajectories (top) and Control inputs (bottom).

  

  Figure 9: Training Loss for Direct Differentiable Physics.

  

  Figure 10: World Model Approach: State trajectories (left) and Control inputs (right).

Discussion and Implications

This framework extends and unifies several active research themes:

• It encompasses classical RL approaches based on information geometry (NPG/TRPO) as limiting cases, but corrects their geometric inconsistency in continuous action spaces by exploiting the physical structure encoded in Wasserstein space.
• Second-order analysis is made tractable via explicit computation of gradients and Hessians in the Wasserstein policy manifold, enabling principled “Newton-like” policy update schemes.
• The paper formally ties the natural policy gradient (in parameter space) to the pullback of the optimal transport metric, justifying the link between function-space geometry and practical stochastic optimization methods.
• The use of ergodic approximations with neural policies scales the approach to high-dimensional applications without sacrificing the geometric rigor of the method.
• The analytical machinery developed (including measurable selection of tangent fields and vector measures) lays a foundation for further work on function-space regularization and policy interpolation in RL.

Conclusion

The Wasserstein-based geometric RL framework presented provides a rigorous foundation for policy optimization that explicitly respects the metric structure of the action space and the dynamical system. The approach supports both theoretical insight (e.g., on the nature of non-convexity and policy landscape curvature) and practical, scalable algorithms through the use of parameter pullbacks and neural approximations. Empirical results confirm that the geometric methodology leads to fast, robust control for both linear and highly nonlinear environments, with and without model knowledge. This work suggests that future developments in RL could exploit optimal transport geometry for both efficient optimization and theoretical guarantees around policy improvement and exploration, especially in structured or high-dimensional control tasks.