Wasserstein Gradient Flows for Optimizing Gaussian Mixture Policies
Hanna Ziesche Leonel Rozo
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 \emph{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 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.