Wasserstein Gradient Flows for Optimizing GMM-based Policies

Hanna Ziesche, Leonel Rozo

Keywords: Optimal transport, policy optimization, gaussian mixture models, robot motion adaptation.

TL;DR: Policy structure-aware Optimization via Wasserstein Gradient Flows for Robot Motion Adaptation

Abstract:

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 $L^2$-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 over a set of 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.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics

Submission Guidelines: Yes

Please Choose The Closest Area That Your Submission Falls Into: Reinforcement Learning (eg, decision and control, planning, hierarchical RL, robotics)