Wasserstein Gradient Flows for Optimizing Gaussian Mixture Policies
Authors: Hanna Ziesche
Leonel Rozo
Venue: NeurIPS 2023 poster (2023)
Keywords: Policy optimization
robot learning
reinforcement learning
Gaussian mixture models
optimal transport
robotics
TL;DR: We leverage the structure of Gaussian mixture policies to cast the policy optimization as a Wasserstein gradient flow over the GMMs space.

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 \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.

Wasserstein Gradient Flows for Optimizing Gaussian Mixture Policies
Hanna Ziesche" and Leonel Rozo"
Bosch Center for Artificial Intelligence (BCAI) Renningen, Germany
name surname@de. www.eedwdeeeb....eooo...oooo.cee com
Abstract
Robots often rely on a repertoire of previously-learned motion policies for perform- ing tasks of diverse complexities. When facing unseen task conditions when new task requirements arise, robots must adapt their motion policies accordingly. In this context, policy optimization is the de facto paradigm adapt robot policies as function of task-specific objectives. Most commonly-used motion policies carry particular structures that are often overlooked in policy optimization algo- rithms. 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 Wasserstein gradient flow over the GMMs space. This naturally allows us to constrain the policy updates via the L -a-sassesteein distance between GMMs to enhance the stability of the policy opti- mization process. Furthermore, manifold to optimize the Gaussian distributions of the GMM policy via Riemannian optimization. We evaluate approach common robotic settings: Reaching motions, collision-avoidance behaviors, and multi-goal tasks. Our results show that method outperforms common policy optimization baselines in terms of task success rate and low-variance solutions.
leverage the geometry of the Bures Wasserstein
Introduction
One of the main premises about autonomous robots is their ability to successfully perform a large range of tasks in unstructured environments. This demands robots to adapt their task models according to environment changes or new task requirements, and consequently to adjust their actions to successfully perform under unseen conditions (.. In general, robotic tasks, such as picking or inserting an object, are usually executed by composing previously-leamee skills [2]. each represented by a motion policy. Therefore, order to successfully perform under new settings, the robot should adapt its motion policies according to the new task requirements and environment conditions.
Research on methods for robot motion policy adaptation is vast 33,4, with approaches mainly building on black- box optimizers [], end-to-end deep reinforcement learning 8), and policy search [9] Regardless of the optimization method, most approaches disregard the intrinsic policy structure in the adaptation strategy. However, several motion policy models (e.g., dynamic movement primitives (DMP) Gaussian mixture models (GMM) [T]. probabilistic movement primitives (ProMPs) [12 13), and neural networks [14), among others), carry specific physical or probabilistic structures that starting learning phase (2). thus the policy structure already encapsulates relevant prior information about the skill. Second, structure-unaware adaptation strategies optimize the policy parameters disregarding
should not be ignored. First, these policy models are often learned from demonstrations
Equal contribution.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

M GMMJ
Adapted GMM policy
Initial GMM policy
New target
Figure 1: Illustration our policy structure- aware adaptation of GMM policies. Left: A robot manipulator tracks a reference trajectory ( (-) ) extracted from a GMM policy being required to reach for a new target (*+). Center: Our policy optimization provides policy updates that follow a Wasserstein gradient flow on the manifold of GMM policies GMM,. Righr: The final GMM policy where some Gaussian components were adapted in order to retricve a new reference trajectory that allows the robot reach for the new target.
initially learned from demonstrations, while
the intrinsic structural characteristics of the policy model (e.g., a DMP represents a second-order dynamical system). In this context, we hypothesize that the policy structure may be leveraged to better control the adaptation strategy via policy structure-aware gradients and trust regions.
Our main idea is to design policy optimization strategy that explicitly builds on particular policy structure. Specifically, we focus on GMM policies, which have been widely used to learn motion skills from human demonstrations [T] 15 [6] [7] [8] [9]. GMMs provide simple but expressive enough
representation for learning a large variety of robot skills such as: Stable dynamic motions (20) 21
collaborative behaviors 23 24 assembly strategies 25 26 and contact-rich manipulation 27
among others. Often, skills learned from demonstrations need to be refined - due to imperfect data or adapted to comply with new task requirements. Existing adaptation strategies for GMM policies either build a kernel method on top of the original skill model [29]. or leverage reinforcement learning (RL) adapt the policy itself [30 31). However, none of these techniques explicitly considered the intrinsic probabilistic structure of the Gaussian mixture policy,
Unlike the aforementioned approaches, we propose policy optimization technique that explicitly considers the underlying GMM structure. To do so, we exploit optimal transport theory [2] [3], which allows us view the set of GMM policies GMM,- Specifically, we leverage the idea of Chen et al.
particular space of probability distributions
(34] and Delon and Desolneux (35 to view
GMM as a set of discrete measures (Dirac masses) on the space of Gaussian distributions G (RR4), which is endowed with a Wasserstein distance (see $
2 This allows us to formulate the policy
optimization as a Wasserstein gradient flow (WGF) over the space of GMMs (as illustrated in Fig. [ and explained in $ where the policy updates are naturally guaranteed to be GMMs. Moreover, we take advantage of the geometry of the Bures-Wasseree Wasserstein manifold optimize the Gaussian distributions of a GMM policy via Riemannian optimization. We evaluate our approach over a set of different GMM policies featuring common robot skills: Reaching motions, collision-avoidame behaviors, and multi-goal tasks (see $ Our results show that our method outperforms common policy optimization baselines in terms of task success rate while providing low-variance solutions.
Related Work: Richemond and Maginnis [36] pioneered the idea of understanding policy opti- mization through the lens of optimal transport. They interpreted the policy iteration as gradient flows by leveraging the implicit Euler scheme under Wasserstein distance (see $ p considering only 1-step return settings. They observed that the resulting policy optimization resembles the gradient flow of the Fokker-Plance equation (JKO scheme) [37] In a similar spirit, Zhang et al. [8 proposed to use WGFs to formulate policy optimization as sequence of policy updates traveling along gradient flow on the space of probability distributions until convergence. To solve the WGF problem, the authors proposed particle-based algorithm to approximate continuous density functions and subsequently derived the gradients for particle updates based on the JKO scheme. Although Zhang al. (38) considered general parametric policies, their method assumed a distribution over the policy parameters and did not consider specific policy structure, which partially motivated their particle- based approximation. Recently, Mokrov et al. [39] tackled the computational burden of particle methods by leveraging input-convex neural networks to approximate the WGFs computation. They reformulated the well-known JKO optimization 37] over probability measures by optimization over convex functions. Yet, this work remains a general solution for WGF computation and it did not address its use for policy optimization problems.
Aside from optimal transport approaches, Arenz et al. 30 proposed a trust-region variational inference for GMMs to approximate multimodal distributions. Although not originally designed

---

for policy optimization, the authors elucidated a connection to learn GMMs of policy parameters in black-box RL. However, their method cannot directly be applied to our GMM policy adaptation setting, nor does it consider the GMM structure from an optimal transport perspective. Nematollahi et al. [3 proposed SAC-GMM, a hybrid model that employs the well-known SAC algorithm [+]] to refine dynamic skills encoded by GMMs. The SAC policy was designed to learn residuals on single vectorized stack of GMM parameters, thus fully disregarding the GMM structure and the geometric constraints of its parameters. Ren et al. 41 introduced PMOE, method to train deep RL policies using a probabilistic mixture of experts via GMMs. They addressed the con-differentiabiiity problem caused by the optimization of categorical distribution parameters associated with the mixture weights via gradient estimators. However, PMOE does not take an optimal transport perspective on the policy optimization problem, nor does it account for the geometry arising from the GMM parameters.
Concerning approaches that explicitly account for the geometry induced by the policy, we may broadly categorize them into two groups: exact models and natural gradient methods. The former category refers to policy optimization methods that fully exploit the geometric structure of the policy, as proposed in this paper. In this case, the state of the art is scarce but recent works such as the Riemannian proximal policy optimization for GMMs proposed in [42], showed the importance of leveraging the geometry induced by the GMM parameters in the policy optimization via Riemannian gradients, similarly to our method. Their policy optimization was regularized by a Wasserstein distance to control the exploration-exple exploitation off. However, their method did not formulate the policy optimization as an optimal transport problem, i.c., the policy updates did not follow Wasserstein gradient flow, in our approach. but it employed instead classical non-convex optimization procedure.
The second category encloses approaches that leverage the well-established technique for enhancing the stability of policy optimization, namely, introducing regularization to the objective function using the Kullback-Leibler (KL) divergence, which induces the so-called natural gradient 443 44 45 46).
In this context, the regularization provides a global measure of the policies similarity, while the natural gradient perspective aims at locally controlling the changes between successive policy updates. In other words, natural gradient approaches leverage the local geometry of the policy parameters space to compute controlled updates. Note that policies similarity can also be quantified via the Wasserstein distance, as proposed in this paper and in recent works [36 38 47 48). Unlike the KL divergence, the Wasserstein distance enjoys powerful geometrical, computational, and statistical For example, Moskovitz et al. (48) recently employed the Wasserstein natural gradient to exploit the local geometry induced by the Wasserstein regularization of behavioral policy optimization (47) but neither the policy nor the behavioral embeddings had particular structures that could be leveraged in the policy optimization.
features [32 33
In this paper we exploit the geometric properties arising from both a Wasserstei--egullarizee objective and a Gaussian mixture policy. This allows us not only to see the policy iteration Wasserstein gradient flow, as in (36), but also to leverage the Riemannian geometry associated to the GMM space GMM, to formulate exact Riemannian gradient descent updates [49]. instead of relying on inexact natural gradients approximations. To achieve this, our method leverages the geometry induced by the structure of the space of GMM policies via the Bures-Wassersse which naturally guarantees that policy updates stay on GMM,. To the best of our knowledge, our optimization of GMM policies based on Wasserstein gradient flows and the Bures-Wasserstein manifold is the first of its kind in the domain of policy adaptation.
manifold,
2 Background
2.1
Wasserstein gradient flows
In Euclidean space a gradient flow is smooth curve x + Rd that satisfies the partial differential equation (PDE) ((t) - (ax(()) for given loss function L Rd - IR and starting point at t 32 50) solution can be found straightforwardly by forward discretization, leading to the well-known explicit Euler update scheme xx+11 _ AV where À denotes the learning rate and xt indicates discretization of the curve ætt) with discretization parameter т. Alternatively,
we can use a backward discretization, which leads to the following implicit Euler scheme
||0 - - + I(x))
: arg min
(1)
2T
x
3

---

is sometimes referred to as Minimizing Movement Scheme and can be used as an alternative characterization of a gradient flow.
This characterization is particularly interesting when we need to extend the concept of gradient flows Non--uuclddea)) general metric setrings, since there is no notion of V V. in these cases 32 [1]. Eq. does not involve any gradients and can be expressed using only metric quantities. In this work, we are particularly interested in gradient flows in the L [F-aasserstein space, defined as the probability measures P((²) on a separable Banach space X' [22] and endowed with the L2-Wasserstein distance W2 defined as
inf / (tEll(u,)) Xxxx || - x2l x=11(d(((11222))
$
W2(,,
where μr, E P(()) and 7 P(A2) is defined to have the two marginals / and v.
Generalized Minimizing Movement scheme characterizing gradient flows in the Wasserstein space can be written in analogy to Eq. ass:
(T, 2T
+ MiT))
T arg min
where L is functional be minimized on the Wasserstein space and m. P(X). In the following, we will omit the superscript t for notational convenience.
2.2 Reinforcement Learning Wasserstein Gradient Flows
Our view of the policy structure-awar optimization builds on the approach outlined [36], which in turn is based on the JKO scheme [37. They proposed a formulation of 1-step RL problems in terms of Wasserstein gradient flows. In particular, they studied the evolution a policy π under the influence of free energy functional J of the form:
((T) (T))+ (H(()) = / dr(als)r(s,a) - B dr(als) log(t(als)),
(4)
where K2 (T) denotes the inner energy of the system, here determined by the reward r(s, a), and H(TT) the entropy of the policy π [a]s), with & and a denoting the state and action, respectively. Thus, Eq. caa be recognized the usual objective in 1-step RL settings with entropy regularization. is well known that the evolution of probability densities under a free energy of this form is properly described by PDE known as the Pokker-Plance equation. Richemond and Maginnis (36) exploited the result of Jordan et al. (37), which stated that this evolution can be interpreted as the gradient flow of the functional J in Wasserstein space. This flow is characterized by the following minimizing movement scheme
(T, ((T) 2r
Th+1 arg min
(5)
which naturally provides iterative updates for the policy æ. While Richemond and Maginnis [36 considered 1-step bandit setting, we extend this approach to full multi-step RL problems and consequently learn policies for long-horizon tasks.
2.3 The I Wasserstein distance between Gaussian Mixture Models (GMMs)
We consider policies π (œ)) that build on GMM structure, i.e., T()) E w.N(;;; 14i E,), where N denotes multivariate Gaussian distribution with mean μı; and covariance matrix Σ,, and wi are the weights of the N individual Gaussian components, which are subject to Σ; wi = 1. In the following, we will write â, & and @ to denote the stacked means, covariance matrices and weights of the N components. Therefore, we do not consider WGFs on the full manifold of probability distributions (Wasserstein space) ((Rd) but rather focus WGFs evolving the submanifold of GMMs, that is we can approximately describe this submanifold as a discrete distribution over the space of Gaussian distributions equipped with the Wasserstein metric. This in turn can be identified with the Bures- Wasserstein manifold which is the product manifold gd x S
GMM, C P(Rd). Following 34 35

---

where S denotes the Riemannian manifold of d-dimensional symmetric positive definite matrices. The corresponding approximated Wasserstein distance between two GMMs T1, T22 given by
N Σ 2
min (2), PEU (W2)
Paj (W Ng),
(6)
where (w1, w2) = P Ply W1. P1NN w = w2} with 1y denoting an N N-dimensional
vector of ones. The Wasserstein distance between two Gaussian distributions in Eq. oaan be computed analytically follows
Σ; + Σn - -2(2/ Val
(W E.). E,)) pi 11ll2 +
(7)
2.4 Learning GMM policies from demonstrations
popular approach in RL- particularly robotics - to reduce the number of policy rollouts the environment to warm-start the policy with of demonstrations provided by an expert. In this work we choose to represent our policy via a GMM. We assume that demonstrations are provided as a set of trajectories r of state-action pairs {S0, an); (81, a1); ... (ST, aT)). ) To initialize our pol- icy. we first use the Expectation-Maxine Maximization (EM) algorithm to fit GMM, in the state-action space, :4,
to the demonstrations. This results in mixture distribution π ((s, a) Σn w,N[[ss ([sa] (4, E,)
from which a policy can be obtained by conditioning on the state 8, as follows
IT (
म. (a 8)
(8)
x. (8, a)da
In the context of GMMs, this is also known as Gaussian Mixture Regression (GMR) 53 The resulting conditional distribution is another GMM action space with state-dependent parameters,
=aa()))) Σ (ss)))(aa))) Ei).
Details on computation of Eq. 9 from the original GMM are given in App. A.I
3 Wasserstein Gradient Flows for GMM Policy Optimization
In this work, we foeus on multi-step RL tasks for policy adaptation. We consider a finite-horizon Markov Decision Process (MDP) with continuous state and action spaces S E R and A E Rm transition and reward functions p(s [s,, a2) and r(s,, az), initial state distribution p(sp) and discount factor y. Further, we assume to have an initial policy π (a, s,), which is to be adapted by optimizing some objective function (z).). As stated in $ T this problem arises robot learning settings where a policy learned via imitation learning (e.g.. LTD) needs to be adapted to new objectives or unseen environmental conditions. To promote exploration and avoid premature convergence to suboptimal policies, we leverage maximum entropy RL (54) by adding entropy term H(T)) to the objective. Thus, the overall objective has the form of a free energy functional (resembling Eq. and can be written as
J(T) K,( + (T)(),,
(10)
where B is hyperparameter and K,(mm) corresponds to the usual cumulative return
K,(rr) FE,
Er(s,, rs,, [ dedgagddddddoolaassssass......ed
Σ
an) Eor(sn,,)) (11)
The evolution of the policy (a, (st) over the course of the optimization can be described as a flow of a probability distribution in Wasserstein space. This formulation comes with three major benefits: (i) We directly leverage the Wasserstein metric properties for describing the evolution of probability distributions; (ii) We exploit the T.-aasserntein distance to constrain the policy updates, which
important to guarantec stability in policy optimization [55 45 56); (ii) By constraining specific
5

---

submanifolds of the Wasserstein space, in this case GMMs, we can impose additional structural properties on the policy optimization.
Since our objective in Eq. T0 has the form of the free energy functional studied by 36 37 we can leverage the iterative updates scheme of Eq. t formulate the evolution of our policy iteration under the flow generated by Eq. [Q] As mentioned previously, we focus on the special case of GMM policies and therefore restrict the Wasserstein gradient flow the submanifold of GMM distributions GMM,. We refer the interested reader to App.
A.3 where we provide the explicit form of (T) of
Eq. 10 for the GMM case.
3.1 Policy optimization
To begin with, we leverage the approximation that describes the GMM submanifold as discrete distribution over the space Gaussian distributions G(R"), equipped with the Wasserstein metric [34] Consequently, our policy optimization problem naturally splits into an optimization over the
K5
(N 1-1) -ii-eennnsio simplex and an optimization on the N-fold product of the Bures- Wasserstein
manifold (BW*), i.e., the product manifold (R& S + The former corresponds to the GMM weights while the latter applies to the set of Gaussian distributions' parameters. Note that the identification with the BW manifold allows us perform the optimization directly on the parameter space. This comes with several benefits: (i) We can leverage the well-known analytic solution of the Wasserstein distance between Gaussian distributions in Eq. 6 greatly reducing the computational complexity of the policy optimization. (ii) As shown in (4), we can guarantee that the policy optimized via Eq. Øeemmins GMM (i.e., satisfies the mass conservation constraint). (iii) Unlike the full Wasserstein S the resulting product manifold is true Riemannian manifold such that we can leverage the machinery Riemannian optimization. Importantly, working in the parameter space allows us to apply an explicit Euler scheme, instead of the implicit formulation of Eq. when optimizing the Gaussian distributions' parameters.
According to the above splitting scheme, we formulate the policy optimization as a two-step procedure that alternates between the Gaussian parameters (i.e. means and covariance matrices) and the GMM weights. To optimize the former, we propose to leverage the Riemannian structure of the BW manifold to reformulate the updates embedding the Gaussian components of the GMM policy in a Riemannian manifold, the Wasserstein gradient flow the implicit form of Eq. can be approximated by an explicit Euler update scheme according to the BW metric (further details are provided in App. A.4). This allows us to leverage the expressions of the Riemannian gradient and exponential map of the BW manifold [59 60 Thus, the optimization boils down to Riemannian gradient descent where the gradient is defined w.r.t the Bures- Wasserstein metrie. particular, we use the expression for Riemannian gradient, metric and exponential map used in [60]. Formally, the resulting updates for the Gaussian parameters of the GMM follow the Riemannian gradient descent scheme given by:
forward discretization, similarly to (58] In other words, by
f Ra, (X grad, J(TK)). and Σt41 Ry, A.grad ((T)))
(12)
where grad denotes the Riemannian gradient w.r.t. the Bures Wasserstein metric, R TaM + M denotes the retraction operator, which maps a point on the tangent space T back to the manifold
M = BW [61] Moreover, À is a learning rate and mı def (fAc, Σt, w..). The Euclidean gradients of (T)) required for computing grad can be obtained using a likelihood ratio estimator (a.k.a score function estimator or REINFORCE) [22] and are provided in A.p
Concerning the GMM weights, we first reparameterize them as Wj exp and optimize j = 1...,, which unlike à are unconstrained. For this optimization we employ the implicit Euler scheme:
w.r.t. the new parameters .
W2 TR). ((TT+1
1 arg min
(13)
27
*Wasserstein space is not true Riemannian manifold, but it can be equipped with a Riemannian structure and formal calculus this manifold [57]. which has been made rigorous in [5]]

---

Algorithm GMM Policy Optimization via Wasserstein Gradient Flows
Input: initial policy T(ass)
1: while not goal reached do
2: Rollout policy π (a]s) in the environment for M episodes to collect interaction trajectories
T = {(50, do, ro). ($1, 01 T1). (8T, Qr, FT) rr)mm=
3: repeat
4: Update Gaussian components parameters û. Σ using Riemannian optimization 12 where
X is determined via line-search (see (32)).
5: until convergence
6: repeat
7: Update GMM weights @ via gradient descent the free energy objective 10, using 14
8:
until converged
9:
end while
where ( Tppt+1 Ex+1, 1). We minimize Eq.
13 by gradient descent w.r.t. 17 as follows:
W W2 TQ)
MR+1 7 AV
((TK+
(14)
The gradient of J(x) can be obtained analytically using likelihood ratio estimator. For the Wasser- stein term, we first compute the gradient w.r.t. the weights via the Sinkhorn algorithm [3]. from which the gradient w.r.t 77 can be then obtained via the chain rule. Note that we have to rely on the Sinkhorn algorithm here since there is analytic solution available for the Wasserstein distance between discrete distributions, unlike the above case of the Gaussian components. Consequently, we cannot compute the corresponding gradients.
3.2 Implementation Pipeline
To carry out the policy optimization, we proceed as in the usual on-policy RL scheme: We first roll out the current policy to collect samples of state-action-reward tuples. Then, interaction trajectories to compute sample- based estimate of the functional Ky- and its gradients w.r.t the policy parameters, explained in $3.1 An optimization step consists of alternating between optimizing the Gaussian parameters using Eq. 12 and updating the weights via Eq. 14 For the optimization of the Gaussian parameters we leverage Pymanopt [64] for Riemannian optimization. We extended this library by implementing the Mures-Wasserster manifold based on the expressions provided by Han et al. [60] (see App. A.2 for details). Furthermore, we added custom line-search routine that accounts for the constraint on the Wasserstein distance between the old and the optimized GMM, to our knowledge such a search method does not exist in out-of-the-box optimizers. The details of this custom tine-search are given in Algorithm 2 App. A.5 Regarding the optimization of the GMM weights, we use POT 655 Python library for optimal transport, from which we obtain the quantities required to compute the gradients of the Wasserstein distance w.r.t. the weights in Eq. [4]
the collected
The full policy optimization finishes if either the objective stops improving or the Wasserstein distance between the old and optimized GMMs exceeds predefined threshold, which is chosen experimentally. Afterwards, fresh rollouts are performed with the updated policy and the aforementioned two-step alternating procedure starts over. This optimization loop is repeated until task-specific success criterion has been fulfilled. We summarize the proposed optimization in Algorithm
Experiments
We tested our approach in three different robotic settings: reaching skill, a collision-free trajectory tracking, and multiple-god -ooall task. All these tasks were represented in 2D operational space. More- over, we tested our approach on two additional versions of the collision-free trajectory tracking task that are formulated for 3D operational- space and 7D joint-space representations of the end-effector state and action. These two latter tasks were performed by a 7-DoF Franka Emika Panda robot in a virtual environment as reported in App. A.6.3 All robot motion policies were initially learned from human demonstrations collected on a Python GUI. We assumed we were given M demon- strations, each of which contained T,,, data points for datasct of T =mm T, total observations
7

---

T= [(s,, a,))?= The state and action a corresponded to the robot end-effector position E R and velocity 4 E TRd, with d being the operational space dimensionality. The GMM models were initially trained via classical Expectation-Maxiizzatiio
The policy rollout consisted of sampling a velocity action a₂ ~ T [aa 8x) using Eq. 9 and subsequently commanding the robot via a Cartesian velocity controller at a frequency of 100Hz. For all the experiments, we used the Roboties Toolbox for Python 666] to simulate the robotic environments.
To show the importance of accounting for the policy structure in RL settings, compared our method against two structure- unaware baselines: PPO 45] and SAC-GMM [I]. Our baselines choice is motivated by the fact that these two RL methods are still widely used and fairly competitive [67], even in real robotic scenarios [68]. We also considered an additional baseline: PMOE [41] which is method train deep RL policies using probabilistic mixture of experts via GMMs. As PPO was not originally designed to directly optimize the parameters of previously-learnee GMM policy, we designed the policy actions to represent (usually small) corrections the GMM parameters, i.e.
(Aw Aù A vec(E) [,), following the same methodology as SAC-GMM [31. The PPO and SAC implementations correspond to the code provided by 3able-Baselines3 69 whose policies are
parametrized by MI.P networks. The PMOF implementation corresponds to the code given in the original paper [1). During policy optimization, we sampled an action from the MLP policy that was then used to update the GMM parameters by adding the computed corrections to the current parameters. Later, we proceeded as described earlier, namely, the updated GMM was used to compute the velocity action via Eq. [ For comparison purposes, we report statistical results for all the settings over 5 runs for the task success rate and solution variance. We tuned the baselines separately for each task using Optuna [70). In addition, to assess the importance of our Riemannian formulation, we performed an ablation where we used the implicit scheme based on Euclidean gradient descent instead of the explicit optimization on the Bures- Wasserstein manifold (see App. A...2)..
4.1 Tasks description
This experiment consists of: (1) learning an initial GMM policy such that the robot end-effector reaches a target by following an L-shape trajectory from its initial position, and (2) adapting the GMM policy to reach new target located midway and above the previously-leamee L-shape trajectories. The initial policy, shown in Fig. left and Fig. [3] was learned from 12 demonstrations and encoded by 7-component GMM. To adapt the policy, we defined a dense reward as a function of the position error between the robot end-effector and the new target. We also added sparse penalty term that punishes rollouts leading significantly divergent trajectories. Each optimization episode comprised 10 rollouts, each of maximum horizon length of 200 iterations. Convergence is achieved when a minimum average position error w.r.t the target- computed over an episode- - is reached.
Reaching Task:
aollision-vvoidance Task: This task consists of: (1) learning an initial GMM policy of linear reaching motion, and (2) adapting the GMM policy to reach a new horizontally-transiatee target while avoiding collide with two spherical obstacles located midway between the initial robot position and the new target. The initial GMM policy was learned from 10 human demonstrations and represented by a 3-component GMM, as shown in Fig. miiddee and Fig. [3b] For policy optimization, we defined a sparse reward as a function of the position error between the robot end-effector position and the target at the end of the rollout. We also included two sparse penalty terms: the first one punishes rollouts leading to collisions with the obstacles, for which the rollout is stopped; the second term penalizes rollouts with significantly divergent trajectories. Each episode consisted of 10 rollouts, each of maximum horizon length of 150 iterations. Convergence is determined by a minimum average position error w.r.t the target computed over an episode.
This setting involves: (1) learning an initial GMM policy where the robot end-effector reaches two different targets (i.c., task goals) starting from the same initial position, and (2) adapting the initial policy to reach a new target located close to one of the previous task goals. The intended adapted behavior should make the robot go through the most relevant GMM components according to the new target location. The initial GMM policy was learned from 12 demonstrations and encoded by a 6-component GMM, as shown in Fig. riight and Fig. 13c To optimize the initial GMM policy, we specified sparse reward based on the position error between the robot end-effector position and the chosen target at the end of the rollout. Similar to the previous experiments, we added a sparse penalty term to penalize rollouts generating significantly divergent trajectories. An
Multiple-goal Task:
8

---

Figure 2: The three tested 2D robotic settings: reaching skill (left). collision- free trajectory tracking (middle). and a multiple-gov goall task (right). The robot color goes from light gray to black to show the evolution the task reproduction. Green Gaussian components (C) depict the initial GMM policy, projected on the 2D Cartesian position space. The end-effector trajectory resulting from the initial GMM policy is shown in dark blue lines (-). Red circles (+) in the collision-avoidan avoidance task represent the obstacles (middle). The different targets of the multiple-goal task (right) are depicted as red stars.
episode comprised 10 rollouts, each of maximum horizon length of 200 iterations. Again, the policy optimization converges when the average position error w.r.t the chosen target reaches a minimum threshold.
4.2 Results Analysis
The reaching task tested our method's ability to adapt previously-learned reaching skill to new goal, located at (6.0, 6.5 (cf. Fig. left). Achieving this required to adapt the Gaussian parameters of mainly the last four GMM components, while the others remained unchanged. We compared all methods in terms of the success rate over environment steps, where the success rate was defined as the percentage of rollouts reaching the new goal. Figure Beft shows that our method achieved success rate of 1 after approximately 70000 environment interactions. Despite PPO was also able to complete the task reliably, it required many more environment steps (cf. Fig. left). In sharp contrast, SAC did not reach any improvement. These observations underline the importance of some kind of trust region or constraint on the policy updates, which allowed both our method and PPO to reach good success rates. The PMOE versions of PPO and SAC showed a better performance than the vanilla counterparts, which aligns with the experimental insights provided in the original PMOE paper 441 regarding the advantages of the probabilistic mixture of experts in deep RL policies. Nevertheless, our approach also outperformed both PMOE versions. This experiment showed that our method much more sample-efficient in adapting the GMM parameters, which we attribute to the fact that our method explicitly takes the GMM structure into account in the formulation of the optimization.
In the colliiinn-aviddanee task, we tested whether our method was able to adapt trajectory tracking skill in order to avoid collisions with newly added obstacles. These were placed in such way that the robot was forced to move its end-effector through a narrow path between the obstacles (cf. Fig. Fngddle))le) While the reaching task could be adapted by mainly varying the means of the GMM components, this task also demands to adapt the covariance of the second GMM component. Figure Baauremmiddl middle shows that our method solved this task reliably after comparatively few environment interactions. Although PPO also achieved a success rate of 1, it took 6 times more environment steps than our method. SAC only reached an average success rate of 0.8, however with high variance (cf. Fig- (ig-ddllele))) Note that the PMOE versions showed a reduced solution variance and higher success rate than their vanilla counterparts, where PMOE-PPP PPO also outperformed PMOE-SAC. These results again show the importance of the constraints on the policy updates. The huge discrepancy in the required environment steps between our method and all the baselines further emphasizes the importance of taking the GMM structure into account in the policy optimization. The explicit consideration of this structure allows our method to identify the required changes to the GMM parameters much more efficiently. These experimental insights can also be observed in the results corresponding to the other formulations of the collision-avoidanee task, namely, the 3D operational and 7D joint space representations, reported in App. A.6.3
While the previous two tasks were accomplished by adapting mostly the Gaussian parameters of the GMM. the multiple-goal task required to adapt the GMM weights. The initial skill comprised

---

FPO
SUCDDS Gat
C.8
SAC
WIFF
0.5
6.8
0.6
PMOE-SAC
PHGE-FPO
200055
0.8
0.2
0.7
0.9
10k
20k
30k
40k
50K
60C
70K
80K
20k
40k
60k
83k
100k
20k
40k
60K BOK 100K 120K 140K 16CK
envrrnmmnt steps
envirorment steps
environment steps
Figure 3: Success rate of our method (WGF) and the baselines on the reaching (leff), the collision-avoidance (middle) and the multiple-goal tasks (right). The shaded area depicts the standard deviation over runs.
Se000
PP1
5AC
PMOT PROFMOT
PUICE PROPARE GAN
P050
1.6
WHGF
PAIDE PROPMIG SAC
Figure 4: Variance of the success rate over the 5 runs for our method (WGF) and the basclines on the reaching task (left), the collision avoidance task (middle) and the multiple-goal task (right). The violine plots overlaid with box plots, quartile lines and swarm plot, where dots indicate the success plots show variance at the following time steps from left to right: 80000, 90000, 95000.
of individual runs. The
reaching motions two different goals and an execution of this skill results in reaching one of them, depending the sampling noise (cf. Fig. 13c). The easiest way to adapt the policy reach only one the two goals is to reduce the GMM weights of the components belonging to the undesired motion and correspondingly increase the weights of the other components. As shown in Fig. righht our method again quickly achieved success rate of PPO required substantially many more environment steps, while SAC not able to solve the task. Again, the PMOE formulation reduced the solution variance and slightly increased the success rate of both PPO and SAC.
In Fig. 4 we report the success rate variance over 5 runs at a fixed time step, which corresponded to the step at which the first method achieved success rate of 1, thus prioritizing sample efficiency. The plots show that our method exhibits very low solution variance. All baselines varied largely, except for the reaching task, where SAC runs collapsed to a success rate of 0. These results show that our method, despite showing large variance at the start, was able to quickly reduce the variance and converge reliably to good success rate. We also provide similar plots of solution variance in Fig. 6 where we report the results for both PPO and SAC using their own convergence time step.
Conclusions and Future Work
We presented novel method for GMM policy optimization, which leverages optimal transport theory to formulate the policy optimization as a Wasserstein gradient flow on the manifold of GMMs. Our formulation explicitly accounts for the GMM structure in the optimization and furthermore enables us to naturally constrain the policy updates by the I Wassseerstein distance between GMMs to enhance the stability of the policy optimization process. Moreover, the embedding of the Gaussian components of the GMM policy in the Bures- Wasserstein manifold greatly reduced the computational several robotic tasks provided strong evidence of the
cost of the policy optimization. Experiments importance of our policy-structure aware optimization against approaches that disregard the GMM structure. A possible limitation of our method is that each optimization loop involves running the Sinkhorn algorithm, which is computationally expensive. This might be improved by employing recent advances on initializing the Sinkhom algorithm (7].. Also, we observed an intricate interplay between the optimization of the GMM weights and the Gaussian parameters, which occasionally resulted in one update hampering the other. In future work we plan to address the latter problem by using separate adaptive learning rates for weights and Gaussian parameters. Another possibility would entail to reformulate the policy optimization as a gradient flow on the space of probability measures endowed with the Wasserstein Fisher-Rao metric [72] 73, 74). This alternative metric may allow us to leverage the Fisher-Rao geometry to optimize the GMM weights, as very recently proposed in 75 to learn isotropic Gaussian mixtures via particle-based approximations. Finally would be interesting to combine our method with an actor-critic formulation and to replace the multi-step cumulative reward by trained Q-function.
10

---

References
[1] Jan Peters, Daniel D. Lee, Jens Kober, Duy Nguyen- Tuong, J. Andrew Bagnell, and Stefan Schaal. Robot Learning, pages 357-398. Springer International Publishing, 2016. ISBN 978-3-319-32552-1, URL
/ttps://doi..or/// 1007//978 319 32552- 1_112-1_11 [
[2] Stefan Schaal, Auke Ijspeert, and Aude Billard. Computational approaches to motor learning by imitation. Phil. Trans. R. Soc. Lond. B, 358:537-547, 2003. http: httppppp//ao...oo11111////////// 1098///t..999//tt/ 2002 1258
[3] Jens Kober, J. Andrew Bagnell, and Jan Peters. Reinforcement learning in robotics: A survey, The International Journal of Roboticy Research (IJRR), 32(11)112818812211 2013. URL httpp::////oo org/
10.
Konstantinos Chatzilygeroudis. Vassilis Vassiliades, Freek Stulp. Sylvain Calinon, and Saan-Baptiste Mouret. survey policy search algorithms for learning robot controllers in handful of trials. IEEE Transactions on Robotics, 36(2):328-34 347, 2020.
1110//....../........ 2019
2958211
1
[5] Freek Stulp and Olivier Sigaud. Policy improvement methods: Between black-box optimization and episodic reinforcement learning, hal-00738463, 2012. http://hal, www.aaas--hivveevvreereeoee fr/
hal ha1-00738463
[6] Noémie Jaquier, Leonel Rozo, Sylvain Calinon, and Mathias Bürger. Bayesian optimization meets Riemannian manifolds in robot learning. In Conference Robot Learning (CoRL), pages 233-246, 2019. URL
mlr prrss//100/jaauierrrrr
html
Shixiang Gu, Ethan Holly, Timothy Lillicrap, and Sergey Levine. Deep reinforcement learning for robotic manipulation with asynchronous off-policy updates. In IEFE International Conference on Robotics and Automation (ICRA), pages 3389-3396, 2017. httppgss//bbo..rr///oo110//aaaa 2017 7989385
[8] Julian Ibarz, Jie Tan, Chelsea Finn, Mrinal Kalakrishnan, Peter Pastor, and Sergey Levine. How to train your robot with deep reinforcement learning: lessons have learned. The International Jourmal Robotics Research (IJRR), 40(4-5))698-721, 2021. URL bttpa::ta:///o.... org/ 1177 /2t1///2222664298585555555
|
Mare Peter Deisenroth, Gerhard Neumann, and Jan Peters. A survey on policy search for robotics. Found. Trends Robot, (112)):1142) 2013. URL
org http://.o......11//00/0000000
[10] Auke Jan Ijspeert, Jun Nakanishi, Heiko Hoffmann, Peter Pastor, and Stefan Schaal. Dynamical primitives: Learning attractor models for motor behaviors. Neural Computation, 25:328-373, 2013. URI.
https::/doi. org/ /1ogg///... 10. 1162/ NECO_ a_ 00393 [
[11] Sylvain Calinon, Florent Guenter, and Aude Billard. On learning, representing, and generalizing task in humanoid robot. IEEE Transactions on Systems, Man, and Cybernetics, Part B (Cybernetics), 37(2): 286-298, 2007.
org/ 10. 1109/TSMCB 2006. 886952
[12] Alexandros Paraschos, Christian Daniel, Jan Peters, and Gerhard Neumann. Using probabilistic movement primitives in robotics. Autonomous Robots, 1215,9-551, 2018. URL httpp:://doi.org/1000
$10514-017-9644 9648-7
[13] Leonel Rozo and Vedant Dave. Orientation probabilistic movement primitives Riemannian manifolds. In Conference on Robot Learning (CoRL.), pages 373-383, 2021. URL httpp:://proceeedingg mlr
press/v164/roozo2a..o html
[14] Shikhar Bahl, Mustafa Mukadam, Abhinav Gupta, and Deepak Pathak. Neural dynamic policies for end-to-end sensorimotor learning. In Neural Information Processing Systems (NewrIPS), pages 5058-5069, 2020. URL
.df
[15] Thomas Cederborg, Ming Li, Adrien Baranes. and Pierre-Yver Oudeyer. Incremental local online Gaussian
mixture regression for imitation learning of multiple tasks. In IEEE/RSI International Conference on Intelligent Robots and Systems (IROS), pages 267-274, 2010. URLpsttt:://///t http::///oooooo/g//// org/10. 1ew./1109RRR0 IROS
1510.6650040 2
[16] Sylvain Calinon. A tutorial on task task- parrammtterizeed movement learning and retrieval. Intelligent Service Robotics, 9(1):1-29 January 2016. ISSN 1861-2776. URL https://doi. http://doooorgg..../ 1007/ 11370-015-0187.

---

[17] Noémie Jaquier. David Ginsbourger, and Sylvain Calinon. Learning from demonstration with model-
based Gaussian process. In Conference on Robot Learning (CoRL), pages 247-257, 2019. http
//proceedings.mie mlr. wttp:///gwg......pp///o///00 jaw.ooroiuuuuub....o htmm
[18] N. Jaquier, I.. Rozo, D. G. Caldwell, and S. Calinon. Calaannn, Cwwmmetty-wwwwree-mmnpllllaallla leaming,. tracking and transfer. 2063):22--650 2021. URL http:p//ooo....rr/.....177777762
[19] Jihong Zhu, Michael Gienger, and Jens Kober. Learning task parameterizeze skills from few demonstrations. IEEE Robotics and Automation Letters (RA-L). 7(2):4063 407000033000000 2022. URL https:://tttoi... org/10
1109/LRA. 202...22002.111
2
[20] S. Mohammad Khansari- Zadeh and Aude Billard. Learning stable nonlinear dynamical systems with Gaussian mixture models. JEEE Transactions Roboticy (T-RO), 27(5):943-957, 2011. URL https
//doi. //doi.org/10.. wttpg//...101///TI 2011.2159411
2
[21] Harish Ravichandar, Iman Salehi, and Ashwin Dani. Learning partially contracting dynamical systems from demonstrations. In Conference on Robot Learming (CoRL), pages 369-378, 2017. URL https
//proceedinge mlr. pptps//v8//raiicaaaaaiiaee
html
[22] Nadia Figueroa and Aude Billard. A thysically-consiieeee Bayesian non- parametramee mixture model for dynamical system learning. In Conference on Robot Learning (CoRL), pages 927-946, 2018. URL
httpp:://proceedinga mlr wttppp/eeeeseeeeeeeeeeeeeeeeeeee httml
[23] Marco Ewerton, Gerhard Neumann, Rudolf Lioutikov, Heni Ben Amor, Jan Peters, and Guilherme Maeda. Learning multiple collaborative tasks with a mixture of interaction primitives. IEEE International Conference on Robotics and Automation (ICRA), pages 1535-1542, 2015. Utthtss:::/ddat org/10
1109/ICRA 2000//CRA.201....... 00315..1.00 7139393 2
[24] Leonel Rozo, Sylvain Calinon, Darwin Caldwell, Pablo Jiménez, and Carme Torras. Learning physical collaborative robot behaviors from human demonstrations. IEEE Transactions on Robotics, 32(3):113-577,
2016. URL atttps://dol org/ 10. 1109/TRO. 2016 2540623 2
and Axel Gräser. Robot learning [25] Maria Kyrarini, Muhammad Abdul Haseeb, Danijela Autonomous Robots, 43:239-222 257, 2019. URL of industrial assembly task via human demonstrations.
httts:///oi org/ 10. 00000//110014401977555555
[26] Leonel Rozo. Meng Guo, Andras G. Kupesik, Marco Todescato, Philipp Schillinger, Markus Giftthaler, Matthias Ochs, Markus Spies, Nicolai Waniek, Patrick Kesper, and Mathias Burger. Learning and sequencing of object-centris manipulation skills for industrial tasks. In JEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pages 9072-9079, 2020. http:ttt///ooorrg/// 1109/
IROS45743. 2020. 9341570 2
[27] Yun Lin, Shaogang Ren, Matthew Clevenger. and Yu Sun. Learning grasping force from demonstration. In IEEE International Conference on Robotics and Automation (ICRA). pages 1526-1531, 2012. URL
https://doi. org//o.../0/.. 1/t///0.111000 1ww.9/19IIIRR,.. 2012. 6225222
[28] Fares J. Abu- Leonel Rozo, and Darwin G. Caldwell. Force-based variable impedance learning for robotic manipulation. Robotics and Autonomous Systems (RAS), 109:156-167, 2018. URL https
//doi.org/00 10. wow.tt.r.oooo..............ooo.. robot. 2018.0112000020 0.07. .008 2
[29] Yanlong Huang, Leonel Rozo, João Silvério, Darwin G. Caldwell. Kernelized movement prim- itives. The International Journal of Robotics Research (I/RR), 38(7)8333-852 2019. doi: 10.1177/
0278364919846386 URL. https: /ooi org/ /tt///1...11//////// 101.1117002782649999888888
[30] Oleg Arenz, Mingjun Zhong, and Gerhard Neumann. Trust-region variational inference with Gaussian mixture models. Journal of Machine Learning Research (JMLR), 21(163):1 160, 2020. URL http:
//jmlr.org hwtpppprrs/pppp///////oo http//ppppsss/////e///pe 121/11111111 524.hmml
[31] Iman Nematollahi, Erick Rosete-Beat Adrian Röfer. Tim Welschehold, Abhinav Valada, and Wolfram Burgard. Robot skill adaptation via soft actor-critic Gaussian mixture models. In International Conference Robotics and Automation (ICRA), pages 8651-86571 2022. URL htttps://doi... org/ 1109/
ICRA46639. 2022. 98811770
3
[32] Filippo Santambrogio. Optimal Transport for Applied Mathematicians, Birkhäuser, Cham, 1 edition, 2015. http/sstp::///oi. org/ 10. 1007 10000/00777-0-- 1-199988---19-888------ 2 2
http://doooorr//////
[33] Gabriel Peyré and Marco Cuturi. Computational optimal transport: With applications to data science. Foundations and Trends in Machine Learning, 1 ((5-6):355-60) 2019. http:: doi org/ 10
1561/2200000073
2
12

---

[34] Yongxin Chen, Tryphon T. Georgiou. and Allen Tannenbaum. Optimal transport for Gaussian mixture models, IEEE Accesx, 7:6269-6278, 2019. httptttss///di. 1ttg///..10//ACESSS. 2018. 2889838
[35] Julie Delon and Agnès Desolneux. A Wasserstein-type distance in the space of Gaussian mixture mod- els. SIAM Journal on Imaging Sciences, 13(2):936-970, 2020. URL htttpa///oo..org/10 1137/
19M1301047
[36] Pierre H. Richemond and Brendan Maginnis. On Wasserstein reinorrcmett learning and the Fokker-Planck equation, 2017. URLpttps://aarxt http:///aaxx...ooora//////////10 arXiv:1712.07185.
[37] Richard Jordan, David Kinderlehrer, and Felix Otto. The variational formulation of the Fokker- Planck equation. Siam Journal on Applied Mathematics, 1996. URL https:///ddo.oo org.//00....r/0... 1137/
90036140096303359
4
[38] Ruiyi Zhang, Changyou Chen, Chunyuan Li, and Lawrence Carin. Policy optimization as Wasserstein gradient flows. In International Conference on Machine Learning (ICML), pages 5737-5746, 2018. URL
https://proceediin mlr, press/v80/ fteess/800/hang182
html
[39] Petr Mokrov, Alexander Korotin, Lingxiao Li, Aude Genevay, Justin Solomon, and Evgeny Burnaev. Large-scale Wasserstein gradient flows. Neural Information Processing Systems (NeurlPS), 2021. URI.
www.tetf/oooee formmm ? LedddnI..JuuHaaa
[40] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International Conference on Machine Learn- ing (ICML). pages 1861-180 1870, 2018. URL/https: http:: httpss:/pproceedingg mlr www.sssws......e...........eeeee 1/88/haaroo_ j18bb html
[41] Jic Ren, Yewen Li, Zihan Ding, Wei Pan, and Hao Dong. Probabilistic mixture- experts for efficient g
deep reinforcement learning, 2021.
09122
9
[42] Shijun Wang, Baocheng Zhu, Chen Li, Mingzhe Wu, James Zhang, Wei Chu, and Yuan Qi. Riemannian proximal policy optimization, 2020. URL https: www.: hbsp://bs200.0... 0911/ arXiv:2005.09195.
[43] Shun-ichi Amari Natural gradient works efficiently in learning. Neural Computation, 10(2):251-276,
1998. https:ttt::///o wwwoorg///10 10. 111111//099766663000777777 3
[44] Sham M Kakade. A natural policy gradient. In Neural Information Processing Systems (NeurIPS), 2001. URL https://proceeddi pdf Sitaaaeelaaaccccee5e5csssoooa300
[45] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms, 2017. httpss////// orgo //////b///00/// 06347
5
[46] Alexandre Galashov, Siddhant Jayakumar, lconard Hasenclever, Dhruva Tirumala, Jonathan Schwarz, Guillaume Desjardins, Wojtek M. Czameeki, Yee Whye Tch, Razvan Pascanu, and Nicolas Heess. In- formation asymmetry in KL-regularized RL. In International Conference on Learning Representations (ICLR), 2019. URL
[47] Aldo Pacchiano, Jack Holder, Yunhao Tang. Krzysztof Choromanski, Anna Choromanska, and Michael Jordan. Learning to score behaviors for guided policy optimization. International Conference on Machine Learning (ICMI), pages 7445-7454, 2020. URL. htttps://proceeddnngg .mr............................. prre
v1 1119/pacchinn200.. html
[48] Ted Moskovitz, Michael Arbel, Ferenc Huszar, and Arthur Gretton, Efficient Wasserstein natural gradients for reinforcement learning. In International Conference on Learning Representations (ICLR), 2021. URL
https:
Shwprrumnn.d--OOnnSSStn2.. 3
[49] Wu Lin, Mark Schmidt, and Mohammad Emtiyaz Khan. Handling the positive-definite constraint in the Bayesian learning rule. In International Conference on Machine Learning (ICML), pages 61166661 16-6126, 2020. URL http:ttts///rrrceeiingg mlr
http://eeeaaagaa.......re//////a wttppessvv19//1120....ttt html
[50] Filippo Santambrogio. Euclidcan, metric, and Wasserstein gradicnt flows: an overview. Bull. Math. Sci., 7: 87-154, 2017. /ttp:///::/ 10. 100.100013/30311000000000
[51] Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradiens flows: in metric spaces and in the space of probability measwres. Springer Science & Business Media, 2005. URL htap:::/ddo..org//// 1007
b137080 4
13

---

[52] Victor M. Panaretos and Yoav Zemel. The Wusserstein Space, pages 37-57. Springer International Publishing, 2020. https:////// wttp://..100/997-3-0000----00000
[53] Zoubin Ghahramani and Michael Jordan. Supervised learning from incomplete data via EM approach. In Advances in Neural Information Processing Systems (NewrIPS), 1994, URL httph::/t hattps://rroceedingg
neurips. wwtpppppp/pcccecccccccccccc..
Paper. pdf
[54] Benjamin Eysenbach and Sergey Levine. Maximum entropy RL (provably) solves some robust RL problems. In International Conference on Learning Representations (ICLR), 2022. URL https:// openreviev.
[55] John Schulman, Sergey Levine. Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In International Conference on Machine Learning (ICML), page 1889-1897, 2015. URL https://proceeddingg mlr
html 5
[56] Fabian Otto, Philipp Becker, Ngo Anh Vien, Hanna Carolin Ziesche, and Gerhard Neumann. Differ- entiable trust region layers for deep reinforcement learning. In International Conference on Learning Representations (ICLR), 2021. [tttttts/////perreeeet htt.:p/awwe.....teeetrrrrrro/o/. net//orum?iddqqy2) A01Vn 5
[57] Felix Otto. The geometry of dissipative evolution equations:: The porus medium equation. Communi- cations Partial Differential Equations, 26(1-2):101-14,, 2001. /ttts:///// 1081/ PDE- 10000200
[58] Yifan Chen and Wuchen Li. Optimal transport natural gradient for statistical manifolds with con- tinuous sample space. Information Geometry, 3, 06 2020. URL Attps://doi.orr/100
s41884-020- 00028-0 0
1
[59] Luigi Malagò, Luigi Montrucchio, and Giovanni Pistone. Gaussian densities. Information Geometry, 1:137-179, 2018. URL. httpp:://aoi..rg//0... 1007/
Wasserstein Riemannian geometry of
041884-018-0014
16 19
[60] Andi Han, Bamdev Mishra, Pratik Jawanpuria, and Junbin Gao, On Riemannian optimization over positive definite matrices with the Bures Wasserstein geometry. In Advances in Neural Information Processing Systems (NeurIPS), 2021. URL lttpppppp///oeeereevvieee httpe/vwwwwweneeeeeeeeeeeeeeeeee
16 19
[61] Nicolas Boumal, An introduction to optimization on smooth manifolds. Cambridge University Press, 2023.
URL http::ttp////tt www.e./icooaabbbooo aab.omaaabmmal....m.mmm.. net/bbokk 6
[62] Ronald Williams.
Simple statistical froliwin-folowing algorithms for connectionist reinforcement learning. Machine Learning, 8:229- 256, 2004.
6
[63] Marco Cuturi and Amaud Doucet. Fast computation of Wasserstein barycenters. In International Conference on Machine Learning (ICMI.,, pages 685-693, 2014. URL. https://proceeddin mlr vress/v32// v32////ur//vvte ctttrrituriii 14. html
[64] James Townsend, Niklas Koep, and Sebastian Weichwald. Pymanopt: Python toolbox for optimization on manifolds using automatic differentiation. Journal of Machine Learning Research (JMLR), 17(137); .ttml
1-5, 2016. http:://///::.oortt wwwppe/rrr.pppppp//eeeeaaeeeeeee
19
[65] Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar 7. Alaya, Aurélie Boishunon, Stanislas Chambon, Lactitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Lóo Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy, Tevgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy. Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. Pot: Python optimal transport. Journal of Muchine Learning Research (JMLR), 22(78):1-8, 2021. URL hntpp
//jmlr org/ httg:/apprs/////
seneers/v22/200 h2////222220200 451 html
[66] Peter Corke and Jesse Haviland. Not your grandmother's 's toolbox - the robotics toolbox reinvented for Python. In IEEE International Conference on Robotics and Automation (ICRA), pages 11357-11388 2021. URL httpss/tt:://10
2021. .9561366
[67] Carlo Drammo Davide Tateo, Andrea Bonarini, Marcello Restelli, and Jan Peters. Mushroomul: Simplify- ing reinforcement learning research. Journal of Machine Learning Research (JMLR), (2(131):1-5, 2021. URL //ttt:////mm.....//////// org wttpg//ppere/v221/ 056.. html
[68] Francesco Braghin Asad Ali Shahid, Dario Piga and Loris Roveda. Continuous control actions learning and adaptation for robotic manipulation through reinforcement learning. Autonomous Robots, 46:483 498. 2022. URL httph:tt//:/tt //doi. org/ 1007/ 120-0444---202----222
14

---

[69] Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dormann. Stable-baselines:: Reliable reinforcement learning implementations. Journal of Machine Learning Research, 22(268): 2021. URL.tttp:/jjmr..oo http:////jaaa.oooopppppa/e////// 1364..01....... .4tml
[70] Takuya Akiba, Shotaro Sano, Toshihiko Yanase, Takeru Ohta, and Masanori Koyama. Optuna: A next- generation hyperparameter optimization framework. In ACM SIGKDD International Conference on Knowledge Discovery and Data Mining (KDD), page 2623-2631, 2019. URL httLttttt:/:/ii...e
1145/3292500. 33307001
[71] James Thortoo and Marco Cuturi. Rethinking initialization of the Sinkhom algorithm, 2022. URL https://arxiv. org/ /bs//206..... 0766007630
07630 arXivv2206.06630. TO
[72] Lenaic Chizat, Bernhard Schmitzer, Gabriel Peyré, and François-Xavier Vialand. An interpolating distance between optimal transport and Fisher-Rao, Rao, 2015. URL pttttp:://aaxvv.org//e../.
06430 T0
[73] Thomas O. Gallouët and Léonard Monsaingeon. JKO splitting scheme for Kantorovich Fisher Rao gradient flows. SIAM Journal on Mathematical Analysis, 49(2):1100-111 11300000011,,,,1100000001100--11 2017. URL https://doi org/10 10. 1137/16M006668
TO
[74] Matthias Licro, Alexander Mielke, and Giuseppe Savaré. Optimal entropy-transport problems and new Hellinger- Kantoroick distance between positive measures. Inventiones mathematicae, 211, 03 2018. URL
https://0o1..org/101100000
0759-8
I0
[75] Yuling Yan, Kaizheng Wang. and Philippe Rigollet. Learning Gaussian mixtures using the Wasserstein- Fisher-Rao gradient flow, 2023. URL atttp::/aarxvv.or//aaeee
T0
15

---

A Appendix
A.1 Details on Gaussian Mixture Regression (GMR)
In GMR we start from a GMM in state-action space TTs a) = Σx= ([sa]] pu Σ,) from which policy, i.. a probability distribution on the action space, can be obtained by conditioning on the state, as follows
;
π (8, a) J (s, a.)da
T (a 8)
(15)
The resulting conditional distribution is another GMM on the action sapce, with state dependent parameters, given by:
π Taa 8x)) Σ
Wi (a:)) (8)), 21), with
(16)
i=1
(s:) = u, + Σ (E))
(st 14))
(17)
" = Di D. (>;) 2000
(18)
(x)) = wiN( (8ti u. -i)
(19)
Σx WeN (8t pi: Eg)
Note that we have split the GMM parameters μ4 and E; into their state and action components according to
14
Σ
D
p1:
E. =
as
(20)
A.2 Riemannian gradients and retractions
For completeness we give here the explicit expressions of the Riemannian gradients and the retractions used in $ [3.1] As the mean vectors are assumed to lie in the Euclidean space, their Riemannian gradients actually coincide with the Euclidean gradients and no retraction is required, so Eq. 12 reduces to the well-known Euclidean gradient descent
p fax + [FJ((T),
(21)
where V denotes the Euclidean gradient w.r.t. á. For the covariance matrices we use the gradient and retraction w.r.t. the Bures- Wasserstein manifold, taken from 59 The gradient is given by
grads J(TA) 2 ((Tx)
(22)
where again V V denotes the Euclidean gradient w.r.t. S and (X}s (X+XT)
Furthermore, the retraction is given by
RE.
x
E
X
Lx
X X
(23)
where Cx [ is the Lyapunov operator, defined as the solution to the matrix linear system
Lx ΣΦ + XL Σ
Σg-
(24)
A.3 Expressions of the free functional ((T) and its Euclidean gradients
For completeness sake, we provide here the explicit expression of the Euclidean gradients for the objective J(TT) w.r.t the parameters of the GMM, which are used in the construction of the Riemannian gradients. Using the policy gradient theorem, obtain the gradient of Eq. [ w.r.t to
16

---

parameter& & as follows
V+J(a) Ve / [ (:)(s ai) Er'(ss,arr, (25)
E,
Ve (elo((T(a)(ss))
Σ Er(se,
E,
Ve log
π (8 an)
Σ Mat
r(st. ai)
[ddaTT(s))a)
LE, E (
(e T(S ar)
[da.Ve ai) [ a.)
Er(su. r(su an)
est
In this work, focus GMM policies, for which the objective (TT) takes the form:
((T) = / p()))) ((:1)))())))s))))
((ss))) oz) Exr(ss,,o ar)
+ 3 f da Σ "
wi (()))))(
log
wi (s)) (i(+:)) 81, ai)
(26)
By inserting Eq. 26 into Eq. 25 we obtain for the individual parameters of the GMM
WIN az ki, Σ E ((8), at) 14)
[H.((T) = E,
(27)
[s ac My. E,)
wi I daN(ss, ai: p1, E E (($), a:) - mi)
ail
E, daN az M.j Mx²)
12t
E, WWN as pl, Σ D (($), at) mi)
(28)
Ej Lo,N((s ari pej Σn}
8 WIN(ss: Miss E1,es 1,50 (st (4,s)
Σ Ent
r(s,, ar)
Ej W.NN Aui,N(sss (885 Misss Ejjos)
E, Σ (i ((s). az) -ui) 8,L1. Di,o (st
m1,e)
Er(sr
Here 8, E {0, 1) indicates which terms the gradient acts on. In this case, the gradient act on the state components and it is absent for the action dimensions.
WiN(sn au ui, Ei) (1 (($t, a.) - 14) ((s), - -mi) )
JIT = Er
Σ
2
E, N E,)
(29)
Wi f daN ari pi, E1)2,1 (1 (1- ((0), a.) - mI) ((8r, ar) - -11) ET
r(si ai)
+
Σ
r
2
E₁ wy daN ati Hg.
P>t
= Er Σ
Sl.on. 1 (1 ((Sr, az) 11) ((Sr, ax) 11) 21-1)
2
+5,
Ç1.0: -1 (1 (1- (s: - m1,) (st - mmx) M
Σ rst
r(s,, an)
2
1,5
17

---

N(st, de: mI, Σi) ΣS Hri 2z})
f ddN azi ur: Ei)
VwiJ(R) Er Σ
Σ
r(s ar)
Σj wy S daN ari Hri En)
(30)
(1.00,01 S1,00)
Er(s,
= E
(31)
W1
Note that we introduced the responsibilities Si and St. which are defined as follows
Si WiN(s,, ar: p4; Ei)
and
(32)
E, (s,, ar Míj m,)²
Wi I dsN as pi Σi) St.s. En wj f daN(s,, azi Mg Mx)
(s Mi.s Mi. D1,0a) EjwNN(Br Mij 25,20)
(33)
A.4 Relation between forward and backward discretization in the Bures- Wasserstein metric
In this section we outline the relation between the implicit and explicit optimization schema w.r.t. the Bures- Wasserstein metric, which leveraged to formulate our policy optimization in $ We closely follow [58]. For the sake of simplicity, we group the Gaussian parameters μu and E into single parameter vector Furthermore, we restrict our explanation which is possible without loosing generality, as each of the N components live in its own manifold Rdd S4+
single Gaussian component,
The Riemannian gradient w.r.t the Gaussian parameters 0, gradg J( satisfies by definition
ge(grad (T) {) = (TT(( E,
(34)
where Vo denotes the Euclidean gradient, is an arbitrary vector on the tangent space TrM, and ge is the Riemannian metric tensor, defining the inner product on T. The Riemannian metric can be written
go( GTG (0)g.
(35)
with two arbitrary tangent vectors Ç. E. and Gw (0) being a positive definite matrix. Moreover, note that the Wasserstein distance W3 (N) + A0)). where AO denotes small perturbation in the Gaussian parameters 0, can be expressed as
1 (N(O0), + A0)) - (2))
(36)
for AO + 0. Similarly, we can approximate the objective evaluated at J(Ø + A0) via the Taylor theorem as
J(O + J(0+10)
(37)
With this, we can approximate
(()) - -J(((0))
0 k++ arg min
(38)
0
2T
(0 02)
Gw (
0x) VoJ (0 0k))
arg min
(39)
2T
from which obtain the update equation for 0 as follows
Q 0 + TGw (Vo
(40)
which corresponds to the Wasserstein natural gradient with respect to the Bures- Wasserstein metric. Note that Eq. 40 in turn corresponds approximation of the exact Riemannian gradient descent
0 RR. .rrdd (0x)))
(41)
18

---

This approximation can be obtained by considering a first-order approximation of the geodesic on the BW manifold. As the exponential map (a.k.a. the retraction) defined via the geodesic, the retraction operator in Eq. 41 turns into simple addition operation under first-order approximation, leading to Eq- 40 Notice that such approximation does not guarantee that the updated parameters stay on the manifold, except for the cases in which O IR This means that the positive- constraints arising from the covariance matrices can be easily violated when using natural gradient approaches as they do not guarantee that the updates stay the underlying manifold. However, it is worth noting that the use of an inexact Riemannian gradient descent often motivated by the difficulty of computing the exponential map (or geodesic) necessary to calculate the exact Riemannian update. In our case, we leverage the retraction and Riemannian gradients [59] 60), which allow us to apply the exact Riemannian gradient descent of This avoids to rely first-order approximations and in turn we can guarantee that the updates of the Gaussian distribution parameters always lie on on the product manifold (R* x S )
A.5 Additional details the implementation
We extended the Pymanopt (64) by adding a custom line-search routine that accounts for constraint on the Wasserstein distance between the old and the optimized GMMs. The details of this line-search can be found in Algorithm 2
Algorithm 2 Constrained line-search. The constraint function -( is arbitrary in general. We use the 2 L2sWasserstein distance between two points on the manifold of GMMs as constraint.
Input: point vo on the manifold, descent direction d, initial step size λo, decrement a, constraint c(aro,)) maximum allowed value for constraint Cmax: minimum step size Amin Output: step size 8, updated point on manifold æ:
1: x = 2o + λo d
X Ao
2: while c(zo. x) > Cmax and À > Amin do
3: decrease step size: >= a
À d update point manifold:
4: end while
5: if X < Amin then
6: return λo To
7: else
8:
return X, a
9: end if
A.6
Additional details on experiments
A.6.1
Additional results
Fig- shows the convergence curves for the two baselines in Fig. of the main paper, however, we extended the horizontal axis up to the maximum number of environment steps used for training.
3.0
LUCC055 c300
0.6
0.8
0.6
0.8
0.4
0.4
0.2
50k
150K
150K confrormmnn staps
25k
175K
cnvirommant stops
envirommmnt staps
Figure 5: The success rate of the two baselines on the reaching task (left), the collision-avoidance task (middle) and the multiple-goal task (right). The shaded area indicates the standard deviation over 5 runs.
Fig. G shows the variance of success rate for the three methods their time step of convergence for all three robotic tasks. Concerning SAC, which did not converge after the maximum number of
19

---

PPU
PPO
PPO
nettod
SAC
6AC
-155
01.0
6..5
1.5 2.0 1000
-05
12.3
61.5
1.00
1.5
2.0
-10
-155
03.11
0.5
1.0
1.5
7.
SU0TPS 1776
Figure 6: Variance of the success rate over the 5 runs for our method (WGF) and the two baselines on the reaching task (leff), the collision avoidance task (middle) and the multiple-goal task (right). The violine plots overlaid with box plots, quartile lines and swarm plot, where dots indicate the success time steps at which we determined the variance right: (280000, 400000, 80000), (275000, 300000, 90000), (130000, 200000, 95000).
of individual runs. PPO, SAC and WGF the three tasks from
WSF
success cao
ron UW
200055 gtee
6.8
S6
0.8
0.4
02
0.2
0.2
0.0
non, BW
ron BW
20K
30k
408
SON
BBK
70K
BOK
20k
60X
BOK 1004
120K
140k
20k
436
60k
50k
100k
120k
envtronment sTeas
envirormen steps
enW ronment steps
Figure 7: The success rate of our method and an ablated version, not using the Sanes-Wasserstein formulation for the reaching task (leff), the colision-avoidance task (middle) and the multiple-goal task (right). The shaded area indicates the standard deviation over 5 rums.
environment steps used for training, we chose the last time step. Specifically, we chose the following time steps for PPO, SAC and WGF, respectively: reaching task (280000, 400000, 80000), collision avoidance task (275000, 300000, 90000), multiple goal task (130000, 200000, 95000). These plots show that PPO may also reach low-variance success rate over the five runs at the time step of convergence, the cost of prohibitively large number of steps. SAC showed huge variance in all tasks, apart from the reaching task, where all runs collapsed to success rate of O.
A.6.2 Ablation study
In order to assess the influence of leveraging a Riemannian optimization approach on the Bures- Wasserstein manifold, we conducted an ablation of our method by eliminating the Riemannian formulation. Instead of the explicit Euler scheme update in Eq. [2] which corresponds to Riemannian gradient descent the Mures-Wasserstehi metrie, we the implicit Euler scheme
W2 (TT (A),TM)
A = arg min
(TT ((T)(()))
(42)
2r
( [), TK) (E),TT)
Σx+1 = arg min
J(TTA(2)
(43)
27
To guarantee that the updated covariance matrices do not leave the manifold of symmetric positive definite matrices, parameterize them in terms of Cholesky factors. The results obtained with this Din-Riemnnnian version of our method are shown in Fig. in direct comparison to our method and Fig. 8 for an extended range.
1.0
BW
1.0
BW
0.8
M
0.8
0.6
AM
0.6
0.4
0.4
50k 100K 150k 200k 250k 300K 350K 400K environment steps
25k
50k
75k 100k 125k 150K 175K 200k environment steps
Figure 8: Extended plot of the based formulation for reaching task (left), the collision-avoidame task (middle) and the multiple-goal task (right). The shaded area indicates the standard deviation over 5 runs.
rate of and ablated version of method, using the Bures- Bures-Wassersee Wasserstein-
20

---

The results clearly show that the non- non-Riemmnn Riemannian method struggles to reach a success rate of 1 for the reaching task and the collision-avoidance task. Furthermore, we observe high variance over different runs in the same settings (see Fig. [ Fig. T0P. We attribute this to the fact that the our method takes exact gradient steps in the direction of steepest descent w.r.t. the underlying BW metric, whereas the implicit scheme only approximates this direction. For this reason the non-Riemannian method is much more noisy, which turn leads to the aforementioned high variance. Nevertheless, the multiple-gov task constitutes exception. Here observed a similar performance for approach and the ablated method. The reason for this that the optimization of this task mainly dominated by the weight updates, which are identical for both methods. This result is therefore expected and confirms that correctness our ablation strategy.
WGF
WGF
method
method
method
BW
non_BW
0.0
-100 0.5
1.5 2.0
00.
0.0
0.5
1.0
1.5
2.0
Figure 9: Variance of the success rate over 5 runs for our method (WGF) and the ablated method (non-BW) on the reaching task (leff), the collision avoidance task (middle) and the multiple-god task (right). The violine plots overlaid with box plots, quartile lines and swarm plot, where dots indicate the success rates of individual runs. The time steps which we determined the variance 80000, 90000,85000.
WGF
WGF
medhod
method
mothod
BW
non BW
non_BW
1.0 D. 0.0 0.5 10 1.5
2.0
-1.0 0.5
0.0
1.0 1.5 success rate
0.5
2.0
110 -0.5
0.0
0.5
1.0
1.5
2.0
success rate
success rate
Figure Variance of the success rate over runs our method (WGF) and ablated method (non (Wo on the reaching task (left), the collision avoidance task (middle) and the multiple goal task (right). The violine plots are overlaid with box plots, quartile lines and a swarm plot, where dots indicate the success rates of individual runs. The time steps at which we determined the variance are (80000, 400000). (90000, 200000), (85000, 90000).
A.6.3 Additional experiments with 7-DoF robotic manipulator
We carried out two additional experiments to show that our method can be employed on tasks per- formed by off-the-shelf robotic manipulators (c.g. 7-DoF Franka Emika Panda robot). Specifically, we extended the collision-avoidaae action depend on the task space representation. The first task was represented in the robot operational space, i.e., the state x E R3 and the action a robot joint space, consequently the state q IR and the action was aimed at assessing the capabilities of our approach adapt robot motion policies state-action spaces of higher dimensions. The initial 3-components GMM policy was trained using 10 human demonstrations featuring linear reaching 3D trajectories. For policy optimization, we used a sparse reward defined as function of the position error between the robot end-effector position and the target at the end the rollout. Moreover, two sparse penalty terms were added to punish collision with obstacles and divergent trajectories.
task described $ 3D environment, where the state and
R The second task was represented the
R?. This experiment
Similarly to the planar task reported in the main paper, we tested whether our method was able to adapt a trajectory tracking skill in order to avoid collisions with newly added obstacles. This means that the robot end-effector needed to pass through narrow path between two spherical obstacles. For the operational space representation, the robot end-effector pose was controlled using full-pose Cartesian velocity controller frequency of 100Hz, where the end-effector orientation was kept constant. For the joint space representation, the robot joint configuration was controlled using joint velocity controller frequency of 100Hz. Figure [T shows results for 3D operational space representation, where method reached success rate of 1.0 very quickly, taking approximately
21

---

Figure 11: success rate our method applied to the 3D narrow-path task performed by the 7-DoF Panda robotic manipulator. The shaded area indi- cates the standard deviation over 5 runs.
Success
0.4
0.2
0.0
WGF
20k
40k
60k
8OK
100k
120k 140k
environment steps
1.0
PPO
The success rate of our method (WFG) and the baselines formed by simulated 7-DoF Panda robotic ma- nipulator. The task is represented in the robot joint space. therefore the state s = q E R", and the ac- tion E IR? correspond to the joint position and velocity, respectively. The shaded standard deviation over 15 runs.
Figure 12:
0.8
success rate
WGF
3D narrow-path task per-
0.6
PMOE-PPO
PMOE-SAC
0.2
indicates
0.0
20k
40k
60k
80k 100k 120k 140k
environment steps
20000 environment S Moreover, the solution variance of our method also very low, which is consistent with our observations concerning the performance of our policy optimization on the three planar tasks analyzed in the main paper. Figure [2 shows the results of the collision avoidance task using the joint space representation. As observed, our approach was able to adapt the robot motion policy so that the robot end-effector safely passes through a narrow path defined by two spherical obstacles (the narrow-path task description is provided in Sec. 4.1) It is evident that our approach outperformed the baselines in this simulated robotic task, providing evidence that our approach scales and it is able to adapt robot motion policies in higher-dimensional tasks.
A.6.4 Initial GMM policies
For the sake completeness, Fig. provides 2D projections the initial GMM policies learned from demonstrations for the three robotic settings considered in the main paper: the reaching motion skill, the collision- free trajectory tracking, and the multiple-goal task. Figure [3 also provides the demonstration data used to train the initial policies. Note that these models are then adapted according to the policy optimization approach introduced in 3.2
As the baselines underperformed in the 2D case, they tested in this specific setting.
22

---

10.0
5.0
-10
-2.5
-5.0
-10
-7.5
-10.0
0.0
7.5
-5.0
-2.5
0.0
2.5
5.
75
10.0
-50
(a) Reaching motion skill
10.0
30
7.5
20
10
-0.0
-7.5
-5
-10.0
15.0-
12.3
10.0
7.5
5.0
-25
0.0
2.5
5.0
20
-15
-10
-5
10
15
(b) Collision-free trajectory tracking
-1
-6
-8
C
(c) Multiple-goal task
Figure 13: Green Gaussian components represent the initial GMM policy learned from demonstrations, projected the Cartesian position (left) and velocity (left) spaces. The recorded position and velocity data depicted as black dots (+).
23


Reviews:
[{"review_id": "ZYzr1Sx3G5", "reviewer": "Reviewer_zWwA", "summary": "This paper proposes an algorithm for RL in continuous action spaces based on the Wasserstein Gradient Flow formulation of Richemond and Maginnis. By specializing to the case of Gaussian Mixture policies, a simpler formulation is obtained, where the Gaussian part of the parameters (mean and covariance for each mixture component) are learned by optimizing a closed-form quadratic loss term rather than having to compute gradients of the Wasserstein_2 distance numerically.", "questions": "What can the authors do in order to improve empirical evaluation ? I would be willing to raise my score if some of the above concerns were addressed.", "limitations": " ", "rating": 7, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "Just like in Zhang et al. although taking another route, Wasserstein_2 distance terms simplify to L2 distances thanks to the Gaussian assumption. In itself the algorithm represents a straightforward application of the principles described in Richemond & Maginnis or Zhang eet al., all the way to the use of the Sinkhorn algorithm, although the optimization part, framed as a splitting operation over two Riemannian gradient descent problems, is nicely executed.\n\nThe paper is well written and easy to follow, even if several typos still remain ('Wassertein', 'WFG' instead of WGF...). Empirical results show that the Wasserstein Gradient Flow outperforms non optimal-transport motivated baselines. ", "weaknesses": "However, we do have several concerns regarding the scope and significance of these results. These span three axes. First, the PPO and SAC comparison baselines, while standard and still competitive, are 2017 or 2018 algorithms and would deserve updating. Second, only three robotics environments are considered, and it would be good to include more, in particular standard control domains. Finally, the Future Work section of the paper mentions multiple conceptual possibilities for improving and modifying Algorithm 1, as details such as the initialization of the Sinkhorn algorithm used for optimizing the mixture parameter can be absolutely critical in practice. We encourage the authors to implement some of those ideas. For these reasons, and being cognizant of the associated compute requirements, we feel the paper would much benefit from a revamped and extended empirical section implementing these ideas and testing them across a wider set of environments.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank you for taking the time to review our work. We are delighted to read that \"our splitting operation over two Riemannian gradients is nicely executed\"! Below, we address some of the key concerns raised as part of the review.\n\n1. **Baselines**:\n    We would like to point out that PPO and SAC continue to uphold their status as prominently competitive methods in RL, even when applied in the context of robot control policies. For example, [R1] reports a very exhaustive evaluation of standard RL algorithms where PPO and SAC are often among the top 3 best performing methods. A similar observation was made in [R2]. \n    Moreover, it is noteworthy to emphasize that the competitive performance displayed by PPO and SAC renders them enticing choices for real robotic manipulation tasks [R3]. \n    Therefore, even though these methods have their origins dating back 5-6 years, their competitive performance establishes them as foundational baselines in our work.\n    Nevertheless, following the concern raised by this reviewer and the suggestion made by reviewer **vzrR**, we added a new baseline: PMOE [R4], which is a method to train deep RL policies using a probabilistic mixture of experts via GMMs (see Figs. 1 and 2 in the attached PDF, the former is an updated version of Fig. 3 of the original paper). \nNote that *our approach also outperforms both versions of this newly added baseline*, which provides additional experimental evidence on the importance of considering the geometric structure arising from the space of GMMs and the associated Bures-Wasserstein metric on the GMM parameters. \n\n\n   [R1]. [F. Helfenstein, Benchmarking Deep Reinforcement Learning Algorithms. MSc. Thesis, 2021.](https://www.ias.informatik.tu-darmstadt.de/uploads/Team/DavideTateo/felix_thesis.pdf)\n    \n   [R2]. [MushroomRL: Simplifying Reinforcement Learning Research, JMLR, 2021.](https://jmlr.org/papers/volume22/18-056/18-056.pdf)\n    \n   [R3]. [Continuous control actions learning and adaptation for robotic manipulation through reinforcement learning. Autonomous robots, 2022.](https://link.springer.com/content/pdf/10.1007/s10514-022-10034-z.pdf?pdf=button)\n    \n   [R4]. [Ren J, et al. Probabilistic mixture-of-experts for efficient deep reinforcement learning. arXiv, 2021.](https://arxiv.org/pdf/2104.09122.pdf)\n\n2. **Robotics environments**:\n    In our paper we reported three different 2D robotic tasks that reflect the most common robot motion generation problems: reaching, obstacle avoidance, and multimodality. Moreover, we also conducted an additional experiment with a simulated 7-DoF robotic manipulator that learned a collision-avoidance behavior in the 3D Cartesian space of the end-effector (see Appendix A.6.3 for details).\n    However, in the spirit of providing further experimental studies as suggested by this reviewer, we tested the same collision-avoidance behavior in the 3D narrow-path setting reported in Appendix A.6.3, with the difference that the robot motion skill is learned in the robot joint space (i.e. the state $\\mathbf{s} = \\mathbf{q} \\in \\mathbb{R}^7$ and the action $\\mathbf{a} = \\dot{\\mathbf{q}} \\in \\mathbb{R}^7$).\nIn this case, the simulated robot was controlled via a joint velocity controller at a frequency of 100Hz.\n    This experiment is aimed at *assessing the capabilities of our approach to adapt robot motion policies in state-action spaces of higher dimensions*. \n    As shown in Fig. 2 in the attached PDF, our approach is able to adapt the robot motion policy so that the robot end-effector safely passes through a narrow path defined by two spherical obstacles (the narrow-path task description is the same as in Sec. 4.1 of the original paper). \nAs observed, it is clear that our approach outperforms all the baselines in this simulated robotic task, providing evidence that our approach scales and it is able to adapt robot motion policies in higher-dimensional tasks.\n\n\n3. **Future work ideas**:\n    We agree with the reviewer on the aspect that the suggested future work directions may improve the performance of our proposed method, although it is worth highlighting that most of them are algorithmic improvements. \n    However, due to time constraints, we decided to focus on improving our comparison against new baselines and on showing the performance of our method in higher-dimensional settings, as reported in the new figures shown in the attached PDF. \n    Note that the newly added experiments support our findings reported in the original paper: leveraging the geometric structure of the probabilistic mixture of Gaussians provides higher success rate and low-variance solutions. \n    \n    \n**Questions**:\n\n1. **Improvement on empirical evaluation**:\n\n    As discussed previously, we added two additional baselines to our comparison and tested our approach in a higher-dimensional robotic setting. \n    These new experiments align with our findings discussed in the original paper, providing strong evidence on the importance of accounting for the geometry of the policy structure in the formulation of policy optimization methods. "}, {"review_id": "dnKmSg2hFi", "reviewer": "Reviewer_UKYq", "summary": "The authors investigate Wasserstein gradient flows (WGF) for a Gaussian mixture model policy. \nThe WGF is a principled natural gradient method for updating the parameters since it follows the metric space defined by the Wasserstein-2 divergence. \nThis approach is compared to deep RL methods on some planar control tasks.", "questions": "1) For the 2D experiments, if the policies are initialized from demonstrations, why is this success rate 0? Are the baselines also initialized with the demonstrations? Also, consider using the *-from demonstrations variant of RL algorithms (e.g. SACfD).\n\n2) I was a bit confused by the action space for the deep RL baselines. It seems the SAC-GMM paper focuses more on a hierarchical RL set up, where the SAC policy defines a GMM policy that runs for several timesteps. There is no reason why you could not have just replace the PPO and SAC policies with the GMR policy, since PPO just needs samples and log probabilities and SAC just needs to do the reparameterization trick.\nThe issue with SAC would be that the max entropy regularization would be too powerful for GMR policies (since the SAC policy is usually clamped to upper-bound the entropy) so you would have to use the KL-regularied version of SAC, e.g. see Marino et al. \n\nIterative amortized policy optimization, Marino et al NeurIPS 2021.\n", "limitations": "The focus on finetuning from demonstration is understated in the title and abstract. It's not clear from the experiments if this method + policy can be used to train a policy from scratch.", "rating": 4, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 1, "strengths": "This was a very nicely written paper with also excelled presentation. I commend the authors for this. \nI also think WGF is an interesting research direction for RL.", "weaknesses": "**Ignores natural gradients in RL**\n\nI was very surprised to see this paper almost completely ignores the work on natural gradient methods for RL. While TRPO is cited, it is only very vaguely with no mention of the KL natural gradient it uses. \nHere are some key references:\n\n[1] Amari, S. I. (1998). Natural gradient works efficiently in learning. Neural computation\n\n[2] A natural policy gradient, Sham M Kakade - Advances in neural information processing, 2001\n\n[3] Natural actor-critic J Peters, S Schaal - Neurocomputing, 2008\n\nRelated to this is the linear programming view of RL, which also uses KL regularization (but uses EM rather than NG) \n\n[4] Relative entropy policy search, J Peters, K Mulling, Y Altun AAAI 2010 \n\n[5] A unified view of entropy-regularized markov decision processes, G Neu, A Jonsson, V Gómez, 2017\n\nand mirror descent as well\n\n[6] Mirror Descent Policy Optimization, Manan Tomar, Lior Shani, Yonathan Efroni, Mohammad Ghavamzadeh, 2020\n\nThese methods are basically Equation 5 of this paper, but with the KL rather than the W_2 divergence. The interesting question for me would be how these two divergences compare in the context of policy optimization. This seems to be central in the recent paper by Moskovitz et al but completely absent in this work. I think there has been a lot of work on doing NG methods for GMMs, e.g.\n\nHandling the Positive-Definite Constraint in the Bayesian Learning Rule, Lin et al 2020.\n\nThis criticism naturally extends to the choice of baselines in the experiments. There should definitely be a natural policy gradient and EM-based method (for example MPO, advantage weighted regression or actor critic) to compare the quality of the update. \n\nI believe this is quite a big issue as it would require a major overhaul of the paper and experiments.\nThe value of the contribution of the paper drops significantly if you cannot place WGF relative to prior work on natural gradients in RL.\n\n**Limited relevancy of GMR policies for NeurIPS**\n\nThis paper focuses exclusively on GMM  / GMR policies, which would also be called 'locally weighted regression' (LWR) back in the day. These models are rarely used because they scale very poorly to high dimensions, because you need either need hierarchy, dimensionality reduction or many many mixtures (e.g. [7]).\n\n[7] Locally Weighted Projection Regression, Vijayakumar et al ICML 2000.\n\nThe cited works for GMR used them for learning from demonstrations for low-dimensional tasks.\nAs a consequence, I feel like this work would be better suited for venues such as CoRL or ICRA where these methods may be still used, unless this method can be extended to use WGF  on models more relevant to the Neurips community.\n\n**Misc**\n\nTypo: Equation 7 is missing a square root on the right-most squared Bures metric term.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your review. We are very pleased to read that the reviewer found our paper \"very nicely written paper with also excelled presentation\"! Below we address the issues raised in the review.\n\n1. **Natural gradients (NGs) in RL**:\n    * NGs leverage a metric to locally control the gradient-based updates of the policy parameters. In this sense, we understand that the reviewer identifies a connection with our approach as it leverages the geometric structure of the Wasserstein space of GMMs. A similar connection may arise when analyzing the NG as an approximation to the implicit optimization scheme. \n    Nevertheless, we would like to point out that: (1) our SOTA revision focused on approaches that exploited the *Wasserstein* geometry of the policy distribution (lines 86-96); (2) our paper did revise the only work (to the best of our knowledge) that exploits the Wasserstein natural gradient in RL problems [36], which focuses on the differences between the KL-div and the Wasserstein distance for NG updates in RL, but *without assuming a specific policy structure*.\n    * In our paper we already discussed in Appendix A.4 the connections between the forward and backward discretization in the Bures-Wasserstein (BW) metric, where we explained that the Wasserstein NG, according to the BW metric in Eq. (39), is an approximation of the exact Riemannian gradient descent in Eq. (40), which is derived from a 1st-order approximation of the geodesic. As noted in Appendix A.4, such approximation has no guarantees that the approximated geodesic stays on the manifold, as it is also brought up in [R1], a paper suggested by this reviewer.\n    * As our problem involves policy updates with, for example, positive-definite constraints arising from the covariance matrices, using NG is mathematically flawed as it does not guarantee that the updates stay on the underlying manifold. Note that the main reason why the method in [R1, Sec 5.3 \" Our Rule as an Inexact RGD Update\"] still used an inexact Riemannian gradient descent is the difficulty of computing the exponential map (or geodesic) necessary to calculate the exact Riemannian update. In sharp contrast, we leverage the Riemannian operations over positive definite matrices identified by the BW geometry as proposed in [49], thus avoiding to use inexact Riemannian gradient descent updates or approximations of the Riemannian retraction operator. In such a way, we also avoid computing numerical estimates for the Fisher information matrix, which is known to break the invariance-to-parametrization properties of NG methods [R2]. \n    * On the experimental front, our paper reported an ablation study in Appendix A.6.2, where we computed the GMM parameter updates following the implicit Euler scheme (Eqs. (41) and (42)). It is worth highlighting that the *Wasserstein NG* approximates such an implicit Euler scheme, as explained in Appendix A.4 (Eqs. (37)-(39)) and detailed in [47]. This ablation study experimentally showed that our exact Riemannian formulation outperforms non-Riemannian methods, suggesting that a proper treatment of the geometry arising from the GMM structure is advantageous over other non-Riemannian approximations. \n    * Although we are aware that part of our discussion appears in the Appendix of the paper and that we did not stress it and refer to it enough in the main text, we believe that our paper did not ignore the relevance and connections between NG and our approach, as suggested by this reviewer.\n\n      [R1] Handling the positive-definite constraint in the bayesian learning rule, ICML, 2020. \n\n      [R2] New insights and perspectives on the natural gradient method, JMLR, 2020.\n\n2. **GMM-based policies relevance**:\n    We respectfully disagree with the reviewer's statement claiming that \"GMM-based models are rarely used because they scale very poorly to high dimensions...\". Note that: \n    * Recent works have shown that it is possible to efficiently train GMMs based on stochastic gradient descent, therefore scaling better to high-dimensional settings (e.g., 30000 dimensions) [R3], which shows that it is possible to use such models in problems of higher-dimensionality.\n    * When learning robot motion policies, the curse of dimensionality often arises in perception modules (e.g., images), which usually represents an observation of the state of the task. Therefore, it is common practice to employ deep NNs to learn *low-dimensional* embeddings representing the task state, which is later used to train robot motion policies. Note that GMM-based policies may also leverage this kind of perception backbones, and thus the corresponding state of the task may still belong to a relatively low-dimensional space. \n    * Note that reviewer **vzrR** brought to our attention the PMOE method, which employs a probabilistic mixture of experts, represented by a GMM, for deep RL settings, and it shows that certain policies are better learned and represented via a mixture of experts. This is just another example of the relevance of such kind of representations. Also, works like [24] show that GMM-based methods are still relevant in the ML community.\n    \n      [R3] Gradient-Based Training of Gaussian Mixture Models for High-Dimensional Streaming Data, Neural Proc Ltrs, 2021.\n\n**Q1**:\nSuccess is measured w.r.t to the new task requirements, not w.r.t the task previously learned from demonstrations. All methods started with the same GMM policy learned from demonstrations. Note that neither the baselines nor out method used training methodologies such as experience replay, demonstrations buffer, etc. Such algorithmic improvements are out of the scope of this paper.\n\n**Q2**: In our understanding, SAC-GMM is not a hierarchical RL method as it is basically a SAC method whose actions correspond to increments applied to the GMM-based policy parameters."}, {"review_id": "b41kWpKCEv", "reviewer": "Reviewer_uu8b", "summary": "This paper proposes to formulate policy optimization as a Wasserstein gradient flow over the Gaussian Mixture Model (GMM) space, which enhances the stability of policy optimization processes. In the proposed GMM policy updates, the mean and variance of gaussians are optimized through Riemann gradient descent via the Bures-Wasserstein metric, and the weight of gaussians are optimized through the implicit Euler scheme. The paper then demonstrates the effectiveness of their proposed approach by conducting experiments on various robotic tasks in simulation.", "questions": "See \"weaknesses\".", "limitations": "Limitations are mostly addressed. Another limitation to include is regarding the current empirical evaluations on low-dimensional robotic tasks. Showing the potential for the proposed approach to scale up to higher dimensional tasks with multimodality natures would significantly enhance the paper.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Overall the paper is well-written and well-presented. Discovering more principled approaches to optimize GMM policies to model more complex policy distributions is a meaningful research direction.", "weaknesses": "My main concern lies in the quality of experiments and empirical evaluations. Specifically,\n\n- The three robotic tasks illustrated in the main experiments are rather simple tasks with low-dimensional observation and action spaces. It would be interesting to demonstrate whether the proposed approach can generalize to higher dimensional tasks that exhibit multimodality solutions.\n\n- More in-depth analysis on why the proposed approach outperforms the baselines (PPO+GMM & SAC+GMM) can be included. Specifically, it would be meaningful to qualitatively compare policy distributions between the proposed approach and the baselines, particularly for the observations where the proposed approach is better than the baselines in modeling policy distributions. This would significantly enrich the  insights provided by the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are very grateful for your time reviewing our work and the provided suggestions. We appreciate the positive feedback about our paper, and are glad to read that our approach \" to optimize GMM policies to model more complex policy distributions is a meaningful research direction\"!  Below we address the issues raised in the review.\n\n1. **Complexity of robotic settings**:\n   * First, we would like to point out that the three 2D robotic tasks cover common adaptation problems in robot motion policy optimization, namely, (1) adapting a reaching motion to a different target, (2) adapting in the presence of obstacles, and (3) adapting a multimodal task. We acknowledge that these experiments did not fully show the capabilities of our approach. To achieve this, we already provided a first evaluation of our method in the narrow-path task using an off-the-shelf 7-DoF robotic manipulator in simulation (see App. A.6.3). \n    * For this rebuttal, we conducted an additional experiment on the simulated 7-DoF robotic manipulator in the narrow-path task. In this case, the distinction lies in the acquisition of the robot's motion skill, which was learned in the space of robot joint configurations, where the state was $\\mathbf{s} = \\mathbf{q} \\in \\mathbb{R}^7$, and the action corresponded to $\\mathbf{a} = \\dot{\\mathbf{q}} \\in \\mathbb{R}^7$. We controlled the simulated robotic arm via a joint velocity controller at a frequency of 100Hz. The objective of this experiment was to evaluate our approach performance in adapting robot motion policies within state-action spaces of higher dimensions. As illustrated in Figure 2 within the attached PDF, our method effectively adapts the robot's motion policy, thus ensuring a collision-free skill execution. Note that the narrow-path task specification aligns with the details outlined in Section 4.1 of the original paper. The observed outcomes distinctly demonstrated our approach's superiority over all benchmarks in this simulated robotic task. This provides further evidence that our approach scales and is proficient in the adaptation of robot motion policies within tasks involving higher dimensions on off-the-shelf robots.\n    \n\n2. **Additional analysis**:\n    Following the suggestions made by this reviewer and other reviewers, we added two baselines to our evaluations, namely, PMOE-PPO and PMOE-SAC, which were proposed in [R1]. As in our approach, PMOE assumes that the policy $\\pi$ is represented by a mixture of Gaussians, and it addresses the problem of estimating the gradients needed to optimize the mixture weights. However, PMOE does not take an optimal transport perspective on the policy optimization problem, nor does it account for the geometry arising from the GMM parameters. The results reported in Fig. 1 of the attached PDF show that our approach also outperforms these two newly added baselines, providing further evidence on the importance of considering the geometry of the GMM space. \n\n    [R1] [Ren J, et al. Probabilistic mixture-of-experts for efficient deep reinforcement learning. arXiv, 2021.](https://arxiv.org/pdf/2104.09122.pdf)\n\n    \n**Limitations**:\n\nAs discussed above, we added an additional experiment showcasing our approach performance in adapting a motion policy in the narrow-path task where the GMM-based policy is learned in a $14$-dimensional state-action space. Moreover, we tested 4 different baselines in the same setting. The results reported in Fig 2. showed that our approach scales to higher-dimensional problems and still outperforms baselines that disregard the geometric structure arising from the policy representation. "}, {"review_id": "nNxtWJ3DiQ", "reviewer": "Reviewer_vzrR", "summary": "The paper proposes a new approach to adapt robot motion policies for different task conditions. It suggests leveraging the structure of probabilistic policies, specifically Gaussian mixture models (GMMs), and formulating policy optimization as an optimal transport problem. By using the L2-Wasserstein distance between GMMs, the policy updates can be constrained to improve the stability of the optimization process. Additionally, the Bures-Wasserstein manifold geometry is utilized for optimizing the Gaussian distributions of the GMM policy through Riemannian optimization. The proposed method is evaluated on various robotic scenarios and demonstrates better performance in terms of task success rate and low-variance solutions compared to common policy optimization baselines.", "questions": "N/A", "limitations": "N/A", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Novel Approach: The paper introduces a new approach to adapt robot motion policies by leveraging the structure of probabilistic policies and formulating policy optimization as an optimal transport problem. This novel perspective can provide insights into enhancing the adaptability and performance of robot motion policies.\n\nUtilization of Gaussian Mixture Models (GMMs): The paper focuses on GMMs, a widely used representation for modeling complex motion policies. By exploiting the specific structure of GMMs, the proposed method offers a tailored solution for policy optimization, which can potentially lead to more effective and efficient adaptation of robot motion policies.\n\nConsideration of Stability: The incorporation of the L2-Wasserstein distance between GMMs as a constraint in policy updates aims to enhance the stability of the optimization process. This consideration addresses a common challenge in policy optimization algorithms and can contribute to more reliable and consistent results.\n\nRiemannian Optimization: The paper leverages the geometry of the Bures-Wasserstein manifold for optimizing the Gaussian distributions of the GMM policy. This utilization of Riemannian optimization techniques showcases a sophisticated mathematical framework to refine the policy parameters, potentially leading to improved performance and convergence properties.\n\nExperimental Evaluation: The proposed method is thoroughly evaluated on common robotic settings, including reaching motions, collision-avoidance behaviors, and multi-goal tasks. The results demonstrate that the proposed approach outperforms common policy optimization baselines in terms of task success rate and low-variance solutions. This empirical validation strengthens the credibility and practical relevance of the proposed method.", "weaknesses": "I noticed that this does not discuss much other GMM-based methods, such as the PMOE[1]. Including a discussion and comparison with existing GMM-based methods would greatly enhance the comprehensiveness and value of your work. Furthermore, I suggest you consider conducting additional experiments to compare your proposed approach with these methods. Because comparing only with non-GMM methods, like PPO and SAC, is not fair.\n\n[1] Ren J, Li Y, Ding Z, et al. Probabilistic mixture-of-experts for efficient deep reinforcement learning[J]. arXiv preprint arXiv:2104.09122, 2021.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are very grateful for your time reviewing our work and the provided suggestion to compare against PMOE, an algorithm we were not aware of. We also appreciate the positive feedback about both the theoretical and practical aspects of our paper. Below we address the main concern of the review:\n\n1. **Discussion and comparison against PMOE**:\n    * In a similar spirit as our paper, PMOE highlights the importance of leveraging a mixture of experts in the policy structure of deep RL methods. However, using GMMs straightforwardly in generic off-policy and on-policy deep RL algorithms introduces a difficulty in their end-to-end training process. This is due to the inherent non-differentiability caused by the optimization of categorical distribution parameters within the GMM optimization procedure. PMOE addresses this problem by proposing a gradient estimator for optimizing the mixture weights. \nAlthough PMOE also assumes that the policy $\\pi$ is represented by a mixture of Gaussians, *its formulation does not provide an optimal transport perspective to the RL problem as in our paper*. PMOE mainly focuses on formulating feasible gradient updates for both the mixture weigths and Gaussians' parameters, but it does not consider their underlying geometry. \nIn sharp contrast, our approach formulates the policy optimization as a gradient flow according to the Wasserstein distance, which leverages the view that *the set of mixtures of Gaussians can be associated with a Wasserstein metric*, leading to a policy optimization with updates following a Wasserstein gradient flow. Moreover, our approach *exploits the fact that a Gaussian can be identified with the Bures-Wasserstein manifold*, which corresponds to the product manifold $\\mathbb{R}^d \\times \\mathcal{S}_{++}^d$. This allows us to formulate an explicit Euler scheme for the *gradient updates that builds on Riemannian optimization*. This means that *our update rule does not depend on inexact Riemannian gradient updates* (like natural gradients), and therefore *our method guarantees that the gradient flow stays on the underlying Riemannian manifold*. \nThe PMOE gradient-based updates [R1, Secs. 3.3, 3.4] do not provide such guarantees. \n\n    * As suggested by the reviewer, we added PMOE [R1] as an additional baseline for our benchmarking experiments. Specifically, we added the PMOE version for both PPO and SAC (i.e., PMOE-PPO and PMOE-SAC). \n    Our implementation is based on the [code](https://github.com/JieRen98/rlkit-pmoe/tree/master/rlkit/torch/PMOEsac) provided by the first author of the paper [R1], and includes several minor changes to comply with recent versions of the code dependencies. \n    The figures in the attached PDF include the results corresponding to the two newly added PMOE methods. \n    In general, PMOE often showed a better performance than the vanilla PPO and SAC baselines, which follows the experimental insights provided in the PMOE paper [R1] regarding the advantages of the probabilistic mixture of experts in deep RL settings. \n    Nevertheless, as shown in the figures, our approach also outperforms both PMOE versions in all three 2D robotic tasks and in the newly added simulated task aimed at testing all the methods in a higher-dimensional setting.\n    The new results corroborate the observations presented in our paper: leveraging the geometric structure of the probabilistic Gaussian mixture yields improved success rates and solutions with reduced variance.\n\n      [R1] [Ren J, et al. Probabilistic mixture-of-experts for efficient deep reinforcement learning. arXiv, 2021.](https://arxiv.org/pdf/2104.09122.pdf)\n\nWe will include the above discussion and the updated figures (appearing in the attached PDF) to the revised version of our paper. "}]

Author rebuttal:
We would like to thank the reviewers and the area chair for their work and feedback on our manuscript. Below we summarize the main points addressed in our rebuttal. We also attach a PDF with two figures showing newly added results for additional experiments as requested by the reviewers. 

**Rebuttal summary**
1. Following the recommendations from reviewers **zWwA**, **uu8b**, and **vzrR**, we performed the following new experiments:
   * Tested our approach (and baselines) on a new simulated robotic task of higher state-action dimensionality. Our results show that *our method is able to adapt the learned policy and outperforms all baselines providing lower-variance solutions and faster convergence*. 
   * Added two new baselines (i.e., PMOE-PPO and PMOE-SAC) based on the method suggested by reviewer **vzrR**. The results show that *our approach still outperforms all methods* despite that the PMOE-based formulations provided a performance gain when compared to vanilla PPO and SAC versions. 
2.  We provided a thorough clarification about the position of our paper w.r.t natural gradient approaches:
    * We clarified that *we did include in our original submission the only paper (to the best of our knowledge) that uses the Wasserstein natural gradient in RL*. As the core of our paper is on Wasserstein gradient flows, we omitted the literature on KL divergence. 
    * We indicated that our original submission included a short discussion (with equations) in App. A.4 to explain the connection between our Riemannian optimization and the approximation provided by an inexact natural gradient approach.
    * We pointed out that our original submission included an ablation study in App. A.6.2 where we compared our method against an optimization solving the implicit Euler scheme (Eqs. 41 and 42), for which a natural gradient method is an approximation. 

3. We discussed the importance of GMM-based policies, their wide use in robotic applications, and the recent efforts in the machine learning community to scale GMMs to very high-dimensional spaces. 
 
We trust that our rebuttal has provided clear explanations for all the questions raised by the reviewers. Finally, we wish to express our advanced appreciation to the reviewers for their time in reviewing our rebuttal and for their forthcoming involvement in the discussion phase.

Meta-review: All the reviews appreciated the innovative integration of OT with GMM, the clarity of writing, and the well-crafted optimization section. The authors' rebuttal further strengthened the paper, particularly concerning the numerical validation. In terms of comparing their approach with natural gradients, the authors at least partially addressed this in their exchange with UKYq, and I encourage them to take this seriously into account when preparing the final version.
Decision: Accept (poster)