Original: W ASSERSTEIN G RADIENT F LOWS FOR O PTIMIZING GMM- BASED P OLICIES
Revised: W ASSERSTEIN G RADIENT F LOWS FOR O PTIMIZING GMM- BASED P OLICIES

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Robots often rely on a repertoire of previously-learned motion policies for performing tasks of diverse complexities.
Revised: Robots often rely on a repertoire of previously-learned motion policies for performing tasks of diverse complexities.

Original: When facing unseen task conditions or when new task requirements arise, robots must adapt their motion policies accordingly.
Revised: When facing unseen task conditions or when new task requirements arise, robots must adapt their motion policies accordingly.

Original: In this context, policy optimization is the de facto paradigm to adapt robot policies as a function of task-specific objectives.
Revised: In this context, policy optimization is the de facto paradigm to adapt robot policies as a function of task-specific objectives.

Original: Most commonly-used motion policies carry particular structures that are often overlooked in policy optimization algorithms.
Revised: Most commonly-used motion policies carry particular structures that are often overlooked in policy optimization algorithms.

Original: We instead propose to leverage the structure of probabilistic policies by casting the policy optimization as an optimal transport problem.
Revised: We instead propose to leverage the structure of probabilistic policies by casting the policy optimization as an optimal transport problem.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Furthermore, we leverage the geometry of the BuresWasserstein manifold to optimize the Gaussian distributions of the GMM policy via Riemannian optimization.
Revised: Furthermore, we leverage the geometry of the BuresWasserstein manifold to optimize the Gaussian distributions of the GMM policy via Riemannian optimization.

Original: We evaluate our approach on common robotic settings: Reaching motions, collision-avoidance behaviors and multi-goal tasks.
Revised: We evaluate our approach on common robotic settings: Reaching motions, collision-avoidance behaviors and multi-goal tasks.

Original: Our results show that our method outperforms common policy optimization baselines in terms of task success rate and low-variance solutions.
Revised: Our results show that our method outperforms common policy optimization baselines in terms of task success rate and low-variance solutions.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: One of the main premises about autonomous robots is their ability to successfully perform a large range of tasks in unstructured environments.
Revised: One of the main premises about autonomous robots is their ability to successfully perform a large range of tasks in unstructured environments.

Original: This demands robots to adapt their task models according to environment changes, and consequently to adjust their actions to successfully perform under unseen conditions (Peters et al., 2016).
Revised: This demands robots to adapt their task models according to environment changes, and consequently to adjust their actions to successfully perform under unseen conditions (Peters et al., 2016).

Original: In general, robotic tasks, e.g. picking or inserting an object, are usually executed by composing previously-learned skills (Schaal et al., 2003), each represented by a motion policy.
Revised: In general, robotic tasks, e.g. picking or inserting an object, are usually executed by composing previously-learned skills (Schaal et al., 2003), each represented by a motion policy.

Original: Therefore, in order to successfully perform under new settings, the robot should adapt its motion policies according to the new task requirements and conditions.
Revised: Therefore, in order to successfully perform under new settings, the robot should adapt its motion policies according to the new task requirements and conditions.

Original: Research on methods for robot motion policy adaptation is vast (Kober et al., 2013; Chatzilygeroudis et al., 2020), with approaches mainly building on black-box optimizers (Stulp & Sigaud, 2012), end-to-end deep reinforcement learning (Ibarz et al., 2021), and policy search (Deisenroth et al., 2013).
Revised: Research on methods for robot motion policy adaptation is vast (Kober et al., 2013; Chatzilygeroudis et al., 2020), with approaches mainly building on black-box optimizers (Stulp & Sigaud, 2012), end-to-end deep reinforcement learning (Ibarz et al., 2021), and policy search (Deisenroth et al., 2013).

Original: Regardless of the optimization method, most approaches rely on policy structure-unaware adaptation strategies.
Revised: Regardless of the optimization method, most approaches rely on policy structure-unaware adaptation strategies.

Original: However, several motion policy models (e.g, dynamic movement primitives (DMP)(Ijspeert et al., 2013), Gaussian mixture models (GMM) (Calinon et 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.
Revised: However, several motion policy models (e.g, dynamic movement primitives (DMP)(Ijspeert et al., 2013), Gaussian mixture models (GMM) (Calinon et 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.

Original: 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.
Revised: 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.

Original: Second, structure-unaware adaptation strategies optimize the policy parameters disregarding the special characteristics of the policy model (e.g., a DMP represents a second-order dynamical system).
Revised: Second, structure-unaware adaptation strategies optimize the policy parameters disregarding the special characteristics of the policy model (e.g., a DMP represents a second-order dynamical system).

Original: In this regard, we hypothesize that the policy structure may be leveraged to better control the adaptation strategy via policy structure-aware gradients and trust regions.
Revised: In this regard, we hypothesize that the policy structure may be leveraged to better control the adaptation strategy via policy structure-aware gradients and trust regions.

Original: Our main idea is to design a policy optimization strategy that explicitly builds on a particular policy structure.
Revised: Our main idea is to design a policy optimization strategy that explicitly builds on a particular policy structure.

Original: Specifically, we focus on GMM policies, which have been widely used to learn motion skills from human demonstrations (Calinon et al., 2007; Cederborg et al., 2010; Calinon, 2016;
Revised: Specifically, we focus on GMM policies, which have been widely used to learn motion skills from human demonstrations (Calinon et al., 2007; Cederborg et al., 2010; Calinon, 2016;

Original: Jaquier et al., 2019).
Revised: Jaquier et al., 2019).

Original: GMMs provide a simple but expressive enough representation for learning a large variety of robot skills: Stable dynamic motions (Khansari-Zadeh & Billard, 2011; Ravichandar et al., 2017; Figueroa & Billard, 2018), collaborative behaviors Ewerton et al.
Revised: GMMs provide a simple but expressive enough representation for learning a large variety of robot skills: Stable dynamic motions (Khansari-Zadeh & Billard, 2011; Ravichandar et al., 2017; Figueroa & Billard, 2018), collaborative behaviors Ewerton et al.

Original: ; Rozo et al.
Revised: ; Rozo et al.

Original: (2016), and contact-rich manipulation Lin et al.
Revised: (2016), and contact-rich manipulation Lin et al.

Original: ; Abu-Dakka et al.
Revised: ; Abu-Dakka et al.

Original: (2018), among others.
Revised: (2018), among others.

Original: Often, skills learned from demonstrations need to be refined – due to imperfect data – or adapted to comply with new task requirements.
Revised: Often, skills learned from demonstrations need to be refined – due to imperfect data – or adapted to comply with new task requirements.

Original: In this context, existing adaptation strategies for GMM policies either build a kernel method on top of the original skill model Huang et al.
Revised: In this context, existing adaptation strategies for GMM policies either build a kernel method on top of the original skill model Huang et al.

Original: (2019), or leverage reinforcement learning (RL) to adapt the policy itself (Arenz et al., 2020; Nematollahi et al., 2022).
Revised: (2019), or leverage reinforcement learning (RL) to adapt the policy itself (Arenz et al., 2020; Nematollahi et al., 2022).

Original: However, none of these techniques explicitly considered the structure of the GMM policy.
Revised: However, none of these techniques explicitly considered the structure of the GMM policy.

Original: Unlike the aforementioned approaches, we propose a policy optimization technique that explicitly considers the underlying GMM structure.
Revised: Unlike the aforementioned approaches, we propose a policy optimization technique that explicitly considers the underlying GMM structure.

Original: To do so, we exploit optimal transport theory (Santambrogio, 2015; Peyr ´ e & Cuturi, 2019), which allows us to view the set of GMM policies as a particular space of probability distributions GMM d .
Revised: To do so, we exploit optimal transport theory (Santambrogio, 2015; Peyr ´ e & Cuturi, 2019), which allows us to view the set of GMM policies as a particular space of probability distributions GMM d .

Original: Specifically, we leverage the idea of Chen et al.
Revised: Specifically, we leverage the idea of Chen et al.

Original: and Delon & Desolneux (2020) to view a GMM as a set of discrete measures (dirac masses) on the space of Gaussian distributions
Revised: and Delon & Desolneux (2020) to view a GMM as a set of discrete measures (dirac masses) on the space of Gaussian distributions

Original: G ( R d ) , which is endowed with a Wasserstein distance (see § 2).
Revised: G ( R d ) , which is endowed with a Wasserstein distance (see § 2).

Original: This allows us to formulate the policy optimization as a Wasserstein gradient flow (WGF) over the space of GMMs (as illustrated in Fig.
Revised: This allows us to formulate the policy optimization as a Wasserstein gradient flow (WGF) over the space of GMMs (as illustrated in Fig.

Original: 1 and explained in §3), where the policy updates are naturally guaranteed to be GMMs.
Revised: 1 and explained in §3), where the policy updates are naturally guaranteed to be GMMs.

Original: Moreover, we take advantage of the geometry of the Bures-Wasserstein manifold to optimize the Gaussian distributions of a GMM policy via Riemannian optimization.
Revised: Moreover, we take advantage of the geometry of the Bures-Wasserstein manifold to optimize the Gaussian distributions of a GMM policy via Riemannian optimization.

Original: We evaluate our approach over a set of different GMM policies featuring common robot skills: Reaching motions, collision-avoidance behaviors and multi-goal tasks (see § 4).
Revised: We evaluate our approach over a set of different GMM policies featuring common robot skills: Reaching motions, collision-avoidance behaviors and multi-goal tasks (see § 4).

Original: Our results show that our method outperforms common policy optimization baselines in terms of task success rate while providing low-variance solutions.
Revised: Our results show that our method outperforms common policy optimization baselines in terms of task success rate while providing low-variance solutions.

Original: Related Work: Richemond & Maginnis (2017) pioneered the idea of understanding policy optimization through the lens of optimal transport.
Revised: Related Work: Richemond & Maginnis (2017) pioneered the idea of understanding policy optimization through the lens of optimal transport.

Original: They interpreted the policy iteration as gradient flows by leveraging the implicit Euler scheme under a Wasserstein distance (see § 2), considering only 1 -step return settings.
Revised: They interpreted the policy iteration as gradient flows by leveraging the implicit Euler scheme under a Wasserstein distance (see § 2), considering only 1 -step return settings.

Original: They observed that the resulting policy optimization resembles the gradient flow of the Fokker-Planck equation (JKO scheme)
Revised: They observed that the resulting policy optimization resembles the gradient flow of the Fokker-Planck equation (JKO scheme)

Original: (Jordan et al., 1996).
Revised: (Jordan et al., 1996).

Original: In a similar spirit, Zhang et al.
Revised: In a similar spirit, Zhang et al.

Original: (2018) proposed to use WGFs to formulate policy optimization as a sequence of policy updates traveling along a gradient flow on the space of probability distributions until convergence.
Revised: (2018) proposed to use WGFs to formulate policy optimization as a sequence of policy updates traveling along a gradient flow on the space of probability distributions until convergence.

Original: To solve the WGF problem, the authors proposed a particle-based algorithm to approximate continuous density functions and subsequently derived the gradients for particle updates based on the JKO scheme.
Revised: To solve the WGF problem, the authors proposed a particle-based algorithm to approximate continuous density functions and subsequently derived the gradients for particle updates based on the JKO scheme.

Original: Although Zhang et al.
Revised: Although Zhang et al.

Original: (2018) considered general parametric policies, their method assumed a distribution over the policy parameters and did not consider a specific policy structure, which partially motivated their particle-based approximation.
Revised: (2018) considered general parametric policies, their method assumed a distribution over the policy parameters and did not consider a specific policy structure, which partially motivated their particle-based approximation.

Original: Recently, Mokrov et al.
Revised: Recently, Mokrov et al.

Original: (2021) tackled the computational burden of particle methods by leveraging input-convex neural networks to approximate the
Revised: (2021) tackled the computational burden of particle methods by leveraging input-convex neural networks to approximate the

Original: WGFs computation.
Revised: WGFs computation.

Original: They reformulated the well-known JKO optimization Jordan et al.
Revised: They reformulated the well-known JKO optimization Jordan et al.

Original: over probability measures by an optimization over convex functions.
Revised: over probability measures by an optimization over convex functions.

Original: Yet, this work remains a general solution for WFG computation and it did not address its use for policy optimization problems.
Revised: Yet, this work remains a general solution for WFG computation and it did not address its use for policy optimization problems.

Original: Aside from optimal transport approaches, Arenz et al.
Revised: Aside from optimal transport approaches, Arenz et al.

Original: (2020) proposed a trust-region variational inference for GMMs to approximate multimodal distributions.
Revised: (2020) proposed a trust-region variational inference for GMMs to approximate multimodal distributions.

Original: Although not originally designed for policy optimization, the authors elucidated a connection to learn GMMs of policy parameters in black-box RL.
Revised: Although not originally designed for policy optimization, the authors elucidated a connection to learn GMMs of policy parameters in black-box RL.

Original: 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.
Revised: 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.

Original: Nematollahi et al.
Revised: Nematollahi et al.

Original: (2022) proposed SAC-GMM, a hybrid model that employs the well-known SAC algorithm (Haarnoja et al., 2018) to refine dynamic skills encoded by GMMs.
Revised: (2022) proposed SAC-GMM, a hybrid model that employs the well-known SAC algorithm (Haarnoja et al., 2018) to refine dynamic skills encoded by GMMs.

Original: The SAC policy was designed to learn residuals on a single vectorized stack of GMM parameters, thus fully disregarding the GMM structure and the geometric constraints of its parameters.
Revised: The SAC policy was designed to learn residuals on a single vectorized stack of GMM parameters, thus fully disregarding the GMM structure and the geometric constraints of its parameters.

Original: Finally,  a Riemannian proximal policy optimization for GMMs wasrecently proposed by Wang et al.
Revised: Finally, two recent works share our idea of leveraging geometry in policy optimization: First, a Riemannian proximal policy optimization for GMMs was proposed by Wang et al.

Original: (2020), where the geometry induced by the GMM parameters was considered in the optimization via Riemannian gradients, similarly to our method.
Revised: (2020), where the geometry induced by the GMM parameters was considered in the optimization via Riemannian gradients, similarly to our method.

Original: The policy optimization was regularized by a Wasserstein distance to control the exploration-exploitation trade-off.
Revised: The policy optimization was regularized by a Wasserstein distance to control the exploration-exploitation trade-off.

Original: Despite this, their method did not formulate the policy optimization as an optimal transport problem (i.e., policy updates do not follow a WGF), asproposed in this paper, but it employed instead a classical non-convex optimization. 
Revised: However, their method did not formulate the policy optimization as an optimal transport problem, i.e. the policy updates do not follow a WGF, as in our approach, but it employed instead a classical non-convex optimization. Second, Moskovitz et al.

Revised: (2021) employed the Wasserstein natural gradient to exploit the local geometry induced by the Wasserstein regularization of behavioral policy optimization (Pacchiano et al., 2020).

Revised: 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 GMM d .

Original:  B ACKGROUND 2. W ASSERSTEIN GRADIENT FLOWS
Revised:  B ACKGROUND 2. W ASSERSTEIN GRADIENT FLOWS

Original: In Euclidean space a gradient flow is a smooth curve x : R → R d that satisfies the partial differential equation (PDE) ˙
Revised: In Euclidean space a gradient flow is a smooth curve x : R → R d that satisfies the partial differential equation (PDE) ˙

Original: x ( t ) = −∇ L ( x ( t )) for a given loss function L : R d → R and starting point x 0 at t = 

Revised: x ( t ) = −∇ L ( x ( t )) for a given loss function L : R d → R and starting point x 0 at t =

Original: (Santambrogio, 2015; 2017).
Revised: (Santambrogio, 2015; 2017).

Original: A solution can be found straightforwardly by forward discretization, leading to the well-known explicit Euler update scheme 
Revised: A solution can be found straightforwardly by forward discretization, leading to the well-known explicit Euler update scheme x τk +1 = x

Original: [Equation] , where λ denotes the learning rate and x τ indicates a discretization of the curve x ( t ) with discretization parameter τ .
Revised: k − λ ∇ L ( x τk ) , where λ denotes the learning rate and x τ indicates a discretization of the curve x ( t ) with discretization parameter τ .

Original: Alternatively, we can use a backward discretization, which leads to the following implicit Euler scheme
Revised: Alternatively, we can use a backward discretization, which leads to the following implicit Euler scheme

Original: [Equation]
Revised: [Equation]

Original: Eq. 1 is sometimes referred to as Minimizing Movement Scheme and can be used as an alternative characterization of a gradient flow.
Revised: Eq. 1 is sometimes referred to as Minimizing Movement Scheme and can be used as an alternative characterization of a gradient flow.

Original: This characterization is particularly interesting when we need to extend the concept of gradient flows to (non-Euclidean) general metric settings, since there is no notion of ∇ L in these cases (Santambrogio, 2015; Ambrosio et al., 2005).
Revised: This characterization is particularly interesting when we need to extend the concept of gradient flows to (non-Euclidean) general metric settings, since there is no notion of ∇ L in these cases (Santambrogio, 2015; Ambrosio et al., 2005).

Original: Eq.
Revised: Eq.

Original: 1 does not involve any gradients and can be expressed using only metric quantities.
Revised: 1 does not involve any gradients and can be expressed using only metric quantities.

Original: In this work, we are particularly interested in gradient flows in the
Revised: In this work, we are particularly interested in gradient flows in the

Original: L 2 -Wasserstein space, defined as the set of probability measures P ( X ) on a separable Banach space
Revised: L 2 -Wasserstein space, defined as the set of probability measures P ( X ) on a separable Banach space

Original: X (Panaretos & Zemel, 2020) and endowed with the L 2 -Wasserstein distance W 2 defined as
Revised: X (Panaretos & Zemel, 2020) and endowed with the L 2 -Wasserstein distance W 2 defined as

Original: [Equation]
Revised: [Equation]

Original: where µ, ν ∈ P ( X ) and γ ∈ P ( X 2 ) is defined to have the two marginals µ and ν .
Revised: where µ, ν ∈ P ( X ) and γ ∈ P ( X 2 ) is defined to have the two marginals µ and ν .

Original: A Generalized Minimizing Movement scheme characterizing gradient flows in the Wasserstein space can be written in analogy to Eq.
Revised: A Generalized Minimizing Movement scheme characterizing gradient flows in the Wasserstein space can be written in analogy to Eq.

Original: 1 as:
Revised: 1 as:

Original: [Equation]
Revised: [Equation]

Original: where L is a functional to be minimized on the Wasserstein space and π k ∈ P ( X ) .
Revised: where L is a functional to be minimized on the Wasserstein space and π k ∈ P ( X ) .

Original: In the following, we will omit the superscript τ for notational convenience.
Revised: In the following, we will omit the superscript τ for notational convenience.

Original: 2. R EINFORCEMENT L EARNING AS W ASSERSTEIN G RADIENT F LOWS
Revised: 2. R EINFORCEMENT L EARNING AS W ASSERSTEIN G RADIENT F LOWS

Original: Our view of the policy structure-aware optimization builds on the approach outlined by Richemond & Maginnis (2017), which in turn is based on the JKO scheme of Jordan et al.
Revised: Our view of the policy structure-aware optimization builds on the approach outlined by Richemond & Maginnis (2017), which in turn is based on the JKO scheme of Jordan et al.

Original: They proposed a formulation of 1 -step RL problems in terms of Wasserstein gradient flows.
Revised: They proposed a formulation of 1 -step RL problems in terms of Wasserstein gradient flows.

Original: In particular, they studied the evolution of a policy π under the influence of a free energy functional J of the form:
Revised: In particular, they studied the evolution of a policy π under the influence of a free energy functional J of the form:

Original: [Equation]
Revised: [Equation]

Original: where K r ( π ) denotes the inner energy of the system, here determined by the reward r ( s , a ) .
Revised: where K r ( π ) denotes the inner energy of the system, here determined by the reward r ( s , a ) .

Original: Moreover, H ( π ) is the entropy of the policy π ( a | s ) , with s and a denoting the state and action, respectively.
Revised: Moreover, H ( π ) is the entropy of the policy π ( a | s ) , with s and a denoting the state and action, respectively.

Original: Thus Eq.
Revised: Thus Eq.

Original: 4 can be recognized as the usual objective in 1 -step RL settings with entropy regularization.
Revised: 4 can be recognized as the usual objective in 1 -step RL settings with entropy regularization.

Original: It is well known that the evolution of probability densities under a free energy of this form is properly described by a PDE known as the Fokker-Planck equation.
Revised: It is well known that the evolution of probability densities under a free energy of this form is properly described by a PDE known as the Fokker-Planck equation.

Original: Richemond & Maginnis (2017) exploited the result of Jordan et al.
Revised: Richemond & Maginnis (2017) exploited the result of Jordan et al.

Original: (1996), which stated that this evolution can be interpreted as the gradient flow of the functional J in Wasserstein space.
Revised: (1996), which stated that this evolution can be interpreted as the gradient flow of the functional J in Wasserstein space.

Original: This flow is characterized by the following minimizing movement scheme
Revised: This flow is characterized by the following minimizing movement scheme

Original: [Equation]
Revised: [Equation]

Original: which naturally provides iterative updates for the policy π .
Revised: which naturally provides iterative updates for the policy π .

Original: 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.
Revised: 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.

Original: 2. T HE L 2 -W ASSERSTEIN DISTANCE BETWEEN G AUSSIAN M IXTURE M ODELS (GMM S )
Revised: 2. T HE L 2 -W ASSERSTEIN DISTANCE BETWEEN G AUSSIAN M IXTURE M ODELS (GMM S )

Original: In this paper, we consider policies π ( x ) that build on a GMM structure, i.e., π ( x )
Revised: In this paper, we consider policies π ( x ) that build on a GMM structure, i.e., π ( x )

Original: = (cid:80)
Revised: = (cid:80)

Original: Ni =1 ω i N ( x ; µ i , Σ i ) , where N denotes a multivariate Gaussian distribution with mean µ i and covariance matrix Σ i , and ω
Revised: Ni =1 ω i N ( x ; µ i , Σ i ) , where N denotes a multivariate Gaussian distribution with mean µ i and covariance matrix Σ i , and ω

Original: i are the weights of the N individual Gaussian components, which are subject to (cid:80) i ω i
Revised: i are the weights of the N individual Gaussian components, which are subject to (cid:80) i ω i

Original: = 1 .
Revised: = 1 .

Original: In the following, we will write ˆ µ , ˆΣ and ˆ ω to denote the stacked means, covariance matrices and weights of the N components.
Revised: In the following, we will write ˆ µ , ˆΣ and ˆ ω to denote the stacked means, covariance matrices and weights of the N components.

Original: Therefore, we do not consider WGFs on the full manifold of probability distributions (Wasserstein space) P ( R d ) but rather focus on WGFs evolving on the submanifold of GMMs, that is GMM d ⊂ P ( R d ) .
Revised: Therefore, we do not consider WGFs on the full manifold of probability distributions (Wasserstein space) P ( R d ) but rather focus on WGFs evolving on the submanifold of GMMs, that is GMM d ⊂ P ( R d ) .

Original: Following Chen et al.
Revised: Following Chen et al.

Original: Delon & Desolneux (2020), we can approximately describe this submanifold as a discrete distribution over the space of Gaussian distributions equipped with the Wasserstein metric.
Revised: Delon & Desolneux (2020), we can approximately describe this submanifold as a discrete distribution over the space of Gaussian distributions equipped with the Wasserstein metric.

Original: This in turn can be identified with the Bures-Wasserstein manifold which is the product manifold R d × S d ++ , where
Revised: This in turn can be identified with the Bures-Wasserstein manifold which is the product manifold R d × S d ++ , where

Original: S d ++ denotes the Riemannian manifold of d -dimensional symmetric positive definite matrices.
Revised: S d ++ denotes the Riemannian manifold of d -dimensional symmetric positive definite matrices.

Original: The corresponding approximated Wasserstein distance between two GMMs π 1 , π 2 is given by
Revised: The corresponding approximated Wasserstein distance between two GMMs π 1 , π 2 is given by

Original: [Equation]
Revised: [Equation]

Original: where U ( ω 1 , ω 2 ) = { P ∈ R N × N + | P 1 N = ω 1 , P T 1 N = ω 2 } with 1 N denoting an N -dimensional vector of ones.
Revised: where U ( ω 1 , ω 2 ) = { P ∈ R N × N + | P 1 N = ω 1 , P T 1 N = ω 2 } with 1 N denoting an N -dimensional vector of ones.

Original: The Wasserstein distance between two Gaussian distributions in Eq.
Revised: The Wasserstein distance between two Gaussian distributions in Eq.

Original: 6 can be computed analytically as follows
Revised: 6 can be computed analytically as follows

Original: [Equation]
Revised: [Equation]

Original: 2. L EARNING GMM POLICIES FROM DEMONSTRATIONS
Revised: 2. L EARNING GMM POLICIES FROM DEMONSTRATIONS

Original: A popular approach in RL – particularly in the robotics domain – to reduce the number of policy rollouts in the environment is to warm-start the policy with a set of demonstrations provided by an expert.
Revised: A popular approach in RL – particularly in the robotics domain – to reduce the number of policy rollouts in the environment is to warm-start the policy with a set of demonstrations provided by an expert.

Original: In this work we choose to represent our policy via a GMM.
Revised: In this work we choose to represent our policy via a GMM.

Original: We assume that demonstrations are provided as a set of trajectories τ of state-action pairs τ
Revised: We assume that demonstrations are provided as a set of trajectories τ of state-action pairs τ

Original: = { ( s 0 , a 0 ) , ( s 1 , a 1 ) , . . . 
Revised: = { ( s 0 , a 0 ) , ( s 1 , a 1 ) , . . . [Equation]

Original: ( s T , a T ) } .
Revised: s T [Equation] T ) } .

Original: To initialize our policy, we first use the Expectation-Maximization (EM) algorithm to fit a GMM, in the joint state-action space, to the demonstrations.
Revised: To initialize our policy, we first use the Expectation-Maximization (EM) algorithm to fit a GMM, in the joint state-action space, to the demonstrations.

Original: This results in a mixture distribution π ( s , a )
Revised: This results in a mixture distribution π ( s , a )

Original: = (cid:80) Ni =1 ω i N (cid:0)
Revised: = (cid:80) Ni =1 ω i N (cid:0)

Original: [ s a ] T ; µ i , Σ i (cid:1) from which a policy can be obtained by conditioning on the state s , as follows
Revised: [ s a ] T ; µ i , Σ i (cid:1) from which a policy can be obtained by conditioning on the state s , as follows

Original: [Equation]
Revised: [Equation]

Original: In the context of GMMs, this is also known as Gaussian Mixture Regression (GMR) (Ghahramani & Jordan, 1994).
Revised: In the context of GMMs, this is also known as Gaussian Mixture Regression (GMR) (Ghahramani & Jordan, 1994).

Original: The resulting conditional distribution is another GMM on action space with statedependent parameters, given by
Revised: The resulting conditional distribution is another GMM on action space with statedependent parameters, given by

Original: [Equation]
Revised: [Equation]

Original: Details on computation of Eq.
Revised: Details on computation of Eq.

Original: 9 from the original GMM are given in App. A.1.
Revised: 9 from the original GMM are given in App. A.1.

Original:  W ASSERSTEIN G RADIENT F LOWS FOR GMM P OLICY O PTIMIZATION
Revised:  W ASSERSTEIN G RADIENT F LOWS FOR GMM P OLICY O PTIMIZATION

Original: In this work, we focus on multi-step RL tasks for policy adaptation.
Revised: In this work, we focus on multi-step RL tasks for policy adaptation.

Original: We consider a finite-horizon
Revised: We consider a finite-horizon

Original: Markov Decision Process (MDP) with continuous state and action spaces S ∈ R n and A ∈ R m , transition and reward functions p ( s t +1 | s t , a t ) and r ( s t , a t ) , initial state distribution ρ ( s 0 ) and a discount factor γ .
Revised: Markov Decision Process (MDP) with continuous state and action spaces S ∈ R n and A ∈ R m , transition and reward functions p ( s t +1 | s t , a t ) and r ( s t , a t ) , initial state distribution ρ ( s 0 ) and a discount factor γ .

Original: Further, we assume to have an initial policy π ( a t | s t ) , which is to be adapted by optimizing some objective function K r ( π ) .
Revised: Further, we assume to have an initial policy π ( a t | s t ) , which is to be adapted by optimizing some objective function K r ( π ) .

Original: As stated in § 1, this problem arises in robot learning settings where a policy learned via imitation learning (e.g., LfD) needs to be adapted to new objectives or unseen environmental conditions.
Revised: As stated in § 1, this problem arises in robot learning settings where a policy learned via imitation learning (e.g., LfD) needs to be adapted to new objectives or unseen environmental conditions.

Original: To promote exploration and avoid premature convergence to suboptimal policies, we leverage maximum entropy RL (Eysenbach & Levine, 2022) by adding an entropy term H ( π ) to the objective.
Revised: To promote exploration and avoid premature convergence to suboptimal policies, we leverage maximum entropy RL (Eysenbach & Levine, 2022) by adding an entropy term H ( π ) to the objective.

Original: Thus, the overall objective has the form of a free energy functional (resembling Eq. 4) and can be written as
Revised: Thus, the overall objective has the form of a free energy functional (resembling Eq. 4) and can be written as

Original: [Equation]
Revised: [Equation]

Original: ) where β is a hyperparameter and K r ( π ) corresponds to the usual cumulative return
Revised: ) where β is a hyperparameter and K r ( π ) corresponds to the usual cumulative return

Original: K r
Revised: K r

Original: [Equation]
Revised: [Equation]

Original: The evolution of the policy π ( a t | s t ) over the course of the optimization can be described as a flow of a probability distribution in Wasserstein space.
Revised: The evolution of the policy π ( a t | s t ) over the course of the optimization can be described as a flow of a probability distribution in Wasserstein space.

Original: 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 L 2 -Wasserstein distance to constrain the policy updates, which is important to guarantee stability in policy optimization (Schulman et al., 2015; 2017; Otto et al., 2021); (iii) By constraining to specific submanifolds of the Wasserstein space, in this case GMMs, we can impose additional structural properties on the policy optimization.
Revised: 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 L 2 -Wasserstein distance to constrain the policy updates, which is important to guarantee stability in policy optimization (Schulman et al., 2015; 2017; Otto et al., 2021); (iii) By constraining to specific submanifolds of the Wasserstein space, in this case GMMs, we can impose additional structural properties on the policy optimization.

Original: Since our objective in Eq.
Revised: Since our objective in Eq.

Original: 10 has the form of the free energy functional studied by Richemond &
Revised: 10 has the form of the free energy functional studied by Richemond &

Original: Maginnis (2017); Jordan et al. (1996), we can leverage the iterative updates scheme of Eq.
Revised: Maginnis (2017); Jordan et al. (1996), we can leverage the iterative updates scheme of Eq.

Original: 5 to formulate the evolution of our policy iteration under the flow generated by Eq.
Revised: 5 to formulate the evolution of our policy iteration under the flow generated by Eq.

Original: 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 d .
Revised: 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 d .

Original: We refer the interested reader to
Revised: We refer the interested reader to

Original: App. A.3, where we provide the explicit form of J ( π ) of Eq.
Revised: App. A.3, where we provide the explicit form of J ( π ) of Eq.

Original: 10 for the GMM case.
Revised: 10 for the GMM case.

Original: 3. P OLICY OPTIMIZATION
Revised: 3. P OLICY OPTIMIZATION

Original: To begin with, we leverage the approximation that describes the GMM submanifold as a discrete distribution over the space of Gaussian distributions G ( R d ) , equipped with the Wasserstein metric (Chen et al., 2019).
Revised: To begin with, we leverage the approximation that describes the GMM submanifold as a discrete distribution over the space of Gaussian distributions G ( R d ) , equipped with the Wasserstein metric (Chen et al., 2019).

Original: Consequently, our policy optimization problem naturally splits into an optimization over the ( N − 1) -dimensional simplex and an optimization on the N -fold product of the
Revised: Consequently, our policy optimization problem naturally splits into an optimization over the ( N − 1) -dimensional simplex and an optimization on the N -fold product of the

Original: Bures-Wasserstein manifold ( BW N ), i.e. the product manifold (cid:0) R d × S d
Revised: Bures-Wasserstein manifold ( BW N ), i.e. the product manifold (cid:0) R d × S d

Original: ++ (cid:1) N .
Revised: ++ (cid:1) N .

Original: The former corresponds to the GMM weights while the latter applies to the set of Gaussian distributions parameters.
Revised: The former corresponds to the GMM weights while the latter applies to the set of Gaussian distributions parameters.

Original: Note that the identification with the BW N manifold allows us to perform the optimization directly on the parameter space.
Revised: Note that the identification with the BW N manifold allows us to perform the optimization directly on the parameter space.

Original: 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.
Revised: 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.

Original: (ii)
Revised: (ii)

Original: As Chen et al.
Revised: As Chen et al.

Original: (2019) show, we can guarantee that the policy optimized via Eq.
Revised: (2019) show, we can guarantee that the policy optimized via Eq.

Original: 6 remains a GMM.
Revised: 6 remains a GMM.

Original: (iii)
Revised: (iii)

Original: Unlike the full Wasserstein space 1 , the resulting product manifold is a true Riemannian manifold such that we can leverage the machinery of Riemannian optimization.
Revised: Unlike the full Wasserstein space 1 , the resulting product manifold is a true Riemannian manifold such that we can leverage the machinery of Riemannian optimization.

Original: Importantly, working in the parameter space allows us to apply an explicit Euler scheme instead of the implicit formulation of Eq.
Revised: Importantly, working in the parameter space allows us to apply an explicit Euler scheme instead of the implicit formulation of Eq.

Original: According to the above-mentioned split, we formulate the policy optimization as an EM-like twostep procedure that alternates between the Gaussian parameters (i.e. means and covariance matrices) and the GMM weights.
Revised: According to the above-mentioned split, we formulate the policy optimization as an EM-like twostep procedure that alternates between the Gaussian parameters (i.e. means and covariance matrices) and the GMM weights.

Original: To optimize the former, we propose to leverage the Riemannian structure of the BW manifold to reformulate the updates as a forward discretization, similarly to Chen &
Revised: To optimize the former, we propose to leverage the Riemannian structure of the BW manifold to reformulate the updates as a forward discretization, similarly to Chen &

Original: Li (2020).
Revised: Li (2020).

Original: In other words, by embedding the Gaussian components of the GMM policy in a Riemannian manifold, the Wasserstein gradient flow in the implicit form of Eq.
Revised: In other words, by embedding the Gaussian components of the GMM policy in a Riemannian manifold, the Wasserstein gradient flow in the implicit form of Eq.

Original: 5 can be approximated by an explicit Euler update scheme according to the BW metric (further details are provided in
Revised: 5 can be approximated by an explicit Euler update scheme according to the BW metric (further details are provided in

Original: App.
Revised: App.

Original: A.4).
Revised: A.4).

Original: This allows us to leverage the expressions of the Riemannian gradient and exponential map of the BW manifold (Malag`o et al., 2018; Han et al., 2021).
Revised: This allows us to leverage the expressions of the Riemannian gradient and exponential map of the BW manifold (Malag`o et al., 2018; Han et al., 2021).

Original: Thus, the optimization boils down to Riemannian gradient descent where the gradient is defined w.r.t the Bures-Wasserstein metric.
Revised: Thus, the optimization boils down to Riemannian gradient descent where the gradient is defined w.r.t the Bures-Wasserstein metric.

Original: In particular, we use the expression for Riemannian gradient, metric and exponential map used in (Han et al., 2021).
Revised: In particular, we use the expression for Riemannian gradient, metric and exponential map used in (Han et al., 2021).

Original: Formally, the resulting updates for the Gaussian parameters of the GMM follow the
Revised: Formally, the resulting updates for the Gaussian parameters of the GMM follow the

Original: Riemannian gradient descent scheme given by:
Revised: Riemannian gradient descent scheme given by:

Original: [Equation]
Revised: [Equation]

Original: where grad denotes the Riemannian gradient w.r.t.
Revised: where grad denotes the Riemannian gradient w.r.t.

Original: the Bures-Wasserstein metric, R x : T x M → M denotes the retraction operator, which maps a point on the tangent space
Revised: the Bures-Wasserstein metric, R x : T x M → M denotes the retraction operator, which maps a point on the tangent space

Original: T x M back to the manifold
Revised: T x M back to the manifold

Original: M ≡ BW (Boumal, 2022).
Revised: M ≡ BW (Boumal, 2022).

Original: Moreover, λ is a learning rate and π k def =
Revised: Moreover, λ is a learning rate and π k def =

Original: π ( ˆ µ k , ˆΣ k , ˆ ω k ) .
Revised: π ( ˆ µ k , ˆΣ k , ˆ ω k ) .

Original: The Euclidean gradients of J ( π ) required for computing grad can be obtained using a likelihood ratio estimator (a.k.a score function estimator or REINFORCE)
Revised: The Euclidean gradients of J ( π ) required for computing grad can be obtained using a likelihood ratio estimator (a.k.a score function estimator or REINFORCE)

Original: (Williams, 2004) and are provided in App. A.3.
Revised: (Williams, 2004) and are provided in App. A.3.

Original: Concerning the GMM weights, we first reparameterize them as ω j =
Revised: Concerning the GMM weights, we first reparameterize them as ω j =

Original: [Equation]
Revised: [Equation]

Original: and optimize w.r.t.
Revised: and optimize w.r.t.

Original: the new parameters η j , j = 1 ...N , which unlike ˆ ω are unconstrained.
Revised: the new parameters η j , j = 1 ...N , which unlike ˆ ω are unconstrained.

Original: For this optimization we employ the implicit Euler scheme:
Revised: For this optimization we employ the implicit Euler scheme:

Original: [Equation]
Revised: [Equation]

Original: where π k +1 ( ˆ η ) def =
Revised: where π k +1 ( ˆ η ) def =

Original: π ( ˆ µ k +1 , ˆΣ k +1 , ˆ η ) .
Revised: π ( ˆ µ k +1 , ˆΣ k +1 , ˆ η ) .

Original: We minimize Eq.
Revised: We minimize Eq.

Original: 13 by gradient descent w.r.t.
Revised: 13 by gradient descent w.r.t.

Original: η as follows:
Revised: η as follows:

Original: [Equation]
Revised: [Equation]

Original: The gradient of J ( π ) can be obtained analytically using a likelihood ratio estimator.
Revised: The gradient of J ( π ) can be obtained analytically using a likelihood ratio estimator.

Original: For the Wasserstein term, we first compute the gradient w.r.t.
Revised: For the Wasserstein term, we first compute the gradient w.r.t.

Original: the weights via the Sinkhorn algorithm (Cuturi &
Revised: the weights via the Sinkhorn algorithm (Cuturi &

Original: Doucet, 2014), from which the gradient w.r.t η can be then obtained via the chain rule.
Revised: Doucet, 2014), from which the gradient w.r.t η can be then obtained via the chain rule.

Original: Note that we have to rely on the Sinkhorn algorithm here since there is no analytic solution available for the
Revised: Note that we have to rely on the Sinkhorn algorithm here since there is no analytic solution available for the

Original: Wasserstein distance between discrete distributions, unlike the above case of the Gaussian components.
Revised: Wasserstein distance between discrete distributions, unlike the above case of the Gaussian components.

Original: Consequently, we cannot compute the corresponding gradients.
Revised: Consequently, we cannot compute the corresponding gradients.

Original: 3. I MPLEMENTATION P IPELINE
Revised: 3. I MPLEMENTATION P IPELINE

Original: 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.
Revised: 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.

Original: Then, we use the collected interaction trajectories to compute a sample-based estimate of the functional K r ( π ) and its gradients w.r.t the policy parameters, as explained in § 3.1.
Revised: Then, we use the collected interaction trajectories to compute a sample-based estimate of the functional K r ( π ) and its gradients w.r.t the policy parameters, as explained in § 3.1.

Original: An optimization step consists of alternating between optimizing the Gaussian parameters using 12, and updating the weights via 14.
Revised: An optimization step consists of alternating between optimizing the Gaussian parameters using 12, and updating the weights via 14.

Original: For the optimization of the Gaussian parameters we leverage Pymanopt (Townsend et al., 2016) for Riemannian optimization.
Revised: For the optimization of the Gaussian parameters we leverage Pymanopt (Townsend et al., 2016) for Riemannian optimization.

Original: We extended this library by implementing the Bures-Wasserstein manifold based on the expressions provided by Han et al.
Revised: We extended this library by implementing the Bures-Wasserstein manifold based on the expressions provided by Han et al.

Original: (2021) (see App.
Revised: (2021) (see App.

Original: A.2 for details).
Revised: A.2 for details).

Original: Furthermore, we added a custom line-search routine that accounts for a constraint on the Wasserstein distance between the old and the optimized GMM, as to our knowledge such a search method does not exist in out-of-the-box optimizers.
Revised: Furthermore, we added a custom line-search routine that accounts for a constraint on the Wasserstein distance between the old and the optimized GMM, as to our knowledge such a search method does not exist in out-of-the-box optimizers.

Original: The details of this custom line-search can be found in Algorithm 2 in
Revised: The details of this custom line-search can be found in Algorithm 2 in

Original: App. A.5.
Revised: App. A.5.

Original: Regarding the optimization of the GMM weights, we 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.
Revised: Regarding the optimization of the GMM weights, we 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.

Original: the weights in Eq.
Revised: the weights in Eq.

Original: The full policy optimization finishes if either the objective stops improving or the Wasserstein distance between the old and optimized GMMs exceeds a predefined threshold, which is chosen experimentally.
Revised: The full policy optimization finishes if either the objective stops improving or the Wasserstein distance between the old and optimized GMMs exceeds a predefined threshold, which is chosen experimentally.

Original: Afterwards, fresh rollouts are performed with the updated policy and the aforementioned two-step procedure starts over.
Revised: Afterwards, fresh rollouts are performed with the updated policy and the aforementioned two-step procedure starts over.

Original: This optimization loop is repeated until a task-dependent success criterion has been fulfilled.
Revised: This optimization loop is repeated until a task-dependent success criterion has been fulfilled.

Original: We summarize the proposed optimization in Algorithm 1.
Revised: We summarize the proposed optimization in Algorithm 1.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: We tested our approach in three different robotic settings: a reaching skill, a collision-free trajectory tracking, and a multiple-goal task.
Revised: We tested our approach in three different robotic settings: a reaching skill, a collision-free trajectory tracking, and a multiple-goal task.

Original: All the tasks are represented in a 2 D operational space.
Revised: All the tasks are represented in a 2 D operational space.

Original: The robot
Revised: The robot

Original:  Input : initial policy π ( a | s ) 1: while not goal reached do 2:
Revised: Algorithm Input : initial policy π ( a | s ) 1: while not goal reached do 2:

Original: Rollout policy π ( a | s ) in the environment for M episodes to collect interaction trajectories τ
Revised: Rollout policy π ( a | s ) in the environment for M episodes to collect interaction trajectories τ

Original: = { ( s 0 , a 0 , r 0 ) , ( s 1 , a 1 , r 1 ) , . . .
Revised: = { ( s 0 , a 0 , r 0 ) , ( s 1 , a 1 , r 1 ) , . . .

Original: , ( s T , a T , r T ) } Mm = 3: repeat 4:
Revised: , ( s T , a T , r T ) } Mm = 3: repeat 4:

Original: Update Gaussian components parameters ˆ µ , ˆΣ using Riemannian optimization (12), where λ ls is determined via line-search (see §3.2).
Revised: Update Gaussian components parameters ˆ µ , ˆΣ using Riemannian optimization (12), where λ ls is determined via line-search (see §3.2).

Original: 5: until convergence 6: repeat 7:
Revised: 5: until convergence 6: repeat 7:

Original: Update GMM weights ˆ ω via gradient descent on the free energy objective 10, using8: until converged 9: end while motion policies were initially learned from human demonstrations collected on a simple Python graphical interface.
Revised: Update GMM weights ˆ ω via gradient descent on the free energy objective 10, using8: until converged 9: end while motion policies were initially learned from human demonstrations collected on a simple Python graphical interface.

Original: We assumed we were given M demonstrations, each of which contained T m data points for a dataset of N = (cid:80) m T m total observations τ
Revised: We assumed we were given M demonstrations, each of which contained T m data points for a dataset of N = (cid:80) m T m total observations τ

Original: = { ( s t , a t ) } N t =1 .
Revised: = { ( s t , a t ) } N t =1 .

Original: The state s and action a correspond to the robot end-effector position x ∈ R 2 and velocity ˙ x ∈ R 2 .
Revised: The state s and action a correspond to the robot end-effector position x ∈ R 2 and velocity ˙ x ∈ R 2 .

Original: The GMM models were trained via classical Expectation-Maximization.
Revised: The GMM models were trained via classical Expectation-Maximization.

Original: The policy rollout consists of sampling a velocity action a t ∼ π ( a t | s t ) using Eq. 9, and subsequently commanding the robot via a Cartesian velocity controller at a frequency of 100Hz .
Revised: The policy rollout consists of sampling a velocity action a t ∼ π ( a t | s t ) using Eq. 9, and subsequently commanding the robot via a Cartesian velocity controller at a frequency of 100Hz .

Original: For all the experiments, we used the Robotics Toolbox for
Revised: For all the experiments, we used the Robotics Toolbox for

Original: Python (Corke & Haviland, 2021) to simulate the robotic environments.
Revised: Python (Corke & Haviland, 2021) to simulate the robotic environments.

Original: To show the importance of accounting for the policy structure in RL settings, we compared our method against two structure-unaware baselines:
Revised: To show the importance of accounting for the policy structure in RL settings, we compared our method against two structure-unaware baselines:

Original: PPO (Schulman et al., 2017) and SAC-GMM (Nematollahi et al., 2022).
Revised: PPO (Schulman et al., 2017) and SAC-GMM (Nematollahi et al., 2022).

Original: As PPO was not originally designed to directly optimize the parameters of a previously-learned GMM policy, we designed the policy actions to represent (usually small) corrections to the GMM parameters, i.e. a = [∆ ω ∆ˆ µ ∆ vec( ˆ Σ )] , following the same methodology as SAC-GMM (Nematollahi et al., 2022).
Revised: As PPO was not originally designed to directly optimize the parameters of a previously-learned GMM policy, we designed the policy actions to represent (usually small) corrections to the GMM parameters, i.e. a = [∆ ω ∆ˆ µ ∆ vec( ˆ Σ )] , following the same methodology as SAC-GMM (Nematollahi et al., 2022).

Original: The PPO and SAC implementations correspond to the code provided by Stable-Baselines

Revised: The PPO and SAC implementations correspond to the code provided by Stable-Baselines

Original: (Raffin et al., 2021), whose policies are parametrized by MLP networks.
Revised: (Raffin et al., 2021), whose policies are parametrized by MLP networks.

Original: 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.
Revised: 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.

Original: Later, we proceed as described earlier, namely, the updated GMM is used to compute the velocity action via Eq. 9.
Revised: Later, we proceed as described earlier, namely, the updated GMM is used to compute the velocity action via Eq. 9.

Original: For comparison purposes, we report statistical results for the three considered settings over 5 runs for the task success rate and solution variance.
Revised: For comparison purposes, we report statistical results for the three considered settings over 5 runs for the task success rate and solution variance.

Original: We tuned the baselines separately for each task using Optuna (Akiba et al., 2019).
Revised: We tuned the baselines separately for each task using Optuna (Akiba et al., 2019).

Revised: 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.6.2).

Revised: Last but not least, we tested our approach on a 3 D version of the collision-free task performed by a 7 -DoF Franka Emika Panda robot in a virtual environment as reported in App. A.6.3.

Original: 4. T ASKS DESCRIPTION
Revised: 4. T ASKS DESCRIPTION

Original: 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 a new target located midway and above the previouslylearned L-shape trajectories.
Revised: 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 a new target located midway and above the previouslylearned L-shape trajectories.

Original: The initial policy, shown in Fig.
Revised: The initial policy, shown in Fig.

Original: 2- left and Fig.
Revised: 2- left and Fig.

Original: 7a, was learned from 12 demonstrations and encoded by a 7 -component GMM.
Revised: 12a, was learned from 12 demonstrations and encoded by a 7 -component GMM.

Original: 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.
Revised: 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.

Original: We also added a sparse penalty term that punishes rollouts leading to significantly divergent trajectories.
Revised: We also added a sparse penalty term that punishes rollouts leading to significantly divergent trajectories.

Original: Convergence is achieved when a minimum average position error w.r.t the target – computed over an episode – is reached.
Revised: Convergence is achieved when a minimum average position error w.r.t the target – computed over an episode – is reached.

Original: Collision-avoidance Task:
Revised: Collision-avoidance Task:

Original: This task consists of: (1) learning an initial GMM policy of a linear reaching motion, and (2) adapting the GMM policy to reach a new horizontally-translated target while avoiding to collide with two spherical obstacles located midway between the initial robot position and the new target.
Revised: This task consists of: (1) learning an initial GMM policy of a linear reaching motion, and (2) adapting the GMM policy to reach a new horizontally-translated target while avoiding to collide with two spherical obstacles located midway between the initial robot position and the new target.

Original: The initial GMM policy was learned from 10 human demonstrations and represented by a 3 -component GMM, as shown in Fig.
Revised: The initial GMM policy was learned from 10 human demonstrations and represented by a 3 -component GMM, as shown in Fig.

Original: 2- middle and Fig. 7b.
Revised: 2- middle and Fig. 12b.

Original: For policy optimization, we defined a sparse reward as a function of the position error between the robot end-effectorposition and the target at the end of the rollout.
Revised: For policy optimization, we defined a sparse reward as a function of the position error between the robot end-effector

Original: 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

Original: [Figure]
Revised: [Figure]

Original:  second term penalizes rollouts with significantly divergent trajectories.
Revised: 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.

Original: Convergence is determined by a minimum average position error w.r.t the target computed over an episode.
Revised: Convergence is determined by a minimum average position error w.r.t the target computed over an episode.

Original: Multiple-goal Task:
Revised: Multiple-goal Task:

Original: 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 a new target located close to one of the previous task goals.
Revised: 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 a new target located close to one of the previous task goals.

Original: The intended adapted behavior should make the robot go through the most relevant GMM components according to the new target location.
Revised: The intended adapted behavior should make the robot go through the most relevant GMM components according to the new target location.

Original: The initial GMM policy was learned fromdemonstrations and encoded by a 6 -component GMM, as shown in Fig.
Revised: The initial GMM policy was learned fromdemonstrations and encoded by a 6 -component GMM, as shown in Fig.

Original: 2- right and Fig. 7c.
Revised: 2- right and Fig. 12c.

Original: To optimize the initial GMM policy, we specified a sparse reward based on the position error between the robot end-effector position and the chosen target at the end of the rollout.
Revised: To optimize the initial GMM policy, we specified a sparse reward based on the position error between the robot end-effector position and the chosen target at the end of the rollout.

Original: Similar to the previous experiments, we added a sparse penalty term to penalize rollouts generating significantly divergent trajectories.
Revised: Similar to the previous experiments, we added a sparse penalty term to penalize rollouts generating significantly divergent trajectories.

Original: Again, the policy optimization converges when the average position error w.r.t the chosen target reaches a minimum threshold.
Revised: Again, the policy optimization converges when the average position error w.r.t the chosen target reaches a minimum threshold.

Original: 4. R ESULTS A NALYSIS
Revised: 4. R ESULTS A NALYSIS

Original: The reaching task tested our method’s ability to adapt a previously-learned reaching skill to a new goal, located at (6 . 0 , − 6 . 5) (cf.
Revised: The reaching task tested our method’s ability to adapt a previously-learned reaching skill to a new goal, located at (6 . 0 , − 6 . 5) (cf.

Original: Fig.
Revised: Fig.

Original: 7a- left ).
Revised: 12a- left ).

Original: Achieving this required to adapt the Gaussian parameters of mainly the last four GMM components, while the other ones remained unchanged.
Revised: Achieving this required to adapt the Gaussian parameters of mainly the last four GMM components, while the other ones remained unchanged.

Original: 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.
Revised: 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.

Original: Figure 3- left shows that our method achieved a success rate of 1 after approximately 70000 environment interactions.
Revised: Figure 3- left shows that our method achieved a success rate of 1 after approximately 70000 environment interactions.

Original: Despite PPO was also able to complete the task reliably, it required many more environment steps (cf. Fig.
Revised: Despite PPO was also able to complete the task reliably, it required many more environment steps (cf. Fig.

Original: 5- left ).
Revised: 5- left ).

Original: In sharp contrast, SAC did not reach any improvement.
Revised: In sharp contrast, SAC did not reach any improvement.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: In the collision-avoidance task, we tested whether our method was able to adapt a trajectory tracking skill in order to avoid collisions with newly added obstacles.
Revised: In the collision-avoidance task, we tested whether our method was able to adapt a trajectory tracking skill in order to avoid collisions with newly added obstacles.

Original: These were placed in such a way that the robot was forced to move its end-effector through a narrow path between the obstacles (cf.
Revised: These were placed in such a way that the robot was forced to move its end-effector through a narrow path between the obstacles (cf.

Original: Fig.
Revised: Fig.

Original: 2- middle ).
Revised: 2- middle ).

Original: 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.
Revised: 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.

Original: Figure
Revised: Figure

Original: 3- middle shows that our method solved this task reliably after comparatively few environment interactions.
Revised: 3- middle shows that our method solved this task reliably after comparatively few environment interactions.

Original: Although PPO also achieved a success rate of 1 , it took 6 times more environment steps than our method.
Revised: Although PPO also achieved a success rate of 1 , it took 6 times more environment steps than our method.

Original: SAC only reached an average success rate of 0 .
Revised: SAC only reached an average success rate of 0 .

Original: 8 , however with high variance (cf.
Revised: 8 , however with high variance (cf.

Original: Fig.
Revised: Fig.

Original: 5- middle ).
Revised: 5- middle ).

Original: These results again show the importance of the constraints on the policy updates.
Revised: These results again show the importance of the constraints on the policy updates.

Original: The huge discrepancy in the required environment steps between our method and PPO further emphasizes the importance of taking the GMM structure into account in the policy optimization.
Revised: The huge discrepancy in the required environment steps between our method and PPO further emphasizes the importance of taking the GMM structure into account in the policy optimization.

Original: [Figure]
Revised: [Figure]

Original: 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.
Revised: 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.

Original: The initial skill comprised reaching motions to two different goals and an execution of this skill results in reaching one of them, depending on the sampling noise (cf. Fig. 7c).
Revised: The initial skill comprised reaching motions to two different goals and an execution of this skill results in reaching one of them, depending on the sampling noise (cf. Fig. 12c).

Original: The easiest way to adapt the policy to reach only one of 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.
Revised: The easiest way to adapt the policy to reach only one of 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.

Original: As shown in Fig.
Revised: As shown in Fig.

Original: 3right , our method again quickly achieved a success rate of 1 .
Revised: 3right , our method again quickly achieved a success rate of 1 .

Original: PPO required substantially many more environment steps, while SAC was not able to solve the task.
Revised: PPO required substantially many more environment steps, while SAC was not able to solve the task.

Original: In Fig.
Revised: In Fig.

Original: 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 a success rate of 1 , thus prioritizing sample efficiency.
Revised: 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 a success rate of 1 , thus prioritizing sample efficiency.

Original: The plots show that our method exhibits a very low solution variance.
Revised: The plots show that our method exhibits a very low solution variance.

Original: Both baselines varied largely, except for the reaching task, where all SAC runs collapse to a success rate of 0 .
Revised: Both baselines varied largely, except for the reaching task, where all SAC runs collapse to a success rate of 0 .

Original: These results show that our method, despite showing large variance at the start, was able to quickly reduce the variance and converge reliably to a good success rate.
Revised: These results show that our method, despite showing large variance at the start, was able to quickly reduce the variance and converge reliably to a good success rate.

Original: We also provide similar plots of solution variance in
Revised: We also provide similar plots of solution variance in

Original: Fig.
Revised: Fig.

Original: 6, where we report the results for each method using its own convergence time step.
Revised: 6, where we report the results for each method using its own convergence time step.

Original:  C ONCLUSIONS AND F UTURE W ORK
Revised:  C ONCLUSIONS AND F UTURE W ORK

Original: We presented a 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.
Revised: We presented a 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.

Original: Our formulation explicitly accounts for the GMM structure in the optimization and furthermore enables us to naturally constrain the policy updates by the L 2 -Wasserstein distance between GMMs to enhance the stability of the policy optimization process.
Revised: Our formulation explicitly accounts for the GMM structure in the optimization and furthermore enables us to naturally constrain the policy updates by the L 2 -Wasserstein distance between GMMs to enhance the stability of the policy optimization process.

Original: Moreover, the embedding of the Gaussian components of the GMM policy in the Bures-Wassertein manifold greatly reduced the computational cost of the policy optimization.
Revised: Moreover, the embedding of the Gaussian components of the GMM policy in the Bures-Wassertein manifold greatly reduced the computational cost of the policy optimization.

Original: Experiments on three robotic tasks provided strong evidence of the importance of our policy-structure aware optimization against approaches that disregard the GMM structure.
Revised: Experiments on three robotic tasks provided strong evidence of the importance of our policy-structure aware optimization against approaches that disregard the GMM structure.

Original: A possible limitation of our method is that each optimization loop involves running the
Revised: A possible limitation of our method is that each optimization loop involves running the

Original: Sinkhorn algorithm, which is computationally expensive.
Revised: Sinkhorn algorithm, which is computationally expensive.

Original: This might be improved by employing recent advances on initializing the Sinkhorn algorithm (Thornton & Cuturi, 2022).
Revised: This might be improved by employing recent advances on initializing the Sinkhorn algorithm (Thornton & Cuturi, 2022).

Original: 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.
Revised: 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.

Original: In future work we plan to address the latter problem by using separate adaptive learning rates for weights and Gaussian parameters.
Revised: In future work we plan to address the latter problem by using separate adaptive learning rates for weights and Gaussian parameters.

Revised: Another possibility would entail to reformulate the approach as a fully dynamical, particle-based optimization on the Bures-Wasserstein manifold, where both the locations and weights of the particles are updated using Wasserstein Fisher-Rao gradient flows Chizat et al.

Revised: ; Chizat (2019); Liero et al.

Original: Another interestingfuture direction involves to combine our method with an actor-critic formulation and to replace the multi-step cumulative reward by a trained Q -function.
Revised: Finally it would be interesting to combine our method with an actor-critic formulation and to replace the multi-step cumulative reward by a trained Q -function.

Original: A A PPENDIX A. D ETAILS ON G AUSSIAN M IXTURE R EGRESSION (GMR)
Revised: A A PPENDIX A. D ETAILS ON G AUSSIAN M IXTURE R EGRESSION (GMR)

Original: In GMR we start from a GMM in state-action space π ( s , a ) = (cid:80) N i =1 ω 
Revised: In GMR we start from a GMM in state-action space π ( s , a ) = (cid:80) Ni =1 ω i N (cid:0)

Original: i N (cid:0) [ s a ] T ; µ i , Σ i (cid:1) from which a policy, i  e. a probability distribution on the action space, can be obtained by conditioning on the state, as follows 
Revised: [ s a ] T ; µ i , Σ i (cid:1) from which a policy, i [Equation] e. a probability distribution on the action space, can be obtained by conditioning on the state, as follows [Equation]

Original: [Equation]
Revised: [Equation]

Original: The resulting conditional distribution is another GMM on the action sapce, with state dependent parameters, given by:
Revised: The resulting conditional distribution is another GMM on the action sapce, with state dependent parameters, given by:

Original: [Equation]
Revised: [Equation]

Original: with
Revised: with

Revised: [Equation]

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: Note that we have split the GMM parameters µ i and Σ i into their state and action components according to
Revised: Note that we have split the GMM parameters µ i and Σ i into their state and action components according to

Original: [Equation]
Revised: [Equation]

Original: A. R IEMANNIAN GRADIENTS AND RETRACTIONS
Revised: A. R IEMANNIAN GRADIENTS AND RETRACTIONS

Original: For completeness we give here the explicit expressions of the Riemannian gradients and the retractions used in § 3.1.
Revised: For completeness we give here the explicit expressions of the Riemannian gradients and the retractions used in § 3.1.

Original: 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.reduces to the well-known Euclidean gradient descent
Revised: 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.reduces to the well-known Euclidean gradient descent

Original: [Equation]
Revised: [Equation]

Original: (21) where ∇ ˆ µ denotes the Euclidean gradient w.r.t.
Revised: 21) where ∇ ˆ µ denotes the Euclidean gradient w.r.t.

Original: ˆ µ .
Revised: ˆ µ .

Original: For the covariance matrices we use the gradient and retraction w.r.t.
Revised: For the covariance matrices we use the gradient and retraction w.r.t.

Original: the Bures-Wasserstein manifold, taken from (Malag ` o et al., 2018; Han et al., 2021).
Revised: the Bures-Wasserstein manifold, taken from (Malag ` o et al., 2018; Han et al., 2021).

Original: The gradient is given by
Revised: The gradient is given by

Original: [Equation]
Revised: [Equation]

Original: (22) where again ∇ ˆΣ denotes the Euclidean gradient w.r.t.
Revised: 22) where again ∇ ˆΣ denotes the Euclidean gradient w.r.t.

Original: ˆΣ and { X } S =
Revised: ˆΣ and { X } S =

Original: [Equation] X + X [Equation] [Equation]
Revised: ( X + X T ).

Original: Furthermore, the retraction is given by
Revised: Furthermore, the retraction is given by

Original: [Equation]
Revised: [Equation]

Original: 23) where L ˆ X (cid:104)
Revised: 23) where L ˆ X (cid:104)

Original: ˆ Σ k (cid:105) is the Lyapunov operator, defined as the solution to the matrix linear system
Revised: ˆ Σ k (cid:105) is the Lyapunov operator, defined as the solution to the matrix linear system

Original: L ˆ X
Revised: L ˆ X

Original: [Equation]
Revised: [Equation]

Original: A. E XPRESSIONS OF THE FREE FUNCTIONAL J ( π ) AND ITS E UCLIDEAN GRADIENTS
Revised: A. E XPRESSIONS OF THE FREE FUNCTIONAL J ( π ) AND ITS E UCLIDEAN GRADIENTS

Original: For completeness sake, we provide here the explicit expression of the Euclidean gradients for the objective J ( π ) w.r.t.
Revised: For completeness sake, we provide here the explicit expression of the Euclidean gradients for the objective J ( π ) w.r.t.

Original: the parameters of the GMM, which are used in the construction of the Riemannian gradients.
Revised: the parameters of the GMM, which are used in the construction of the Riemannian gradients.

Original: Using the policy gradient theorem, we obtain the gradient of Eq.
Revised: Using the policy gradient theorem, we obtain the gradient of Eq.

Original: 11 w.r.t to a parameter ξ as follows
Revised: 11 w.r.t to a parameter ξ as follows

Original: [Equation]
Revised: [Equation]

Original: In this work
Revised: In this work

Original: [Equation]
Revised: [Equation]

Original: we focus on GMM policies
Revised: we focus on GMM policies

Original: [Equation]
Revised: [Equation]

Original: for which the objective J ( π ) takes the form:
Revised: for which the objective J ( π ) takes the form:

Original: [Equation]
Revised: [Equation]

Original: By inserting Eq.
Revised: By inserting Eq.

Original: 25 into Eq.
Revised: 25 into Eq.

Original: 24 we obtain for the individual parameters of the GMM ∇ µ l 
Revised: 24 we obtain for the individual parameters of the GMM ∇ µ l J

Original: [Equation]
Revised: [Equation]

Original: Here δ s ∈ { 0 , 1 } indicates which terms the gradient acts on.
Revised: Here δ s ∈ { 0 , 1 } indicates which terms the gradient acts on.

Original: In this case, the gradient act on the state components and it is absent for the action dimensions. 
Revised: In this case, the gradient act on the state components and it is absent for the action dimensions. [Equation]

Revised: r ( s t , a t ) (cid:35) .

Original: [Equation]
Revised: [Equation]

Original: Note that we introduced the responsibilities ζ l, s t , a t and ζ l, s t , which are defined as follows
Revised: Note that we introduced the responsibilities ζ l, s t , a t and ζ l, s t , which are defined as follows

Original: [Equation]
Revised: [Equation]

Original: and
Revised: and

Original: [Equation]
Revised: [Equation]

Original: A. R ELATION BETWEEN FORWARD AND BACKWARD DISCRETIZATION IN THE B URES -W ASSERSTEIN METRIC
Revised: A. R ELATION BETWEEN FORWARD AND BACKWARD DISCRETIZATION IN THE B URES -W ASSERSTEIN METRIC

Original: In this section we outline the relation between the implicit and explicit optimization schema w.r.t.
Revised: In this section we outline the relation between the implicit and explicit optimization schema w [Equation]

Original:  the Bures-Wasserstein metric, which is leveraged to formulate our policy optimization in § 3.
Revised: r.t. the Bures-Wasserstein metric, which is leveraged to formulate our policy optimization in § 3.

Original: We closely follow Chen & Li (2020).
Revised: We closely follow Chen & Li (2020).

Original: For the sake of simplicity, we group the Gaussian parameters µ and Σ into a single parameter vector θ .
Revised: For the sake of simplicity, we group the Gaussian parameters µ and Σ into a single parameter vector θ .

Original: Furthermore, we restrict our explanation to a single Gaussian component, which is possible without loosing generality, as each of the N components live in its own manifold R d ×S d ++ .
Revised: Furthermore, we restrict our explanation to a single Gaussian component, which is possible without loosing generality, as each of the N components live in its own manifold R d ×S d ++ .

Original: The Riemannian gradient w.r.t the Gaussian parameters θ , grad θ J ( π ( θ )) , satisfies by definition
Revised: The Riemannian gradient w.r.t the Gaussian parameters θ , grad θ J ( π ( θ )) , satisfies by definition

Original: [Equation]
Revised: [Equation]

Original: (33) where ∇ θ denotes the Euclidean gradient, ξ is an arbitrary vector on the tangent space T θ M , and g θ is the Riemannian metric tensor, defining the inner product on T θ M .
Revised: where ∇ θ denotes the Euclidean gradient, ξ is an arbitrary vector on the tangent space T θ M , and g θ is the Riemannian metric tensor, defining the inner product on T θ M .

Original: The Riemannian metric g θ can be written as
Revised: The Riemannian metric g θ can be written as

Original: [Equation]
Revised: [Equation]

Original: 34) with two arbitrary tangent vectors ζ , ξ , and G W ( θ ) being a positive definite matrix.
Revised: 34) with two arbitrary tangent vectors ζ , ξ , and G W ( θ ) being a positive definite matrix.

Original: Moreover, note that the Wasserstein distance W 22 (cid:0) N ( θ ) , N ( θ + ∆ θ ) (cid:1) , where ∆ θ denotes a small perturbation in the Gaussian parameters θ , can be expressed as 
Revised: Moreover, note that the Wasserstein distance W 22 (cid:0) N ( θ ) , N ( θ + ∆ θ ) (cid:1) , where ∆ θ denotes a small perturbation in the Gaussian parameters θ , can be expressed as [Equation]

Original: [Equation]
Revised: [Equation]

Original: ) for ∆ θ → 0 .
Revised: for ∆ θ → 0 .

Original: Similarly, we can approximate the objective evaluated at J ( θ + ∆ θ ) via the Taylor theorem as
Revised: Similarly, we can approximate the objective evaluated at J ( θ + ∆ θ ) via the Taylor theorem as

Original: [Equation]
Revised: [Equation]

Original: With this, we can approximate
Revised: With this, we can approximate

Original: [Equation]
Revised: [Equation]

Original: from which we obtain the update equation for θ as follows
Revised: from which we obtain the update equation for θ as follows

Original: [Equation]
Revised: [Equation]

Original: Note that Eq.
Revised: Note that Eq.

Original: 39 in turn corresponds to an approximation of the exact Riemannian gradient descent
Revised: 39 in turn corresponds to an approximation of the exact Riemannian gradient descent

Original: [Equation]
Revised: [Equation]

Original: This approximation can be obtained by considering a first-order approximation of the geodesic on the BW manifold.
Revised: This approximation can be obtained by considering a first-order approximation of the geodesic on the BW manifold.

Original: As the exponential map (a.k.a.
Revised: As the exponential map (a.k.a.

Original: the retraction) is defined via the geodesic, the retraction operator in Eq.
Revised: the retraction) is defined via the geodesic, the retraction operator in Eq.

Original: 40 turns into a simple addition operation under a first-order approximation, leading to Eq. 
Revised: 40 turns into a simple addition operation under a first-order approximation, leading to Eq. 39.

Original: Notice that such approximation does not guarantee that the updated parameters θ stay on the manifold, except for the cases in which θ ∈ R d .
Revised: Notice that such approximation does not guarantee that the updated parameters θ stay on the manifold, except for the cases in which θ ∈ R d .

Original: In our case, we leverage the retraction and Riemannian gradients of (Malag`o et al., 2018; Han et al., 2021), which allow us to apply the exact Riemannian gradient descent of 40.
Revised: In our case, we leverage the retraction and Riemannian gradients of (Malag`o et al., 2018; Han et al., 2021), which allow us to apply the exact Riemannian gradient descent of 40.

Original: 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 (cid:0)
Revised: 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 (cid:0)

Original: R d × S d ++ (cid:1) N .
Revised: R d × S d ++ (cid:1) N .

Original:  Algorithm 
Revised: We extended the Pymanopt (Townsend et al., 2016) by adding a custom line-search routine that accounts for a 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.

Original: The  · ) is arbitrary in general.
Revised: The constraint function c ( x 0 . · ) is arbitrary in general.

Original: We use the L 2 -Wasserstein distance between two points on the manifold of GMMs as constraint.
Revised: We use the L 2 -Wasserstein distance between two points on the manifold of GMMs as constraint.

Original: Input :point x 0 on the manifold, descentdirection d , initialstep size λ 0 , decrementα , constraint c ( x 0 , · ), maximum allowed value for constraint c max, minimum step size λ min
Revised: Input : the manifold descent initial decrement constraint c ( x 0 , · ) maximum allowed value for constraint c max minimum step size λ min

Original: Output : stepsize s , updatedpoint on manifold x 1: x = x 0 + λ 0 · d λ = λ2: while c( x 0 , x ) > c maxand λ> λ min do 3: decreasestep size: λ = α· λ updatepoint on manifold: x= x0 +λ · d 4: end while5: if λ < λ min then 6: return λ 0 , x7: else 8: return λ , x 9: end if
Revised: Output : step updated x 1: x = x 0 + λ 0 · d λ = λ2: while c c max λ λ min do 3: decrease λ = α λ update x x + d 4: end while

Original: We extended the Pymanopt (Townsend et al., 2016) by adding a custom line-search routine that accounts for a constraint on the Wasserstein distance between the old and the optimized GMMs. The details of this line-search can be found in Algorithm 2.
Revised: : else 8: x 9: end if

Original: A. A DDITIONAL DETAILS ON EXPERIMENTS A.6. A DDITIONAL RESULTS
Revised: A. A DDITIONAL DETAILS ON EXPERIMENTS A.6. A DDITIONAL RESULTS

Original: Fig.
Revised: Fig.

Original: 5 shows the convergence curves for the two baselines as in Fig.
Revised: 5 shows the convergence curves for the two baselines as in Fig.

Original: 3 of the main paper, however, we extended the horizontal axis up to the maximum number of environment steps used for training.
Revised: 3 of the main paper, however, we extended the horizontal axis up to the maximum number of environment steps used for training.

Original: Fig.
Revised: Fig.

Original: 6 shows the variance of the success rate for the three methods at their time step of convergence for all three robotic tasks.
Revised: 6 shows the variance of the success rate for the three methods at their time step of convergence for all three robotic tasks.

Original: Concerning SAC, which did not converge after the maximum number of environment steps used for training, we chose the last time step.
Revised: Concerning SAC, which did not converge after the maximum number of environment steps used for training, we chose the last time step.

Original: 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) .
Revised: 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) .

Original: These plots show that PPO may also reach low-variance success rate over the five runs at the time step of convergence, at the cost of a prohibitively large number of steps.
Revised: These plots show that PPO may also reach low-variance success rate over the five runs at the time step of conver- gence, at the cost of a prohibitively large number of steps.

Original: SAC showed huge variance in all tasks, apart from the reaching task, where all runs collapsed to a success rate of 0 .
Revised: SAC showed huge variance in all tasks, apart from the reaching task, where all runs collapsed to a success rate of 0 .

Original: [Figure]
Revised: [Figure]

Revised: A.6. A DDITIONAL A BLATIONS

Revised: In order to assess the influence of leveraging a Riemannian optimization approach on the BuresWasserstein manifold, we conducted an ablation of our method by eliminating the Riemannian formulation.

Revised: Instead of the explicit Euler scheme update in Eq.

Revised: 12, which corresponds to Riemannian gradient descent w.r.t.

Revised: the Bures-Wasserstein metric, we use the implicit Euler scheme

Revised: [Equation]

Revised: To guarantee that the updated covariance matrices do not leave the manifold of symmetric positive definite matrices, we parameterize them in terms of Cholesky factors.

Revised: The results obtained with this non-Riemannian version of our method are shown in Fig.

Revised: 7 in direct comparison to our method and

Revised: Fig.

Revised: 8 for an extended range.

Revised: The results clearly show that the non-Riemannian method struggles to reach a success rate offor the reaching task and the collision-avoidance task.

Revised: Furthermore, we observe a high variance over different runs in the same settings (see Fig.

Revised: 9 and Fig. 10).

Revised: We attribute this to the fact that the our method takes exact gradient steps in the direction of steepest descent w.r.t.

Revised: the underlying

Revised: BW metric, whereas the implicit scheme only approximates this direction.

Revised: For this reason the nonRiemannian method is much more noisy, which in turn leads to the aforementioned high variance.

Revised: Nevertheless, the multiple-goal task constitutes an exception.

Revised: Here we observed a similar performance for our approach and the ablated method.

Revised: The reason for this is that the optimization of this task is mainly dominated by the weight updates, which are identical for both methods.

Revised: This result is therefore expected and confirms that correctness of our ablation strategy.

Revised: A.6. A DDITIONAL E XPERIMENT WITH 7 -D O F ROBOTIC MANIPULATOR

Revised: We carried out an additional experiment to show that our method can be employed on tasks performed by off-the-shelf robotic manipulators (e.g. a 7 -DoF Franka Emika Panda robot).

Revised: Specifi-

Revised: [Figure]

Revised: cally, we extended the collision-avoidance task described in § 4 to a 3 D environment (i

Revised: [Figure]

Revised: e. the state s = x ∈ R 3 and the action a = ˙ x ∈ R 3 ).

Revised: The initial 3 -components GMM policy was trained using 10 human demonstrations featuring linear reaching 3 D trajectories.

Revised: 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 at the end of the rollout.

Revised: Moreover, two sparse penalty terms were added to punish collision with obstacles and divergent trajectories

Revised: [Figure]

Revised: Similarly to the planar task reported in the main paper

Revised: [Figure]

Revised: we tested whether our method was able to adapt a trajectory tracking skill in order to avoid collisions with newly added obstacles.

Revised: This means that the robot end-effector needed to pass through a narrow path between two spherical obstacles.

Revised: The robot end-effector pose was controlled using a full-pose Cartesian velocity controller at a frequency of 100Hz , where the end-effector orientation was kept constant.

Revised: Figure 11 shows that our method reached a success rate of 1 . 0 very quickly, taking approximately 20000 environment steps.

Revised: Moreover, the solution variance of our method was 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.

Original: A.6. I NITIAL GMM POLICIES
Revised: A.6. I NITIAL GMM POLICIES

Original: For the sake of completeness, Fig. 7 provides 2 D projections of 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.
Revised: For the sake of completeness, Fig. 12 provides 2 D projections of 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.

Original: Figure 7 also provides the demonstration data used to train the initial policies.
Revised: Figure 12 also provides the demonstration data used to train the initial policies.

Original: Note that these models are then adapted according to the policy optimization approach introduced in § 3.2.
Revised: Note that these models are then adapted according to the policy optimization approach introduced in § 3.2.

Original: [Figure]
Revised: [Figure]

Revised: [Figure]

Revised: [Figure]