# Wasserstein Policy Optimization

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

## Key facts

- Cited by 13 papers in the map (data through 14 Oct 2026).
- Placed in RL theory and sample complexity, one of 15 directions in Deep reinforcement learning methods.
- Builds on 28 of its 53 references that are in the map.
- First posted on arXiv 1 May 2025.

## 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, 28 of 28)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Continuous control with deep reinforcement learning](https://mlsanity.com/paper/1509.02971/) (2015)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Trust Region Policy Optimization](https://mlsanity.com/paper/1502.05477/) (2015)
- [High-Dimensional Continuous Control Using Generalized Advantage Estimation](https://mlsanity.com/paper/1506.02438/) (2015)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [Fast and Accurate Deep Network Learning by Exponential Linear Units (ELUs)](https://mlsanity.com/paper/1511.07289/) (2015)
- [DeepMind Control Suite](https://mlsanity.com/paper/1801.00690/) (2018)
- [Sigmoid-Weighted Linear Units for Neural Network Function Approximation in Reinforcement Learning](https://mlsanity.com/paper/1702.03118/) (2017)
- [Optimizing Neural Networks with Kronecker-factored Approximate Curvature](https://mlsanity.com/paper/1503.05671/) (2015)
- [A Distributional Perspective on Reinforcement Learning](https://mlsanity.com/paper/1707.06887/) (2017)
- [dm_control: Software and Tasks for Continuous Control](https://mlsanity.com/paper/2006.12983/) (2020)
- [Maximum a Posteriori Policy Optimisation](https://mlsanity.com/paper/1806.06920/) (2018)
- [Distributed Distributional Deterministic Policy Gradients](https://mlsanity.com/paper/1804.08617/) (2018)
- [Learning Continuous Control Policies by Stochastic Value Gradients](https://mlsanity.com/paper/1510.09142/) (2015)
- [Revisiting Natural Gradient for Deep Networks](https://mlsanity.com/paper/1301.3584/) (2013)
- [Acme: A Research Framework for Distributed Reinforcement Learning](https://mlsanity.com/paper/2006.00979/) (2020)
- [Learning Agile Soccer Skills for a Bipedal Robot with Deep Reinforcement Learning](https://mlsanity.com/paper/2304.13653/) (2023)
- [V-MPO: On-Policy Maximum a Posteriori Policy Optimization for Discrete and Continuous Control](https://mlsanity.com/paper/1909.12238/) (2019)
- [SVGD as a kernelized Wasserstein gradient flow of the chi-squared divergence](https://mlsanity.com/paper/2006.02509/) (2020)
- [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)
- [Efficient Wasserstein Natural Gradients for Reinforcement Learning](https://mlsanity.com/paper/2010.05380/) (2020)
- [S2AC: Energy-Based Reinforcement Learning with Stein Soft Actor Critic](https://mlsanity.com/paper/2405.00987/) (2024)
- [Wasserstein Quantum Monte Carlo: A Novel Approach for Solving the Quantum Many-Body Schr"odinger Equation](https://mlsanity.com/paper/2307.07050/) (2023)
- [Towards practical reinforcement learning for tokamak magnetic control](https://mlsanity.com/paper/2307.11546/) (2023)

## Built on by (most cited first, 13 of 13)

- [Wasserstein Proximal Policy Gradient](https://mlsanity.com/paper/2603.02576/) (2026)
- [Diffusion-Augmented Markov Decision Processes for Maximum Entropy Reinforcement Learning](https://mlsanity.com/paper/2512.02019/) (2025)
- [Equivalence of stochastic and deterministic policy gradients](https://mlsanity.com/paper/2505.23244/) (2025)
- [Global Convergence of Wasserstein Policy Gradient for Entropy-Regularized Reinforcement Learning](https://mlsanity.com/paper/2605.26078/) (2026)
- [A Gradient Flow Perspective on Minimum MMD Estimation](https://mlsanity.com/paper/2607.03871/) (2026)
- [Beyond Fixed Morphologies: Learning Graph Policies with Trust Region Compensation in Variable Action Spaces](https://mlsanity.com/paper/2508.14102/) (2025)
- [Trust Regions Sell, But Who's Buying? Overlap Geometry as an Alternative Trust Region for Policy Optimization](https://mlsanity.com/paper/2602.06627/) (2026)
- [Wasserstein Formulation of Reinforcement Learning. An Optimal Transport Perspective on Policy Optimization](https://mlsanity.com/paper/2604.14765/) (2026)
- [Ratio-Variance Regularized Policy Optimization](https://mlsanity.com/paper/2605.26784/) (2026)
- [Policy Gradient for Continuous-Time Robust Markov Decision Processes](https://mlsanity.com/paper/2606.04335/) (2026)
- [Difference of Convex Programming in the Wasserstein Space with Applications to MMD Optimization](https://mlsanity.com/paper/2606.27767/) (2026)
- [Wasserstein Policy Gradient for Entropy-Regularized Linear-Quadratic Control](https://mlsanity.com/paper/2608.07433/) (2026)
- [Multi-step Proximal Policy Improvement in Offline Reinforcement Learning](https://mlsanity.com/paper/2609.03842/) (2026)

## Abstract (arXiv)

We introduce Wasserstein Policy Optimization (WPO), an actor-critic algorithm for reinforcement learning in continuous action spaces. WPO can be derived as an approximation to Wasserstein gradient flow over the space of all policies projected into a finite-dimensional parameter space (e.g., the weights of a neural network), leading to a simple and completely general closed-form update. The resulting algorithm combines many properties of deterministic and classic policy gradient methods. Like deterministic policy gradients, it exploits knowledge of the gradient of the action-value function with respect to the action. Like classic policy gradients, it can be applied to stochastic policies with arbitrary distributions over actions -- without using the reparameterization trick. We show results on the DeepMind Control Suite and a magnetic confinement fusion task which compare favorably with state-of-the-art continuous control methods.
