Wasserstein Gradient Flows for Optimizing GMM-based Policies
Authors: Hanna Ziesche
Leonel Rozo
Venue: Submitted to ICLR 2023 (2023)
Keywords: Optimal transport
policy optimization
gaussian mixture models
robot motion adaptation.

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

Under review as a conference paper at ICLR 2023
WASSERSTEIN GRADIENT FLOWS FOR OPTIMIZING GMM-BASED POLICIES
Anonymous authors Paper under double-blind review
ABSTRACT
Robots often rely on a repertoire of previously-leamed motion policies for per- forming tasks of diverse complexities. When facing unseen task conditions or when new task requirements arise, robots must adapt their motion policies accord- ingly. In this context, policy optimization is the de facto paradigm to adapt robot policies function of task-specific objectives. Most commonly-used motion policies carry particular structures that are often overlooked in policy optimiza- tion algorithms. We instead propose leverage the structure of probabilistic poli- cies by casting the policy optimization as an optimal transport problem. Specif- ically, focus robot motion policies that build Gaussian mixture models (GMMs) and formulate the policy optimization as a Wassertein gradient flow over the GMMs space. This naturally allows
constrain the policy updates via the - Wasserersnnin distance between GMMs to enhance the stability of the pol- icy optimization process. Furthermore, 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 set- tings: Reaching motions, collision-avoidancee 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.
INTRODUCTION
One of the main premises about autonomous robots is their ability to successfully perform large range of tasks in unstructured environments. This demands robots adapt their task models accord- ing to environment changes, and consequently to adjust their actions to successfully perform under unseen conditions (Peters et al. 2016). In general, robotic tasks, e.g. picking or inserting an object, usually executed by composing previously-eeame by a motion policy. Therefore, order to successfully perform under new settings, the robot should adapt its motion policies according the new task requirements and conditions.
skills (Schaal et al. 2003), each represented
Research on methods for robot motion policy adaptation is vast (Kober et al. [2013] Chatzilygeroudis al., 2020), with approaches mainly building black-box optimizers (Stlpp & Sigaud (0122) end-to-end deep reinforcement learning (Tbarz et al. (021), and policy search (Deisenroth et al, 2013). Regardless of the optimization method, most approaches rely on policy structure-unawarc adaptation strategies. However, several motion policy models (e.g, dynamic movement primitives (DMP)([jsperrt al. 2013), Gaussian mixture models (GMM) (Calinon al. 2007), probabilistic movement primitives (ProMPs) (Paraschos et al. 2018), and neural networks (Bahl et al. 2020), among others), carry specific physical or probabilistic structures that should not be ignored. First, these policy models are often learned from demonstrations in a starting learning phase (Schaal et al.) 2003), thus the policy structure already encapsulates relevant prior information about the skill. Sec- ond, crructure-unwware adaptation strategies optimize the policy parameters disregarding the special characteristics of the policy model (c.g., a DMP represents a second-order dynamical system). In this regard, we hypothesize that the policy structure may be leveraged to better control the adaptation strategy via policy structure-awar gradients and trust regions.
Our main idea is to design policy optimization strategy that explicitly builds particular policy structure. Specifically, we focus on GMM policies, which have been widely used to learn motion
skills from human demonstrations (Calinon al. 2007 Cederborg al. 2010 Calinon 2016

---

Under review as a conference paper at ICLR 2023
M GMM4
Adapted GMM policy
Initial GMM policy
target
Illustration our policy structure- aware adaptation of GMM policies. Policy updates follow Wasser- stein gradient flow on the manifold of GMM policies GMMa.
Figure 1:
Jaquier et al.
2019p.
GMMs provide simple but expressive enough representation for learning a
large variety of robot skills: Stable dynamic motions (Khamsari-Zaue Zadeh & Billard 2011: Ravichandar
ctt al. 20177 Figueroa & Billard 2018), collaborative behaviors Ewerton et al. (2015): S) Rozo et al.
(2016), and contact-rich manipulation Lin et al. (2012); Abu-Daaka et al. (2018), among others.
Often, skills learned from demonstrations need to be refined due to imperfect data or adapted to
comply with new task requirements. In this context, existing adaptation strategies for GMM policies either build kernel method on top of the original skill model Huang et al. (2019), or leverage reinforcement learning (RL) to adapt the policy itself (Arenz al. 2020, Nematollahi al. 2022) However, none of these techniques explicitly considered the structure of the GMM policy.
Unlike the aforementioned approaches, we propose a policy optimization technique that explicitly considers the underlying GMM structure. To do so, we exploit optimal transport theory (Santambro-
gio 2015 Peyré & Cuturi 2019), which allows us to view the set of GMM policies as particular
space of probability distributions GMM2. Specifically, we leverage the idea of Chen et al. (2019) and Delon Desolneux (2020) view a GMM as a set of discrete measures (dirac masses) on the space of Gaussian distributions G(R)), which is endowed with 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 $3). where the policy updates are naturally guaranteed to be GMMs. Moreover, we take advantage the geometry of the Bures- Wasserstein manifold to 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: Reach- ing motions, collision-aviidance behaviors and multi-goal tasks (see 4. 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 Maginnis (0017) pioncered the idea of understanding policy op- timization through the lens of optimal transport. They interpreted the policy iteration as gradient flows by leveraging the implicit Euler scheme under Wasserstein distance (see $ D considering only 1-step return settings. They observed that the resulting policy optimization resembles the gradi- ent flow of the Bokker--Plance equation (JKO scheme) (Jordan et al.) 1996). In a similar spirit, [Zhang et al. (2018) proposed to use WGFs to formulate policy optimization as sequence of policy updates traveling along a gradient flow on the space of probability distributions until convergence. To solve the WGF problem, the authors proposed a particle- basscll--based based algorithm to approximate continuous den- sity functions and subsequently derived the gradients for particle updates based on the JKO scheme. Although Zhhang et al. (2018) considered general parametric policies, their method assumed dis- tribution over the policy parameters and did not consider a specific policy structure, which partially motivated their particle-based approximation. Recently, Mokrov et al. (202I) tackled the compu- lational burden of particle methods by leveraging input-convex neural networks to approximate the WGFs computation. They reformulated the well-known JKO optimization [ordan al. (199)) over probability measures by an optimization over convex functions. Yet, this work remains a general solution for WFG computation and did not address its use for policy optimization problems.
Aside from optimal transport approaches, Arenz et al. (2020) 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 adapta- tion setting, nor does it consider the GMM structure from an optimal transport perspective. Ne- matollahi et al. (2022) proposed SAC-GMM, a hybrid model that employs the well-known SAC algorithm (Haamoja
al.) 2018) to refine dynamic skills encoded by GMMs. The SAC policy was designed to learn residuals on a single vectorized stack of GMM parameters, thus fully dis- regarding the GMM structure and the geometric constraints of its parameters. Finally, two recent
2

---

Under review as a conference paper at ICLR 2023
works share our idea of leveraging geometry in policy optimization: First, Riemannian proximal policy optimization for GMMs was proposed by [Wang et al. (2020), where the geometry induced by the GMM parameters was considered in the optimization via Riemannian gradients, similarly to our method. The policy optimization was regularized by a Wasserstein distance to control the exploration-exploitttonn trade-off. However, their method did not formulate the policy optimization as an optimal transport problem, i.c. the policy updates do not follow a WGF, as in our approach, but employed instead a classical non-convex optimization. Second, Moskovitz al. (202 employed the Wasserstein natural gradient to exploit the local geometry induced by the Wasserstein regular- ization of behavioral policy optimization (Pacchiano al. 2020). In contrast, our method exploits the geometry induced by the structure of the space of GMM policies via the Bures-Wasserstein manifold, which naturally guarantees that policy updates stay on GMMa.
2
BACKGROUND
2.1 WASSERSTEIN GRADIENT FLOWS
In Euclidean space a gradient flow is smooth curve x E RR that satisfies the partial differ- ential equation (PDE) (t)
V (t())) for given loss function L : R - RR and starting point
To at Santambrogio 2015 2017 solution can be found straightforwardly by forward discretization, leading to the well-known explicit Euler update scheme 2x
2x AVL(x
where ò denotes the learning rate and x" indicates discretization of the curve rrtt) with discretiza- tion parameter т. Alternatively, we can use a backward discretization, which leads to the following implicit Fuler scheme
| - 2T11 L(x))
1 arg min
2T
Eq. [is is sometimes referred to as Minimizing Movement Scheme and can be used as alternative characterization of a gradient flow.
This characterization is particularly interesting when we need to extend the concept of gradient flows to (oon-Euclidean) general metric settings, since there is no notion of VL in these cases (Santam brogio 2015 Ambrosio et al.) 2005). 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 I.2 Wassersrst space, defined the set of probability measures P(X) separable Banach space
X Panaretos & Zemel 2020) and endowed with the - Wassessrnnnninn distance W2 defined as
W2(1, = ( ( inf /
1/2
|m x2l x2l12d
(2)
(TEI(H)))) xxx
where n1 PX) and 7 P(X²) is defined to have the two marginals y and v.
Generalized Minimizing Movement scheme characterizing gradient flows in the Wasserstein space can be written in analogy to Eq. ass:
(T, + C(r)).
TH arg min
2r
where L is a functional to be minimized on the Wasserstein space and ग. E P((). In the following, we will omit the superscript 7 for notational convenience.
2.2 REINFORCEMENT LEARNING AS WASSERSTEIN GRADIENT FLOWS
Our view of the policy structure-aware optimization builds on the approach outlined by Richemond & Maginnis (0077). which in turn is based on the JKO scheme of Jordan et al.) (1996). They proposed formulation of 1-step RL problems in terms of Wasserstein gradient flows. In particular, they studied the evolution a policy π under the influence of a free energy functional J of the form:
( = KT(T)+ + BH(T)) = - dr((a)))(ss B J.
(4)
3

---

Under review as a conference paper at ICLR 2023
where (x) denotes the inner energy of the system, here determined by the reward r(s, a). More- over, TT( is the entropy of the policy т T(ass), with 8 and a denoting the state and action, re- the usual objective in 1-step RL settings with entropy
spectively. Thus Eq. 4 can be recognized regularization. It is well known that the evolution of probability densities under a free energy of this form is properly described by PDE known the Tokker-Plance equation. |Richemond & Maginnis (217)) exploited the result Jordan et al. (1996), which stated that this evolution can be interpreted as the gradient flow of the functional J in Wasserstein space. This flow characterized by the following minimizing movement scheme
(2 mT) J(T))
T++1 arg min
(5)
2r
which naturally provides iterative updates for the policy и, While Richemond & Maginnis (2017)
considered a 1-step bandit setting, we extend this approach to full multi-step RL problems and learn policies for long-horizon tasks.
2.3 THE W2-SERSERN DISTANCE BETWEEN GAUSSIAN MIXTURE MODELS (GMMs)
In this paper, we consider policies Tσx that build on a GMM structure, i.e., π(x;) Σn (a: p4; E,), where N denotes a multivariate Gaussian distribution with mean μ4; and covariance matrix E, and w; are the weights of the N individual Gaussian components, which are subject to Σ; wi = 1. In the following, we will write f., E 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) (IR but rather focus on WGFs evolving the submanifold of GMMs, that is GMM, C )( Following Chen al. (019); Delon & Desolneux (2220), we can approximately describe this submanifold as a discrete distribu- tion over the space of Gaussian distributions equipped with the Wasserstein metric. This in turn can manifold which is the product manifold Rd x s where St denotes the Riemannian manifold of d-dimensional symmetric positive definite matrices. The corresponding approximated Wasserstein distance between two GMMs m., T2 is given by
be identified with the
(T)
min (WU
Σ P. 3N (.) My E,)), i.j
(6)
where U(wi, (P I WI, P w2} with 1N denoting an N-dimensional vector of ones. The Wasserstein distance between two Gaussian distributions in Eq. 6 can be com- puted analytically follows
W2 (Ni M4, Ex), N2 My E,)) = 11|u - My tr + + E, 2 ( ) (7)
2.4 LEARNING GMM POLICIES FROM DEMONSTRATIONS
A popular approach in RL - particularly in the roboties domain to reduce the number of policy rollouts in the environment is to warm-start the policy with set of demonstrations provided by an expert. In this work we choose to represent our policy via GMM. We assume that demonstrations are provided as set of trajectories of state-action pairs {(S0. ap), (.), 1.....(e (ST, (sr,arr) ) To initialize our policy, we first use the Sipeetation-Maxxmizattee in the joint state-action space, to the demonstrations. This results in mixture distribution (ss a) = D,,,,, ([N[[s ;4, state 8, as follows
(EM) algorithm to fit a GMM,
p4, E,) from which policy can be obtained by conditioning the
7T (s.aa)
(a 8)
J (s, a))aa
In the context of GMMs, this is also known as Gaussian Mixture Regression (GMR) GGahramani & Jordan 1994). The resulting conditional distribution is another GMM on action space with state- dependent parameters, given by
N
क T 81) = Σ wi (ar) E.).
(9)
Details on computation of Eq. from the original GMM are given in App. A.I
4

---

Under review as a conference paper at ICLR 2023
3 WASSERSTEIN GRADIENT FLOWS FOR GMM POLICY OPTIMIZATION
In this work, we focus on multi-step RL tasks for policy adaptation. We consider a finite-horizon Markov Decision Process (MDP) with continuous state and action spaces S R" and A E I""", transition and reward functions p( Stt az) and r(st, at), initial state distribution p(su) and a discount factor 7. Further, assume to have initial policy I ((a,(s)), which is to be adapted by optimizing some objective function (T). As stated in $ 1 this problem arises in robot learning settings where policy learned via imitation learning (e-g., LfD) needs to be adapted to new objec- tives or unseen environmental conditions. To promote exploration and avoid premature convergence to suboptimal policies, we leverage maximum entropy RL (Eysenbach & Levine 2022) by adding an entropy term H(TT) to the objective. Thus, the overall objective has the form of a free energy functional (resembling Eq. 4 and can be written as
J(x) = (T) + BH(()),
(10)
where B is hyperparameter and K2 (T) corresponds to the usual cumulative return
(T) E Σ r((,, an)
a.) Σ (nrr(s,, (11)
The evolution of the policy T(aa)s)) over the course of the optimization can be described as a flow of probability distribution 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 J'-Wasserstein distance to constrain the policy updates, which is important to guarantee stability in policy optimization (Schulman al. 2015 2017| Otto et al.) 202I]: (iii) By constraining to specific submanifolds of the Wasserstein space, in this case GMMs, can impose additional structural properties on the policy optimization.
Since our objective in Eq. 10 has the form of the free energy functional studied by Richemond & Maginnis (2017); Jordan al. (1996), can leverage the iterative updates scheme of Eq. 5 to formulate the evolution of our policy iteration under the flow generated by Eq. [0] As mentioned previously, we focus on the special case of GMM policies and therefore restrict the Wasserstein gradient flow to the submanifold of GMM distributions GMM,. We refer the interested reader to App. A.3
where we provide the explicit form of (Tz) of Eq. 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 of Gaussian distributions G(R"), equipped with the Wasserstein met- ric (Chen et al. 2019). Consequently, our policy optimization problem naturally splits into op- timization over the (N- didmeennso simplex and an optimization on the N-fold product of the
Bures- Wasserstein manifold (BW*N, i.e. the product manifold (Rd x ) The former corre- sponds to the GMM weights while the latter applies to the set of Gaussian distributions parameters. perform the optimization directly
Note that the identification with the BWN manifold allows on the parameter space. This comes with several benefits: (i) We can leverage the well-known analytic solution of the Wasserstein distance between two Gaussian distributions in Eq. 6 greatly reducing the computational complexity of the policy optimization. (ii) As Chen et al. (2019) show, we can guarantee that the policy optimized via Eq. remains GMM. (iif) Unlike the full Wasser- stein space]] the resulting product manifold is a true Riemannian manifold such that we can leverage the machinery of Riemannian optimization. Importantly, working the parameter space allows us to apply an explicit Euler scheme instead of the implicit formulation of Eq.
According to the above-mentioned split, we formulate the policy optimization as an EM-like 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 as forward discretization, similarly to Chen & Li 2020). In other words, by embedding the Gaussian components of the GMM policy in a Rie- mannian manifold, the Wasserstein gradient flow in the implicit form of Eq. can be approximated
Wasserstein space is not true Riemannian manifold, but it can he equipped with a Riemannian structure and formal calculus this manifold (Oitol 200T), which has been made rigorous by (Ambrosio 2005)
5

---

Under review as a conference paper at ICLR 2023
by explicit Euler update scheme according 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 (Malagó et al. 2018 Han et al. 202I). Thus, the optimization boils down to Riemannian gradient descent where the gradient is defined w.r.t the Bures-Wassersterr metric. In particular, we use the expression for Riemannian gradient, metric and exponential map used in (Han) et al. 2021). Formally, the resulting updates for the Gaussian parameters of the GMM follow the Riemannian gradient descent scheme given by:
Rx, p Ría (X Aggrad, ((+))) and Et+1= RR. (A A.gradæ J(T))),
(12)
where grad denotes the Riemannian gradient w.r.t. the werss-Wasserrteiin metric, Rs Ta.M- + denotes the retraction operator, which maps a point the tangent space Ta- back to the manifold
M
M = BW (Boumal 2022). Moreover, À is a learning rate and TT. def T(β Σx. ax.). The Euclidean gradients of J(xT) required for computing grad can be obtained using likelihood ratio estimator (a.k.a score function estimator or REINFORCE) (Williams) 2004) and are provided in App. A.3
Concerning the GMM weights, we first reparameterize them as wy Σ exp
exp
and optimize
w.r.t the new parameters 7/j.J = I...N, which unlike ws are unconstrained. For this optimization
we employ the implicit Euler scheme:
W2 (TT (f), TT)
+i++1 = arg min A
(()++())))
(13)
27
where (f) def x (fx+11, Ea+1, n).
We minimize Eq. 13 by gradient descent w.r.t 77 as follows:
+ 7 XV; (T T&) J(T++1
(14)
T
The gradient of (TT) can be obtained analytically using a likelihood ratio estimator. For the Wasser- stein term, we first compute the gradient w.r.t the weights via the Sinkhorn algorithm (Cuturi & Doucet (0144), from which the gradient w.r.t 17 can be then obtained via the chain rule. Note that we have to rely on the Sinkhorn algorithm here since there is no analytic solution available for the Wasserstein distance between discrete distributions, unlike the above case of the Gaussian compo- nents. Consequently, we cannot compute the corresponding gradients.
3.2 IMPLEMENTATION PIPELINE
To carry out the policy optimization, we proceed in the usual on-policy RL scheme: We first roll out the current policy to collect samples of state-action-reward tuples. Then, we use the collected interaction trajectories to compute a sample-basee estimate of the functional (x) and its gradi- ents w.r.t the policy parameters, as explained in $ 3.1 An optimization step consists of alternating between optimizing the Gaussian parameters using 12 and updating the weights via 14 For the
optimization of the Gaussian parameters we leverage Pymanopt (Townsend et al. 2016) for Rie- mannian optimization. We extended this library by implementing the Bures- Wasserstein manifold based on the expressions provided by Han et al. (20TT) (see App. A.2 for details). Furthermore, we added a custom line-search routine that accounts for a constraint on the Wasserstein distance between the old and the optimized GMM, to our knowledge such search method does not exist in out-of-the-box optimizers. The details of this custom line-search can be found in Algorithm 2 in use POT (Flamary et al.) 2021) a Python library for optimal transport, from which we obtain the quantities required for computing the gradients of the Wasserstein distance w.r.t. the weights in Eq. 14
App. A.5 Regarding the optimization of the GMM weights,
The full policy optimization finishes if either the objective stops improving or the Wasserstein dis- tance between the old and optimized GMMs exceeds a predefined threshold, which is chosen exper- imentally. Afterwards, fresh rollouts are performed with the updated policy and the aforementioned two-step procedure starts over. This optimization loop repeated until task-dependent success criterion has been fulfilled. We summarize the proposed optimization in Algorithm
4 EXPERIMENTS
We tested our approach in three different robotic settings: a reaching skill, a collision- trajectory tracking, and multiple-goal task. All the tasks represented in 2D operational space. The robot

---

Under review as a conference paper at ICLR 2023
Algorithm GMM Policy Optimization via Wasserstein Gradient Flows
Input: initial policy T{ass)
1: while not goal reached do
2: Rollout policy π T[aass) in the environment for M episodes collect interaction trajectories
I r= [((0, do; ro); (81, 01,
ar, (T) )
3: repeat
4: Update Gaussian components parameters ft. Σ using Riemannian optimization X is determined via line-search (see 3.2).
where
5: until convergence
6: repeat
7:
Update GMM weights @ via gradient descent on the free energy objective 10 using 14
8:
until converged
9:
end while
motion policies were initially learned from human demonstrations collected on a simple Python graphical interface. We assumed were given M demonstrations, each of which contained Tm, data points for dataset of N = = Tm total observations = {($t, ) The state and action a correspond to the robot end-effector position æ E IR² and velocity ø E R². The GMM models were trained via classical Expectation- Catiatiion-Maiiizzttiin.n The policy rollout consists of sampling velocity action a2 ~ TT T(aa)s)) using Eq. and subsequently commanding the robot via a Cartesian velocity controller frequency of 100Hz. For all the experiments, used the Robotics Toolbox for Python (Corke & Haviland, (021) to simulate the robotic environments.
To show the importance of accounting for the policy structure in RL settings, we compared method against two structure- matollahi al. 2022). As PPO was not originally designed directly optimize the parameters of a previously-eeareed GMM policy, we designed the policy actions to represent (usually small) cor- rections to the GMM parameters, i.c. a = [Aw Aú A vec E} following the same methodology as SAC-GMM (Nematollahi et al. 2022). The PPO and SAC implementations correspond to the code provided by Stable- Baselines. (Raffin et al.) (021), whose policies are parametrized by MLP networks. During policy optimization, we sample an action from the MLP policy that is then used to update the GMM parameters by adding the computed corrections to the current parameters. Later, we proceed as described earlier, namely, the updated GMM is used to compute the velocity action via Eq. [ For comparison purposes,
baselines: PPO (Schulman et (017) and SAC-GMM (Ne-
report statistical results for the three considered settings over 5 runs for the task success rate and solution variance. We tuned the baselines separately for each task using Optuna (Akiba al., 2019). In addition, to assess the importance of our Riemannian formulation, we performed an ablation where we used the implicit scheme based on Euclidean gradi- ent descent instead of the explicit optimization on the Bures- Wasserstein manifold (see App.[A.6. [A.6.... Last but not least, we tested our approach on a 3D version of the collision-free task performed by 7-DoF Franka Emika Panda robot in a virtual environment as reported in App. A.6.3
4.1 TASKS DESCRIPTION
Reaching Task: 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- learned L-shape trajectories. The initial policy, shown in Fig. Zteett and Fig. 12a, was learned from 12 demonstrations and encoded by 7-component GMM. To adapt the policy, we defined 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 to significantly divergent trajectories. Convergence is achieved when minimum average position error w.r.t the target computed over an episode- - is reached.
Collision-avoidanee Task: This task consists of: (I) learning an initial GMM policy of a linear target while avoiding to collide with two spherical obstacles located midway between the initial robot po- sition 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. 2 and Fig. [2b For policy optimiza- tion, we defined sparse reward as function of the position error between the robot end-effector
reaching motion, and (2) adapting the GMM policy to reach new horizontally
7

---

Under review as a conference paper at ICLR 2023
Figure 2: The three tested rohotic settings: a reaching skill (leff), a collision-free trajectory tracking (middle), and multiple-goal task (right). The robot color goes from light gray to black to show the evolution of the task reproduction. Green Gaussian components
depict the initial GMM policy, projected on the 2D Cartesian position space. The end-effector trajectory resulting from initial GMM policy is shown in dark blue lines (-). Red circles (O) in the collision-avoidaree task represent the obstacles (middle). The different targets of the multiple-goat task (right) are depicted as red stars.
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. Convergence is determined by minimum average position error w.r. the target computed over episode.
Multiple-goal Task: This setting involves: (1) learning an initial GMM policy where the robot end-effector reaches two different targets (i.e., task goals) starting from the same initial position, and (2) adapting the initial policy to reach 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. right and Fig. 12c 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 sparse penalty term penalize rollouts generating significantly divergent trajectories. Again, the policy optimization converges when the average position error w.r.t the chosen larget reaches minimum threshold.
4.2 RESULTS ANALYSIS
The reaching task tested our method's ability to adapt a previously- learned reaching skill new goal, located at (..)) (ef. Fig. 12f Achieving this required to adapt the Gaussian pa- rameters of mainly the last four GMM components, while the other ones remained unchanged. We compared all methods in terms of the success rate over environment steps, where the success rate is defined as the percentage of rollouts that reach the new goal. Figure Heft shows that our method achieved success rate of after approximately 70000 environment interactions. Despite PPO was also able to complete the task reliably, required many more environment steps (cf. Fig, (efft). 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 reach good success rates. Furthermore, this experiment showed that our method is 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 collision-avoiiance task, we tested whether our method was able to adapt a trajectory track- ing 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. midddle.. 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. Fig- ure middde shows that our method solved this task reliably after comparatively few environment interactions. Although PPO also achieved 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. middle.. These results again show the importance of the constraints The huge discrepancy in the required environment steps between our method and PPO further em- phasizes the importance of taking the GMM structure into account in the policy optimization.
the policy updates.
8

---

Under review as a conference paper at ICLR 2023
1.0
PPO
1.0
PPO
1,0
0.0
S4C
0.8
WGGF
SAC
SUCDESS 380
0.8
0.6
SAC
0.4
WKR
0.2
0.2
0.0
10K
20R
30x
40c
50k
60k
70k
BOk
20k
40%
60kc
80k
100k
20k
40k
60l
BOK
100K
120K
14CK 1502
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.
PPO
PPO
PPU
meted
WCF
WSF
WGF
-055
0.5
1.5
-10
-055
0.5
1.5
2.0
-05
0.0
1.0
7.0
010TE ma
Figure 4: Variance of the success rate over the 5 runs for method (WGF) and the two baselines 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 swarm plot, where dots indicate the success rates of individual runs. The plots show the variance at the following time steps from left to right: 80000, 90000, 95000.
While the previous two tasks were accomplished by adapting mostly the Gaussian parameters of the GMM, the multiple-goal task requires to adapt the GMM weights. The initial skill comprised reaching motions two different goals and an execution of this skill results in reaching one of them, depending on the sampling noise (cf. Fig. [2c). The easiest way to adapt the policy to reach only one of the two goals is 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. 3+ right, our method again quickly achieved success rate environment steps, while SAC was not able to solve the task.
PPO required substantially many more
In Fig. 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. Both baselines varied largely, except for the reaching task, where all SAC runs collapse that our method, despite showing large variance at the start, was able to quickly reduce the variance and converge reliably good success rate. We also provide similar plots of solution variance in Fig. 6 where we report the results for each method using its own convergence time step.
success rate of O. These results show
CONCLUSIONS AND FUTURE WORK
We presented a novel method for GMM policy optimization, which leverages optimal transport the- ory to formulate the policy optimization Our formulation explicitly accounts for the GMM structure in the optimization and furthermore en- ables to naturally constrain the policy updates by the L L 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- Wasssetenn manifold greatly reduced the computational cost of the policy optimization. Experiments on three robotic tasks provided strong evidence of the 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 re- cent advances on initializing the Sinkhorn algorithm (Thornton & Cuturi 2022). Also, we observed an intricate interplay between the optimization of the GMM weights and the Gaussian parameters, which sometimes 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 approach timization on the Bures- Wasserstein manifold, where both the locations and weights of the particles updated using Wasserstein Fisher-Rat gradient flows Chizat et al. (2015); Chizat (2019): Liero jet al. (2018). Finally it would be interesting to combine our method with an actor-critic formulation and to replace the multi-step cumulative reward by trained Q-function.
Wasserstein gradient flow on the manifold of GMMs.
fully dynamical, particle-based op-

---

Under review as a conference paper at ICLR 2023
REFERENCES
Fares J. Abu-Dakka, Leonel Rozo, and Darwin G. Caldwell. Force-based variable impedance learn- ing for robotic manipulation. Robotics and Autonomous Systems (RAS), 109:156-167, 2018. URL
robot 2018. 07.008
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), pp. 2623-2631, 2019. URL
http::::/ddo..oggee
10e.1115//33 11111//2//2200.... 3330701
Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media, 2005. URL https://
doi oog.org/// 10. 1007/b137080
Oleg Arenz, Mingjun Zhong, and Gerhard Neumann. Trust-region variational inference with Gaus- sian mixture models. Journal of Muchine Learning Reseurch (JMLR), 21(163): 2020. URL
http::/mmml wwwpoogpgggg/////...
524.htm
Shikhar Bahl, Mustafa Mukadam, Abhinav Gupta, and Deepak Pathak. Neural dynamic policies for end-to-end sensorimotor learning. In Neural Information Processing Systems (NeurIPS), pp. 5058-5069, 2020. URL /ttptttsss//prooceddnggg...
nerips wtaesscccpppeee pc/ppppe/2220//ilee/
Nicolas Boumal. An introduction to optimization smooth manifolds. To appear with Cambridge University Press, 2022. UtL/
net /bok/
Sylvain Calinon. A tutorial on task- parammaretttr movement learning and retrieval. Intelligent
Service Robotics, 9(1):1-29, January 2016. ISSN 1861-2776. htt.httss::///// //tttp::://d.... wor.o...../......
0187-9
Sylvain Calinon, Florent Guenter, and Aude Billard. On learning, representing, and generaliz- ing task in humanoid robot. IEEE Transactions Systems, Man, and Cybernetics, Part (Cyberneties), 37(2):286-298, 2007. URL pettp::/ddoi.oggee 1019///011...11 2006.
886952
Thomas Cederborg, Ming Li, Adrien Baranes, and Pierre- Yves Oudeyer. Incremental local online Gaussian mixture regression for imitation learning of multiple tasks. In JEEE/RSJ International Conference on Intelligent Robots and Systems (IROS). pp. 267-274. 2010. URL https://
Konstantinos Chatzilygeroudis, Vassilis Vassiliades, Freek Stulp, Sylvain Calinon, and Jean- Baptiste Mouret. survey on policy search algorithms for learning robot controllers in a handful of trials. IEEE Transactions on Robotics, 36(2):328-347, 2020. URI. htt.tttp:: //do org/
10. www.11109...O.. 2019 29988111
Yifan Chen and Wuchen Li. Optimal transport natural gradient for statistical manifolds with contin- uous sample space. Information Geometry, 3, 06 2020. URL http:::t////ot 00tp://///001..rr1111000077
001884-020-00022
Yongxin Chen, Tryphon T. Georgiou, and Allen Tannenbaum. Optimal transport for Gaussian mixture models. IEEE Access, 7:6269-6278, 2019, URL https://doi.org/10..11
ACCESS 2018 2889838
Lenaic Chizat. Sparse optimization on measures with over-parameterized gradient descent, 2019. URL
10300
Lenaic Chizat, Bernhard Schmitzer, Gabriel Peyré, and François-Xavier Xavier Vialard. An interpolating distance between optimal transport and fisher-rao, 2015. URL https::/aaxxii org/ abs
1506. 00430]
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), pp. 11357- 11363, 2021. URL httl/tttss:// //10: 10000/1009110 200.009/TTRA4880000 2021 99611666)
10

---

Under review as a conference paper at ICLR 2023
Marco Cuturi and Arnad Doucet. Fast computation of Wasserstein barycenters. In Inter-
national Conference on Machine Learning (ICML), pp. 685-693, 2014. URL https://
proceedings ptw.eediiing..m....ess..ese www................m...........e v32 wuturrr4..... html
Mare Peter Deisenroth, Gerhard Neumann, and Jan Peters. A survey on policy search for robotics. Found. Trends Robot, 2(1-2):1-142, 2013. URL http::///cc..or//1.....
23C0000021
Julie Delon and Agnès Desolneux. A Wasserstein-type distance in the space of Gaussian mixture models. SIAM Journal on Imaging Sciences, 13(2):936-977)) 2020.
Marco Ewerton, Gerhard Neumann, Rudolf Lioutikov, Heni Ben Amor, Jan Peters, and Guilherme Maeda. Learning multiple collaborative tasks with mixture of interaction primitives. In IEEE International Conference on Robotics and Automation (ICRA), pp. 1535-1542, 2015. URL
2015. 7139393
Benjamin Eysenbach and Sergey Levine. Maximum entropy RL (provably) solves some robust RL problems. International Conference on Learning Representations (ICLR), 2022. URL
https://openrevvew net forumridd f
Nadia Figueroa and Aude Billard. A physically-coossitenn bayesian non-parametric mixture model for dynamical system learning. In Conference on Robot Learning (CoRL), pp. 927-946, 2018.
URL https::/proceeei.pppp
www..press//ee
html
Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar 7. Alaya, Aurélie Boishunon, Stanis- las Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy, levgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. Pot: Python optimal transport. Journal of Machine Learning Research (JMLR), 22(78):1-8, 2021. URL http::/ www.lr.orgppppppppee v22 html
Zoubin Ghahramani and Michael Jordan. Supervised learning from incomplete data via an EM approach. In Advances in Neural Information Processing Systems (NeurIPS), 1994.
URL Intps::/proceesinngg newr.is.......ppp. wttp:/rriip..cpppeeee/ee//////// cc/paper/193//ille E2201 222011199.ccee222catt3333ebbbbb2
Paper.
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 Confer- ence on Machine Learning (ICML). pp. 1861-1870, 2018. URL https::/prrcced/
mlr. wil.spress/ve v88 wareess/80/haarnoo ja hb.html
Andi Han, Bamdev Mishra, Pratik Jawanpuria, and Junbin Gao. On Riemannian optimization over positive definite matrices with the Bures-Wi Senss-Wsserssttr geometry. In Advances in Neural Informa- tion Processing Systems (NeurIPS), 2021. URL
net/ foret//oommmmmm
ECHxGFmc62al
Yanlong Huang, Leonel Rozo, João Silvério, and Darwin G. Caldwell. Kernelized movement primitives. The International Journal of Robotics Research (IJRR), 38(7):833 852, 2019. doi: 10.1177702783669999886683
URL http:/ttss///oi.ogg////////
1177 0tt..//.........1.00000789993333
Julian Ibarz, Jie Tan, Chelsea Finn, Mrinal Kalakrishnan, Peter Pastor, and Sergey Levine. How to train your robot with deep reinforcement learning: lessons we have learned. The International Journal of Robotics Research (IJRR), 40(4-5):698-7211 2021. [Rt.ttts:: //doi orrg
Auke Jan Ijspeert, Jun Nakanishi, Heiko Hoffmann, Peter Pastor, and Stefan Schaal. Dynamical movement primitives: Learning attractor models for motor behaviors. Neural Computation, 25: 328-373, 2013. URL httpttttt///o....mo org 10
00393
Noémie Jaquier, David Ginsbourger, and Sylvain Calinon. Learning from demonstration with model-based gaussian process. In Conference on Robot Learning (CoRL), pp. 247-257, 2019. URL|ttt:: http:httpp//prccceeeiiee w 00/ http:///1000///aeae//ee//00 html

---

Under review as a conference paper at ICLR 2023
Richard Jordan, David Kinderlehrer, and Felix Otto. The variational formulation of the fokker- planck equation. Siam Journal on Applied Mathematics, 1996.
Mohammad Khansari-Zadeh and Aude Billard. Learning stable nonlinear dynamical systems with Gaussian mixture models. IEEE Transactions on Robotics (T-RO), 27(5):943-99) 957, 2011. URL htt.thpss://oo.or//////p//
TwT.0//TRR.../ 2t1./TOO..01011.10 2011.2159412
Jens Kober, J. Andrew Bagnell, and Jan Peters. Reinforcement learning in robotics: A survey.
The International Journal of Robotics Research (IJRR), 32(11):1238 1274, 2013. URL https:
Matthias Liero, Alexander Mielke, and Giuseppe Savaré. Optimal entropy-transport problems and a new hellinger-kantoroovie distance between positive measures. Inventiones mathematicae, 211, 03 2018. doi: 00.10077/00222--17707775999
Yun Lin, Shaogang Ren, Matthew Clevenger, and Yu Sun. Learning grasping force from demon- stration. In IEEE International Conference Robotics and Aulomation (ICRA), pp. 1526-1531, 2012. URL pectps:/ddoi.oorr
1109/ICRA, 2012. 6225222
Luigi Malagò, Luigi Montrucchio, and Giovanni Pistone. Wasserstein Riemannian geometry of Gaussian densities. Information Geometry, 1:137-179, 2018.
00677//18884018-00000 0014-4
Petr Mokrov, Alexander Korotin, Lingxiao Li, Aude Genevay, Justin Solomon, and Evgeny Burnaev. Large-scale Wasserstein gradient flows. In Neural Information Processing Systems (NeurIPS), 2021.
Ted Moskovitz, Michael Arbel, Ferenc Huszar, and Arthur Gretton. Efficient Wasserstein natural gradients for reinforcement learning. In International Conference on Learning Representations (ICLR), 2021. /ttptttss::///eeetttta
net
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 on Robotics and Automation (ICRA), pp. 8651-8657, 2022. URL.
https:://oll.org10...110 http:///.1100///14666333.3331
ICRA46639 2022. 9811770
Fabian Otto, Philipp Becker, Ngo Anh Vien, Hanna Carolin Ziesche, and Gerhard Neumann. Dif- ferentiable trust region layers for deep reinforcement learning. In International Conference on Learning Representations (ICLR), 2021. URL. https: http::::////eeeeeeeeeeee net. / forum?id
GYZD-AO1VM
Felix Otto. The geometry of dissipative evolution equations:: The porus medium equation. Com- munications Partial Differential Equations, 26(1-2):101-174, 2001. URL httpts://doi......
org/10.105] 101111011PPEEEE1 1EE111PDDE--E- 103002243.
Aldo Pacchiano, Jack Parker-Holder, Yunhao Tang, Krzysztof Choromanski, Anna Choroman- ska, and Michael Jordan. Learning to score behaviors for guided policy optimization. In ternational Conference on Machine Learning (ICML), pp. 7445-7454, 2020. URL https:
//proceedings mlr. stteas//l19/ppccchaano///. html
html
Victor M. Panaretos and Yoav Zemel. The Wasserstein Space, pp. 37-57. Springer International Publishing, 2020. /tL/httts///do.crrot http:////oo..001111000100//1
Alexandros Paraschos, Christian Daniel, Jan Peters, and Gerhard Neumann. Using probabilistic movement primitives in robotics. Autonomous Robots, 42:529-551, 2018. URL https::/
doi. foi.org//00
Jan Peters. Daniel D. Lee, Jens Kober, Duy Nguyen- J. Andrew Bagnell, and Stefan Schaal.
Robot Learning, pp. 357-398. Springer International Publishing, 2016. ISBN 978-3-319-32552.
1. URL /oi wow.loorrol... 1007 3- - 32552-1
Gabriel Peyré and Marco Cuturi. Computational optimal transport: With applications to data sci- ence. Foundations and Trends in Machine Learning, 11(5-6)3355667, 2019. URL http:
7/dx.dol .ww..org//11 1561 015612200000033
12

---

Under review as a conference paper at ICLR 2023
Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dormann. babble-asselness:: Reliable reinforcement learning implementations. Journal of Machine Learning Research, 22(268):1-88 2021. URL http://jmiz http://mm....rr/a////ee/oe/oe//a pttpg//ppps//2////// 20 1364. html
Harish Ravichandar, Iman Salehi, and Ashwin Dani. Learning partially contracting dynamical sys- tems from demonstrations. In Conference on Robot Learning (CoRL), pp. 369-378, 2017. URL
https://proceette .orssmmlrr www..m.........emmmmmmmmmmm prwss/v18////eppp ravichandary 17a.117a.. html
Pierre H. Richemond and Brendan Maginnis. the Fokker-Planck equation, 2017. URL attpss//arrxx..orrr/////122..... arXiv:1712.07185.
On Wasserstein reinforcement learning and
Leonel Rozo, Sylvain Calinon, Darwin G. Caldwell, Pablo Jiménez, and Carme Torras. Lear- ing physical collaborative robot behaviors from human demonstrations. IEEE Transactions on Robotics, 32(3):513-527, 2016. URL. htttps::/dbi.orgg// 1101 2016.
2540623
Filippo Santambrogio. Optimal AptiiiaiiTTeeeapoooo/ooooaaAaaaa 2015. http:////o...ooooo111100//////10
Birkhäuser, Cham, edition,
Filippo Santambrogio. Euclidean, metric, and Wasserstein gradient flows: an overview. Bull. Math. Sci., 7:87-154, 2017.
Stefan Schaal, Auke ljspeert, and Aude Billard. Computational approaches motor learning by imitation. Phil. Trans. R. Soc. Lond. B, 358:537-547,, 2003. URL http:://do/ /C0 1098/ratb. 2002 1258.
John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In International Conference on Machine Learning (ICML), pp. 1889-1897, 2015.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms, 2017. http::/s.........p..s/s/////aaaa
html
Freek Stulp and Olivier Sigaud. Policy improvement methods: Between black-box opti- mization and episodic reinforcement learning. hal-007334663, 2012. URL http://hal,
James Thorton and Marco Cuturi. Rethinking initialization of the Sinkhorn algorithm, 2022. URL https://arxiv wttp:///////..ooooo///// / 07630 arXiv:2206.076300
James Townsend, Niklas Koep, and Sebastian Weichwald. Pymanopt:
Python toolbox for opti- mization on manifolds using automatic diffcrentiation. Journal of Machine Learning Research
(JMLR), 17(137):1-5, 2016. /ttpa////////////oo//o//////////
Shijun Wang. Baocheng Zhu, Chen Li, Mingzhe Wu, James Zhang, Wei Chu, and Yuan Qi. Rieman- nian proximal policy optimization, 2020. URL
arXiv:2005.09195.
Ronald J. Williams. Simple statistical gradient-following algorithms for connectionist reinforce- ment learning. Machine Learning, 8:229-256 2004. URL. 1tttpps://oci.or/11..11111 1007/
BFC0992696
Ruiyi Zhang, Changyou Chen, Chunyuan Li, and Lawrence Carin. Policy optimization as Wasser- stein gradient flows. In International Conference on Machine Learning (ICML), pp. 5737-5746,
2018. httpt http://::///poreeeenngee mlr press v80 pptpsss/v80/haaggg..... 2att:///aaaaaaaaaaaaaaaaa html
13

---

Under review as a conference paper at ICLR 2023
A APPENDIX
A.1 DETAILS ON GAUSSIAN MIXTURE REGRESSION (GMR)
In GMR we start from a GMM in state-action space iTs a) Ex= (Isaaa pr Σ.) from which a policy, ie. probability distribution on the action space, can be obtained by conditioning on the state, as follows
π (8, [ a)da
(a s)
(15)
The resulting, conditional distribution another GMM on the action sapce, with state dependent parameters, given by:
N
T T 81) = Σ
Ei), 14 (8)) = 1h Dg (2)) (8t -4)),
wi
with
(16)
iml
(17)
Σ Ex E.e (Σ,) E500
(18)
(N p Σ;}
(s)) Σ (sri pi: Ei}
(19)
Note that we have split the GMM parameters μM and E, into their state and action components according to
Σ
p.
E.
(20)
A.2 RIEMANNIAN GRADIENTS AND RETRACTIONS
For completeness we give here the explicit expressions of the Riemannian gradients and the retrac- tions used in 33.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 the well-known Euclidean gradient descent
p = fax + (TJ((T),
(21)
where Vn denotes the Euclidean gradient w.r.t. ft. For the covariance matrices we use the gradient and retraction w.r.t. the Mures-Wasserstein manifold, taken from (Malagò al. 2018 Han et al., 2021
The gradient is given by
grad J(xx) 4 Ets
(22)
(x+x") 2
where again Vs denotes the Euclidean gradient w.r.t. S and {X)s ss = retraction given by
Furthermore, the
RE
x Σx X + Cx ΣS
XL
(23)
where Cx
'S$
is the Lyapunov operator, defined as the solution to the matrix linear system Ez.
L
XLx
A.3 EXPRESSIONS OF THE FREE FUNCTIONAL ((T) AND ITS EUCLIDEAN GRADIENTS
For completeness sake, provide here the explicit expression of the Euclidean gradients for the objective (TT) w.r.t. the parameters of the GMM, which are used the construction of the Rie- mannian gradients. Using the policy gradient theorem, obtain the gradient of Eq. 11 w.r.t
14

---

Under review as a conference paper at ICLR 2023
parameter & as follows
VT)(T)) Ve / [Idesdeedanppe))))(( (1) ai)
Exr(or,arr) (24)
872
r(s,, Er(st, as)
E,,
Ve (og(T(ar(s1))
E,
Σ Ve log
π ( as)
Σ r(s and Mst
Ertonaad)
ddant(s)) ar)
Σ E, (
(sT ((8, ar)
(e ai) [s,,
r( an)
In this work, we focus on GMM policies, for which the objective ((T) takes the form:
((x) / p()))) ((s))))( E(ss)))ssss oz) Σ 7'r(s;, ar)
f da
+ 3
wi Mule:
a.)
log
(25)
By inserting Eq. 25
into Eq. 24
obtain for the individual parameters of the GMM
WiN o1 M1, Σ E ((8), at) mm)
V [ = E Σ
(26)
E,w,N(st, aii Hj: Ey)
wi I daN(ss, ar aii E, 1 (($), at) - mi)
Σ est
r(s ai)
Σnw, w. J daN (8t; a+; p.j; Ej)
= E, WIN(s ae MI, Σn) 5 ((1), at) 14))
(27)
Eyw,N(ss ari pej Ej)
WiN WiN(s M1,s E1.oa) Diid (st - (4is)
8
Σ r(sn, an)
E, Ww,N Mjjss Egnos)
ESt
E, Σ (G ((s), az) - mi) 8.51.0. Ei,do (st m1,e)
Σ rr( an)
Mat
Here 8, E {0,1) indicates which terms the gradient acts on. this case, the gradient act on the state components and it is absent for the action dimensions.
WIN(s ar ui Σ,)Σ,' 1 (1 (()),a)) - mi) ((s), m1) ΣT
J = Er Σ
2
E, W,,NN(s, ari Mij Eny)
(28)
1 Wi I daN aii Σn) ET (1))
((s az) - mi) ((8), ar) - m1) Σj
Σ r rsse as)
+
2
E, w, daN as Hij Ey)
t>t
Er Σ
Gl D) (1- ((s), az) mi) ((sr, az) u)) _T-1)))
2
+0,
GL, -1 (1 (1-
(1-(s: (8: - m1,, (st - mlw) Σ1Ah
Σ e>1
r(st, an)
2
115
15

---

Under review as a conference paper at ICLR 2023
N(s du: ,4, Σ.)
S daN azi ui, Ei)
[wi E Σ
Σ
r(se: an)
Σj (80, py 2zj)
Σj w3 ds, ar poy Eij)
(29)
E Er(s).a.)
E
(30)
W1
Note that we introduced the responsibilities Sieeca and Sr.a,, which are defined as follows
Ciaaa = WIN(Bi. a:: p4; Ei)
and
(31)
Σj (s at Míj E,)²
wi I daN(ss au Mi Σ,) Cr.s Σn wi daV Mz Ej)
W/N(Sti Wilss Mi,ss Ei,as)
(32)
Σ, Ws phj.s; Ej.ss)
A.4 RELATION BETWEEN FORWARD AND BACKWARD DISCRETIZATION IN THE BURES- WEEEES-WSSSRSTEEN METRIC
In this section we outline the relation between the implicit and explicit optimization schema w.r.t the Bures- Wasserstein metric, which is leveraged to formulate our policy optimization in 3 We closely follow Chen & Li (2020). For the sake of simplicity, we group the Gaussian parameters μ and Σ into single parameter vector A. Furthermore, we restrict our explanation to a single Gaussian component, which possible without loosing generality, as each of the N components live in its own manifold Rd S The Riemannian gradient w.r.t the Gaussian parameters 0, gradg
satisfies by definition
go(grad, ((T(O), (T)( &,
(33)
where Vo denotes the Euclidean gradient, & is arbitrary vector the tangent space ToM, and Ge is the Riemannian metric tensor, defining the inner product on M. The Riemannian metric can be written
go(C. C
(34)
with two arbitrary tangent vectors Ç. E, and Gw (0) being positive definite matrix. Moreover, note that the Wasserstein distance W (N + A0)). where AO denotes small perturbation in the Gaussian parameters 0, can be expressed as
1 | Gw
W W2 (1 + A0)
Gw (0) (((0))))
(35)
2
for A0 0. Similarly, can approximate the objective evaluated at J[@ A0) via the Taylor theorem as
J( + +00 = (()) - Vol(0) -0 O((0))))))
(36)
With this, we can approximate
(O), - -((()))
Oz+1 = arg min
(37)
2r
0k)
Gw 0k) VoJ (0 Bk))
~ arg min 0
(38)
2r
from which we obtain the update equation for 0 as follows
0 x++ 0g + GG (0x) (V)
(39)
Note that Eq.
39 turn corresponds to an approximation of the exact Riemannian gradient descent
k Ro. (A grad, (+)))))))
(40)
16

---

Under review as a conference paper at ICLR 2023
This approximation can be obtained by considering first-order approximation of the geodesic the BW manifold. As the exponential map (a.k.a. the retraction) is defined via the geodesic, the retraction operator in Eq. 40 into simple addition operation under first-order upproximation, leading to Eq. 39 Notice that such approximation does not guarantee that the updated parameters stay on the manifold, except for the cases in which 0 E Rd In case, we leverage the retraction and Riemannian gradients of (Malagó et al. 2018 Han et al. 2021), which allow us to apply the exact Riemannian gradient descent of [40] This avoids to rely on first-order approximations and in
turn we can guarantee that the updates of the Gaussian distribution parameters always lie on on the
product manifold (I2 Sd
A.5 ADDITIONAL DETAILS ON THE IMPLEMENTATION
We extended the Pymanopt Towwsend et al. 2016) by adding a custom line-search routine that accounts for constraint 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 L Wasserstein distance between two points on the manifold of GMMs as constraint.
Input: point Zo the manifold, descent direction d, initial step size λo decrement a, con- straint c( -). maximum allowed value for constraint Cmax minimum step size Amin Output: step size updated point manifold
1:
To + Ao d
x= Ào
2: while c(æo. x) Cmax and À > Amin do
3: decrease step size: X = a
update point on manifold: zo
4: end while
5: if Amin then
6: return λo,
7: else
8:
return A,
9: end if
A.6 ADDITIONAL DETAILS ON EXPERIMENTS
A.6.1 ADDITIONAL RESULTS
Fig. 5 shows the convergence curves for the basclines as Fig. 3 the main paper, however,
we extended the horizontal axis up to the maximum number of environment steps used for training.
1.0
SU0005 ga00
0.6
0.6
0.6
0.4
0.4
0.4
150K
250k
175K
InVVrommment steps
envirtnmment steps
envirromment StRCS
Figure 5: The success rate of the two baselines on the reaching task (Ieff), the collision-avoidanee avoidance task (middle) and the multiple-goal task (right). The shaded area indicates the standard deviation over 5 runs.
Fig. 6 shows the variance of the success rate for the three methods at their time step of convergence for all three robotic tasks. Concerning SAC, which did not converge after the maximum number of 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 the five the time step of conver-
17

---

Under review as a conference paper at ICLR 2023
gence, at the cost of a prohibitively large number of steps. SAC showed huge variance in all tasks,
apart from the reaching task, where all runs collapsed to a success rate of
PPO
netel
SAC
SAC
SAC
-100
-155
0.0
0.5 1.0
1.5 -100
-05
01.0
61.5
1.0
1.5
2.0
-10
-05
0.01
3.0 0.5
1.0
1.5
ALCTHO 13te
Figure 6: Variance of the success rate over the 5 runs for our method (WGF) and the two baselines 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 rates of individual runs. The time steps which we determined the variance are for PPO. SAC WGF for the three tasks from left to right: (280000, 400000, 80000), (275000, 300000, 90000), (130000, 200000, 95000).
A.6.2 ADDITIONAL ABLATIONS
In order to assess the influence of leveraging Riemannian optimization approach on the Bures- Wasserstein manifold, we conducted an ablation of our method by eliminating the Riemannian for- mulation. Instead the explicit Euler scheme update in Eq. [2. which corresponds to Riemannian gradient descent w.r.t the Saree-Waasersttinn metric, we use the implicit Euler scheme
W (A), mT)
f arg min
((Tx(()))
(41)
2r
(T L), mk)
Σx+1 arg min
(T (ε))
(42)
To guarantee that the updated covariance matrices do not leave the manifold of symmetric positive definite matrices, we parameterize them terms of Cholesky factors. The results obtained with this non non-Riemann Riemannian version of our method are shown in Fig. in direct comparison to our method and Fig. for an extended range.
1.0
WGT
1.0
1.0
2-00055 0200
0.8
0.4
0.4
0.4
0.2
0.2
VISI
0.2
noY_EW
non_BM
10k
20k
30k
40%
50K
63k 70K
80k
20c
40
50N
50%
100c
120% 140K
20k
40k
65K
80k
100K
120k
enviro-ment staps
enTAFTammennt stapps
emwironmmem
staps
Figure 7: The success rate of our method and an ablated version, not using the Bures- Wasserstein formulation for the reaching task (left), the collision collision-avoidanee avoidance task (middle) and the multiple-ge task (right). The shaded area indicates the standard deviation over 5 runs.
The results clearly show that the non-Riennee non-Riemannian method struggles to reach success rate of for the reaching task and the collision-avoidance task. Furthermore, we observe a high variance different runs in the 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 in turn leads to the aforementioned high variance. Nevertheless, the multiple-godl task constitutes an exception. Here we observed a similar perfor- mance for our approach and the ablated method. The reason for this is that the optimization of this task mainly dominated by the weight updates, which are identical for both methods. This result therefore expected and confirms that correctness of our ablation strategy.
settings (see Fig. 9 and Fig. T0p. We attribute this the fact that
A.6.3 ADDITIONAL EXPERIMENT WITH 7-DoF ROBOTIC MANIPULATOR
We carried out an additional experiment to show that our method can be employed formed by off-the-shelf robotic manipulators (e.g. 7-DoF Franka Emika Panda robot). Specifi-
on tasks per-
18

---

Under review as a conference paper at ICLR 2023
1.0
non BW
1.0
von_BW
0.8
success rate
0.6
Success
0.4
0.2
0.2
0.0
0.0
50k
100k 150k 200k 250k 300k 350k 400k environment steps
0
25k 50k 75k 100k 125k 150k 175k 200k environment steps
Figure 8: Extended plot of the success rate of and ablated version of our method, not using the Bures- Wasserstein-baset formulation for the reaching task (left), the collision-avoidance task (middle) and the multiple-goal task (right). The shaded area indicates the standard deviation over runs.
WGF
WGF
WGF
method
method
non_BW BW
BW
non BW
-10 -0.5 0.0 0.5 1.0 1.5 2.0 succesS rate
1100 -0.5 0.0 0.5 1.0
1.5
2.0
-1.0 -0.
0.0
0.5
1.0 1.5
2.0
SUCCeSS rate
SUccess rate
Figure 9: Variance of the on the reaching task (left), collision avoidance task (middle) and multiple-goal task (right). The violine plots are overlaid with hox 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, 90000, 85000.
runs for method (WGF) and the ablated method (non-BW)
cally, we extended the collision-avoidance task described in $ $ to a 3D environment (i.e. the state = x E R3 and the action a E R*). 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 a function of the position error between the robot end-effector position and the target the end of the rollout. Moreover, two sparse penalty terms were added to punish collision with obstacles and divergent trajectories.
Similarly the planar task reported in the main paper, tested whether method was able to adapt a trajectory tracking skill 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. The robot end-effector pose was controlled using a full-pose Cartesian velocity controller at a fre- quency of 100Hz, where the end-effector orientation method reached success rate of 1.0 very quickly, taking approximately 20000 environment steps. Moreover, the solution variance of our method was also very low, which consistent with our ob- servations concerning the performance of our policy optimization on the three planar tasks analyzed in the main paper.
kept constant. Figure [T shows that
A.6.4 INITIAL GMM POLICIES
For the sake of completeness, Fig. 12 provides 2D projections of the initial GMM policies learned from demonstrations for the three robotic settings considered in the main paper: the reaching mo- tion skill, the collision-fres trajectory tracking. and the multiple-goal task. Figure [2 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
19

---

Under review as a conference paper at ICLR 2023
mathod
method
nen, BW
non_BW
nen_BW
-1.0 -0.5 0.0 0.5 1.0 1.5 2.0 success rate
-1.0 -0.5 0.0 0.5 1.0 1.5
2.0
-1.0 -0.5 0.0 0.5 1.0 1.5 2.0
Success rate
SUccesS rate
Figure 10: Variance of the success over 5 for method (WGF) and the ablated method (non-BW) the reaching task (left), the collision avoidance task (middle) and the multiple-goal task (right). The violine plots are overlaid with box plots, quartile dicate the success rates of individual runs. The time steps at which we determined the variance are (80000, 400000), (90000, 200000), 900000,
swarm plot, where dots
Figure 11: The success rate of our method applied to the 3D narrow-path task performed by the 7-DoF Panda robotic manipulator. The shaded indi- cates the standard deviation over 5 runs.
Success
0.4
0.2
0.0
WGF
20k
40k
60k
80k
100k
120k
140k
anvironment-steg
20

---

Under review as a conference paper at ICLR 2023
10.0
10
5.0
0.0
-30
-5.0
-7.3
-10.0
11.0 7.5
-5.0
-2.5
0.0
2.5
5.3
7,5
10.0
-00
10
10
20
30
10
(a) Reaching motion skill
10.0
7.3
25
5.0
2.5
15
2.5
-10.0
2.5
-10
(b) Collision- free trajectory tracking
2
-6
-8
-10
-1
2
2
1
-2
Multiple-goal task
Figure 12: Green Gaussian components represent the initial GMM policy learned from demonstrations, projected on the Cartesian position (lefi) and velocity (lefi) spaces. The recorded position and velocity data are depicted as black dots (+).
21


Reviews:
[{"review_id": "dTxe18cEqpn", "reviewer": "Reviewer_8XBw", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors consider the problem of policy optimization; they propose using Bures-Wasserstein gradient flows for optimizing policies which are constrained to the class of Gaussian mixture models. Existing methods in this setting don't exploit the structure enforced by the Gaussian mixture models, and the authors demonstrate the advantage of this perspective through empirical analyses on 2$d$ benchmarks. ", "summary_of_the_review": "The authors propose an interesting method for policy optimization through the lens of Bures-Wasserstein gradient flows. While the theoretical underpinnings are unclear, the experiments show useful improvement over _some_ benchmarks. It would have been even more compelling if the experiments were more exhaustive. ", "strength_and_weaknesses": "**Strengths:**\n* The paper is well-written and easy to follow. I appreciate the informative, yet concise, introduction on Wassserstein gradient flows for the non-expert reader. \n\n* Even though the idea of formulating policy optimization using Wasserstein gradient flows is already known, the use of the Bures-Wasserstein geometry in the specific setting of Gaussian mixture models is novel and interesting. \n\n* Despite the drawbacks in the experiments (see comments below) the authors do a nice job interpreting the results and providing useful insights. \n\n* The authors claim that their perspective leads to significant computational advantage. (This would have been more compelling if it were reported as a benchmark of time taken for each method). \n\n\n**Weaknesses:**\n* My biggest concern is on the theoretical front. The objective of the paper is to study Wasserstein gradient flows which don't provide any dynamics for the weight update. A weight update, as in Eqs. (13 & 14), will warrant machinery from Wasserstein Fisher-Rao gradient flows instead. See, for example,  [Chizat et al. (2018)](https://link.springer.com/article/10.1007/s10208-016-9331-y) [Liero et al. (2018)](https://link.springer.com/article/10.1007/s00222-017-0759-8) [Chizat (2022)](https://link.springer.com/article/10.1007/s10107-021-01636-z) Therefore, it is unsurprising that the authors note, in their conclusions section, that the \"weight update often hampered with the update of the Gaussian parameters\". \n\n* Notwithstanding, for the proposed method using the Bures-Wasserstein geometry, I would really have liked to see some guarantee that the proposed methodology leads to policy improvement, c.f., Wang et al. (2020, Theorem 2). \n\n* On the experimental front, I found the benchmarks to be lacking. For example, it would have been instructive if the authors compared their method to other parametric policy updates such as [Moskovitz (2020)](https://arxiv.org/abs/2010.05380) ; especially since the notion of Wasserstein natural gradients (Chen & Li, 2020) is mentioned in §3.1 and §A.4. \n\n* Furthermore, it would have been even more enlightening if the authors contrasted their performance to that of Wang et al. (2020), since this would have demonstrated the advantage of the Bures-Wasserstein perspective in contrast to their perspective which relies on the JKO scheme.\n\n\n\n**Minor Clarifications:**\n* I'm unclear on how the number of Gaussian components, $K$, are chosen for the mixture model. For example, Delon & Desolneux (2020, §6.3) note that performance is sensitive to the choice of $K$. Alternatively, Figueroa & Billard (2018) tackle this issue by choosing a Dirichlet process prior on $K$. How sensitive are the results to this choice?\n\n* The authors implement the JKO scheme for the weight update using the Sinkhorn algorithm, which includes an entropic regularization. How does this interplay with the entropic regularization in the free energy functional J? The proximal map is effectively subtracting  a factor $\\beta$  of the entropy term due to $J_{r, \\beta}(\\pi)$ and adding a $1/2\\tau$ factor due to Sinkhorn?\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "0n6LT8cjLA9", "reviewer": "Reviewer_1kyr", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper presents GMM-based robot policy optimization formulated as a Wasserstein gradient flow, resulting in constraining policy updates for a stable optimization process. The proposed method is compared against two baselines PPO-based GMM update and SAC-GMM. The evaluations are performed in three tasks with a toy robot arm to reach a target, avoid collision during trajectory tracking, and reach multiple targets. The results demonstrate that their proposed approach (WGF) outperforms PPO and SAC in all three tasks. ", "summary_of_the_review": "This paper presents a stable approach to optimize GMM policies for robot control which also outperforms prior methods such as PPO and SAC. However, the experiments could have more complex environments to highlight the proposed approach's scalability better. Even if it does not scale, those evaluations will highlight a limitation to address in future work. \n\nUpdated: The paper presents an interesting idea but lacks experiments that could validate the proposed approach under standard collision avoidance constraints for motion planning problems.", "strength_and_weaknesses": "*Strength*\n+ A novel Gradient flow approach for GMM optimization\n+ Algothrim seems sound with theoretical backing.\n+ The results demonstrate better performance than PPO and SAC\n\n*Weaknesses*\n- Experiment section is relatively weak. The robot tasks are too simple. Although WGF outperforms PPO and SAC on those tasks, a more cluttered scenario with a realistic robot arm (URDFs of UR5, Panda, etc.) would exhibit the applicability of the proposed approach to practical environments. Perhaps a relevant baseline for simulation environment setup could be [1].\n\n[1] Continuous-time Gaussian process motion planning via probabilistic inference. The International Journal of Robotics Research, 37(11), 1319-1340.     ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "-1i_yoM06E", "reviewer": "Reviewer_zFF5", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a Riemannian optimization method to optimize GMM policies. It uses GMM to represent the policy structure and demonstrates an EM-like approach to optimize such a policy with hidden variables with a maximum entropy RL objective. In each iteration, the Gaussian parameters are optimized with the Riemannian gradients, while the GMM parameters are optimized later with an additional Wasserstein distance constraint. Optimizing GMM with RL objectives enables the authors to transfer a GMM policy from demonstrations to new tasks. The authors considered several reach or collision avoidance tasks and showed better performance.", "summary_of_the_review": "The paper proposes an interesting idea of how to optimize a policy with prior structures. However, I feel the paper lacks enough ablation study and can not be accepted given its current form. I am happy to increase the score if additional results are provided. ", "strength_and_weaknesses": "Strength:\n\n- This paper attacks a very important direction, optimizing a policy with hidden structures. Though the paper only studied Gaussian policy here, I can imagine that a similar idea can be applied to other policy classes, like a neural network.\n- The EM-like optimization approach for optimizing GMM with REINFORCE is novel to me. I think this might be the key to its success, as directly optimizing the likelihood of GMM models are not easy.\n- Please correct me if I am wrong. The proposed Riemannian ensures the optimization does not break the semidefinite constraints of Gaussian models. Optimizing GMM parameters in the BW manifolds constrain the policy update, enabling better transfer performance.\n\nWeakness:\n\n- The GMM policy can hardly generalize to high-dimensional tasks, which will constrain its application. It will be interesting if it is possible to optimize broader policies.\n- I feel the paper lacks the necessary ablation study. The baseline comparison is reasonable. However, it is unclear why the proposed method works by comparing it with SAC. I do not fully understand why the Reinmannian gradient is necessary. Is it possible to optimize the policy by parameterizing the Sigma as a form of V^TV to guarantee its positive semidefinite? What would happen? Besides, what would happen if the Wasserstein part is removed from eq (13)? An ablation study is required to illustrate the contribution of each component.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject