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

arXiv: https://arxiv.org/abs/2604.14765 · HTML: https://mlsanity.com/paper/2604.14765/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in RL theory and sample complexity, one of 15 directions in Deep reinforcement learning methods.
- Builds on 6 of its 23 references that are in the map.
- First posted on arXiv 16 Apr 2026.

## Placed in

- Deep reinforcement learning methods / [RL theory and sample complexity](https://mlsanity.com/deeprl/rl-theory-and-sample-complexity/)

## Builds on (most cited first, 6 of 6)

- [Trust Region Policy Optimization](https://mlsanity.com/paper/1502.05477/) (2015)
- [Lipschitz Continuity in Model-based Reinforcement Learning](https://mlsanity.com/paper/1804.07193/) (2018)
- [Policy Optimization as Wasserstein Gradient Flows](https://mlsanity.com/paper/1808.03030/) (2018)
- [Learning to Score Behaviors for Guided Policy Optimization](https://mlsanity.com/paper/1906.04349/) (2019)
- [Wasserstein Policy Optimization](https://mlsanity.com/paper/2505.00663/) (2025)
- [Wasserstein Gradient Flows for Optimizing Gaussian Mixture Policies](https://mlsanity.com/paper/2305.10411/) (2023)

## Abstract (arXiv)

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, providing rigorous guarantees for their existence in a general context. We then define the tangent space of policies and characterize the geodesics, specifically addressing the measurability of vector fields mapping 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 (i.e., the expected cumulative cost), providing a formal second-order analysis. Finally, we illustrate the method with numerical examples: we compute the gradient directly from our theoretical formalism for low-dimensional problems, and we illustrate the numerical applicability of our formal framework to high-dimensional continuous control by parameterizing the policy with a neural network optimized via an ergodic approximation of the cost.