# Wasserstein Gradient Flows for Optimizing Gaussian Mixture Policies

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

## Key facts

- Cited by 8 papers in the map (data through 14 Oct 2026).
- Builds on 37 of its 77 references that are in the map.
- First posted on arXiv 17 May 2023.

## Builds on (most cited first, 30 of 37)

- [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)
- [Trust Region Policy Optimization](https://mlsanity.com/paper/1502.05477/) (2015)
- [Optuna: A Next-generation Hyperparameter Optimization Framework](https://mlsanity.com/paper/1907.10902/) (2019)
- [Deep Reinforcement Learning for Robotic Manipulation with Asynchronous Off-Policy Updates](https://mlsanity.com/paper/1610.00633/) (2016)
- [Fast Computation of Wasserstein Barycenters](https://mlsanity.com/paper/1310.4375/) (2013)
- [How to Train Your Robot with Deep Reinforcement Learning; Lessons We've Learned](https://mlsanity.com/paper/2102.02915/) (2021)
- [Maximum Entropy RL (Provably) Solves Some Robust RL Problems](https://mlsanity.com/paper/2103.06257/) (2021)
- [Optimal Entropy-Transport problems and a new Hellinger-Kantorovich distance between positive measures](https://mlsanity.com/paper/1508.07941/) (2015)
- [Pymanopt: A Python Toolbox for Optimization on Manifolds using Automatic Differentiation](https://mlsanity.com/paper/1603.03236/) (2016)
- [Large-Scale Wasserstein Gradient Flows](https://mlsanity.com/paper/2106.00736/) (2021)
- [A Wasserstein-type distance in the space of Gaussian Mixture Models](https://mlsanity.com/paper/1907.05254/) (2019)
- [Neural Dynamic Policies for End-to-End Sensorimotor Learning](https://mlsanity.com/paper/2012.02788/) (2020)
- [Information asymmetry in KL-regularized RL](https://mlsanity.com/paper/1905.01240/) (2019)
- [MushroomRL: Simplifying Reinforcement Learning Research](https://mlsanity.com/paper/2001.01102/) (2020)
- [Optimal transport for Gaussian mixture models](https://mlsanity.com/paper/1710.07876/) (2017)
- [Kernelized Movement Primitives](https://mlsanity.com/paper/1708.08638/) (2017)
- [A survey on policy search algorithms for learning robot controllers in a handful of trials](https://mlsanity.com/paper/1807.02303/) (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)
- [On Riemannian Optimization over Positive Definite Matrices with the Bures-Wasserstein Geometry](https://mlsanity.com/paper/2106.00286/) (2021)
- [Handling the Positive-Definite Constraint in the Bayesian Learning Rule](https://mlsanity.com/paper/2002.10060/) (2020)
- [Probabilistic Mixture-of-Experts for Efficient Deep Reinforcement Learning](https://mlsanity.com/paper/2104.09122/) (2021)
- [Learning Gaussian Mixtures Using the Wasserstein-Fisher-Rao Gradient Flow](https://mlsanity.com/paper/2301.01766/) (2023)
- [Differentiable Trust Region Layers for Deep Reinforcement Learning](https://mlsanity.com/paper/2101.09207/) (2021)
- [Bayesian Optimization Meets Riemannian Manifolds in Robot Learning](https://mlsanity.com/paper/1910.04998/) (2019)
- [Efficient Wasserstein Natural Gradients for Reinforcement Learning](https://mlsanity.com/paper/2010.05380/) (2020)
- [Optimal transport natural gradient for statistical manifolds with continuous sample space](https://mlsanity.com/paper/1805.08380/) (2018)
- [Trust-Region Variational Inference with Gaussian Mixture Models](https://mlsanity.com/paper/1907.04710/) (2019)
- [Rethinking Initialization of the Sinkhorn Algorithm](https://mlsanity.com/paper/2206.07630/) (2022)

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

- [SAC Flow: Sample-Efficient Reinforcement Learning of Flow-Based Policies via Velocity-Reparameterized Sequential Modeling](https://mlsanity.com/paper/2509.25756/) (2025)
- [The Art of Imitation: Learning Long-Horizon Manipulation Tasks from Few Demonstrations](https://mlsanity.com/paper/2407.13432/) (2024)
- [Slicing the Gaussian Mixture Wasserstein Distance](https://mlsanity.com/paper/2504.08544/) (2025)
- [Slicing Wasserstein Over Wasserstein Via Functional Optimal Transport](https://mlsanity.com/paper/2509.22138/) (2025)
- [DDEQs: Distributional Deep Equilibrium Models through Wasserstein Gradient Flows](https://mlsanity.com/paper/2503.01140/) (2025)
- [Robot Skill Generalization via Keypoint Integrated Soft Actor-Critic Gaussian Mixture Models](https://mlsanity.com/paper/2310.15059/) (2023)
- [Beyond Fixed Morphologies: Learning Graph Policies with Trust Region Compensation in Variable Action Spaces](https://mlsanity.com/paper/2508.14102/) (2025)
- [Wasserstein Formulation of Reinforcement Learning. An Optimal Transport Perspective on Policy Optimization](https://mlsanity.com/paper/2604.14765/) (2026)

## Abstract (arXiv)

Robots often rely on a repertoire of previously-learned motion policies for performing tasks of diverse complexities. When facing unseen task conditions or when new task requirements arise, robots must adapt their motion policies accordingly. In this context, policy optimization is the de facto paradigm to adapt robot policies as a function of task-specific objectives. Most commonly-used motion policies carry particular structures that are often overlooked in policy optimization algorithms. We instead propose to leverage the structure of probabilistic policies by casting the policy optimization as an optimal transport problem. Specifically, we focus on robot motion policies that build on Gaussian mixture models (GMMs) and formulate the policy optimization as a Wassertein gradient flow over the GMMs space. This naturally allows us to constrain the policy updates via the L2-Wasserstein distance between GMMs to enhance the stability of the policy optimization process. Furthermore, we leverage the geometry of the Bures-Wasserstein manifold to optimize the Gaussian distributions of the GMM policy via Riemannian optimization. We evaluate our approach on common robotic settings: Reaching motions, collision-avoidance behaviors, and multi-goal tasks. Our results show that our method outperforms common policy optimization baselines in terms of task success rate and low-variance solutions.