# Policy Optimization as Wasserstein Gradient Flows

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

## Key facts

- Cited by 34 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 9 Aug 2018.

## Built on by (most cited first, 30 of 34)

- [Large-Scale Wasserstein Gradient Flows](https://mlsanity.com/paper/2106.00736/) (2021)
- [Recent Advances in Optimal Transport for Machine Learning](https://mlsanity.com/paper/2306.16156/) (2023)
- [Sqrt(d) Dimension Dependence of Langevin Monte Carlo](https://mlsanity.com/paper/2109.03839/) (2021)
- [Wasserstein Unsupervised Reinforcement Learning](https://mlsanity.com/paper/2110.07940/) (2021)
- [Wasserstein Policy Optimization](https://mlsanity.com/paper/2505.00663/) (2025)
- [Distributional Hamilton-Jacobi-Bellman Equations for Continuous-Time Reinforcement Learning](https://mlsanity.com/paper/2205.12184/) (2022)
- [Trust Region Policy Optimization with Optimal Transport Discrepancies: Duality and Algorithm for Continuous Actions](https://mlsanity.com/paper/2210.11137/) (2022)
- [Coin Sampling: Gradient-Based Bayesian Inference without Learning Rates](https://mlsanity.com/paper/2301.11294/) (2023)
- [Wasserstein Gradient Flows for Optimizing Gaussian Mixture Policies](https://mlsanity.com/paper/2305.10411/) (2023)
- [Provably Convergent Policy Optimization via Metric-aware Trust Region Methods](https://mlsanity.com/paper/2306.14133/) (2023)
- [On the Benefit of Optimal Transport for Curriculum Reinforcement Learning](https://mlsanity.com/paper/2309.14091/) (2023)
- [Optimal Transport-Assisted Risk-Sensitive Q-Learning](https://mlsanity.com/paper/2406.11774/) (2024)
- [Linear convergence of proximal descent schemes on the Wasserstein space](https://mlsanity.com/paper/2411.15067/) (2024)
- [Wasserstein Proximal Policy Gradient](https://mlsanity.com/paper/2603.02576/) (2026)
- [Nested-Wasserstein Self-Imitation Learning for Sequence Generation](https://mlsanity.com/paper/2001.06944/) (2020)
- [Batch Bayesian Optimization via Particle Gradient Flows](https://mlsanity.com/paper/2209.04722/) (2022)
- [Constrained Sliced Wasserstein Embedding](https://mlsanity.com/paper/2506.02203/) (2025)
- [Hessian-guided Perturbed Wasserstein Gradient Flows for Escaping Saddle Points](https://mlsanity.com/paper/2509.16974/) (2025)
- [Learn to Evolve: Self-supervised Neural JKO Operator for Wasserstein Gradient Flow](https://mlsanity.com/paper/2601.05583/) (2026)
- [Distributionally-Constrained Policy Optimization via Unbalanced Optimal Transport](https://mlsanity.com/paper/2102.07889/) (2021)
- [Generalized Talagrand Inequality for Sinkhorn Distance using Entropy Power Inequality](https://mlsanity.com/paper/2109.08430/) (2021)
- [Proximal Mean Field Learning in Shallow Neural Networks](https://mlsanity.com/paper/2210.13879/) (2022)
- [Wasserstein gradient flow for optimal probability measure decomposition](https://mlsanity.com/paper/2406.00914/) (2024)
- [Computational and Statistical Asymptotic Analysis of the JKO Scheme for Iterative Algorithms to update distributions](https://mlsanity.com/paper/2501.06408/) (2025)
- [A Distributed Algorithm for Measure-valued Optimization with Additive Objective](https://mlsanity.com/paper/2202.08930/) (2022)
- [Information-Theoretic Equivalence of Entropic Multi-Marginal Optimal Transport: A Theory for Multi-Agent Communication](https://mlsanity.com/paper/2208.10256/) (2022)
- [Wasserstein Diversity-Enriched Regularizer for Hierarchical Reinforcement Learning](https://mlsanity.com/paper/2308.00989/) (2023)
- [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)

## Abstract (arXiv)

Policy optimization is a core component of reinforcement learning (RL), and most existing RL methods directly optimize parameters of a policy based on maximizing the expected total reward, or its surrogate. Though often achieving encouraging empirical success, its underlying mathematical principle on em policy-distribution optimization is unclear. We place policy optimization into the space of probability measures, and interpret it as Wasserstein gradient flows. On the probability-measure space, under specified circumstances, policy optimization becomes a convex problem in terms of distribution optimization. To make optimization feasible, we develop efficient algorithms by numerically solving the corresponding discrete gradient flows. Our technique is applicable to several RL settings, and is related to many state-of-the-art policy-optimization algorithms. Empirical results verify the effectiveness of our framework, often obtaining better performance compared to related algorithms.