# On Wasserstein Reinforcement Learning and the Fokker-Planck equation

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

## Key facts

- Cited by 10 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 19 Dec 2017.

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

- [Large-Scale Wasserstein Gradient Flows](https://mlsanity.com/paper/2106.00736/) (2021)
- [Categorical SDEs with Simplex Diffusion](https://mlsanity.com/paper/2210.14784/) (2022)
- [Differentiable Trust Region Layers for Deep Reinforcement Learning](https://mlsanity.com/paper/2101.09207/) (2021)
- [Trust Region Policy Optimization with Optimal Transport Discrepancies: Duality and Algorithm for Continuous Actions](https://mlsanity.com/paper/2210.11137/) (2022)
- [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)
- [Wasserstein Adaptive Value Estimation for Actor-Critic Reinforcement Learning](https://mlsanity.com/paper/2501.10605/) (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)
- [Optimal Transport Meets Reinforcement Learning: A Survey](https://mlsanity.com/paper/2610.01413/) (2026)

## Abstract (arXiv)

Policy gradients methods often achieve better performance when the change in policy is limited to a small Kullback-Leibler divergence. We derive policy gradients where the change in policy is limited to a small Wasserstein distance (or trust region). This is done in the discrete and continuous multi-armed bandit settings with entropy regularisation. We show that in the small steps limit with respect to the Wasserstein distance W2, policy dynamics are governed by the Fokker-Planck (heat) equation, following the Jordan-Kinderlehrer-Otto result. This means that policies undergo diffusion and advection, concentrating near actions with high reward. This helps elucidate the nature of convergence in the probability matching setup, and provides justification for empirical practices such as Gaussian policy priors and additive gradient noise.