SlowGuess/NeurIPS_2023_OCR huggingface dataset:
To begin with, we leverage the approximation that describes the GMM submanifold as a discrete distribution over the space of Gaussian distributions  $\mathcal{G}(\mathbb{R}^d)$ , equipped with the Wasserstein metric [34, 35]. 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 Bures-Wasserstein manifold  $(BW^{N})$ , i.e., the product manifold  $\left(\mathbb{R}^d\times \mathcal{S}_{++}^d\right)^N$ . The former corresponds to the GMM weights while the latter applies to the set of Gaussian distributions' parameters. Note that the identification with the  $BW^{N}$  manifold allows us to perform the optimization directly on the parameter space. This comes with several benefits: (i) We can leverage the well-known analytic solution of the Wasserstein distance between Gaussian distributions in Eq. 6, greatly reducing the computational complexity of the policy optimization. (ii) As shown in [34], we can guarantee that the policy optimized via Eq. 6 remains a GMM (i.e., it satisfies the mass conservation constraint). (iii) Unlike the full Wasserstein space<sup>2</sup>, the resulting product manifold is a true Riemannian manifold such that we can leverage the machinery of Riemannian optimization. Importantly, working in the parameter space allows us to apply an explicit Euler scheme, instead of the implicit formulation of Eq. 3, when optimizing the Gaussian distributions' parameters.
According to the above splitting scheme, we formulate the policy optimization as a two-step procedure that alternates between the Gaussian parameters (i.e. means and covariance matrices) and the GMM weights. To optimize the former, we propose to leverage the Riemannian structure of the  $BW$  manifold to reformulate the updates as a forward discretization, similarly to [58]. 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. 5 can be approximated by an explicit Euler update scheme according to the  $BW$  metric (further details are provided in App. A.4). This allows us to leverage the expressions of the Riemannian gradient and exponential map of the  $BW$  manifold [59, 60]. Thus, the optimization boils down to Riemannian gradient descent where the gradient is defined w.r.t the Bures-Wasserstein metric. In particular, we use the expression for Riemannian gradient, metric and exponential map used in [60]. Formally, the resulting updates for the Gaussian parameters of the GMM follow the Riemannian gradient descent scheme given by:
$$
\hat {\boldsymbol {\mu}} _ {k + 1} = \mathrm {R} _ {\hat {\boldsymbol {\mu}} _ {k}} (\lambda \cdot \operatorname {g r a d} _ {\hat {\boldsymbol {\mu}}} J (\pi_ {k})) ， \quad \text {a n d} \quad \hat {\boldsymbol {\Sigma}} _ {k + 1} = \mathrm {R} _ {\hat {\boldsymbol {\Sigma}} _ {k}} (\lambda \cdot \operatorname {g r a d} _ {\hat {\boldsymbol {\Sigma}}} J (\pi_ {k})) ， \tag {12}
$$
where  $\mathrm{grad}$  denotes the Riemannian gradient w.r.t. the Bures-Wasserstein metric,  $\mathrm{R}_x:\mathcal{T}_x\mathcal{M}\to \mathcal{M}$  denotes the retraction operator, which maps a point on the tangent space  $\mathcal{T}_{\pmb{x}}\mathcal{M}$  back to the manifold  $\mathcal{M}\equiv \mathrm{BW}$  [61]. Moreover,  $\lambda$  is a learning rate and  $\pi_k\stackrel {\mathrm{def}}{=}\pi (\hat{\pmb{\mu}}_k,\hat{\pmb{\Sigma}}_k,\hat{\pmb{\omega}}_k)$ . The Euclidean gradients of  $J(\pi)$  required for computing grad can be obtained using a likelihood ratio estimator (a.k.a score function estimator or REINFORCE) [62] and are provided in App. A.3.
Concerning the GMM weights, we first reparameterize them as  $\omega_{j} = \frac{\exp\eta_{j}}{\sum_{k = 1}^{N}\exp\eta_{k}}$  and optimize w.r.t. the new parameters  $\eta_j,j = 1\dots N$ , which unlike  $\hat{\omega}$  are unconstrained. For this optimization we employ the implicit Euler scheme:
$$
\hat {\eta} _ {k + 1} = \underset {\hat {\eta}} {\arg \min } \left(\frac {W _ {2} ^ {2} \left(\pi_ {k + 1} (\hat {\eta}) , \pi_ {k}\right)}{2 \tau} - J \left(\pi_ {k + 1} (\hat {\eta})\right)\right), \tag {13}
$$
$^{2}$ Wasserstein space is not a true Riemannian manifold, but it can be equipped with a Riemannian structure and formal calculus on this manifold [57], which has been made rigorous in [51]
Algorithm 1 GMM Policy Optimization via Wasserstein Gradient Flows
Input: initial policy  $\pi (\pmb {a}|\pmb {s})$    
1: while not goal reached do   
2: Rollout policy  $\pi (\pmb {a}|\pmb {s})$  in the environment for  $M$  episodes to collect interaction trajectories  $\tau = \{(s_0,a_0,r_0),(s_1,a_1,r_1),\dots ,(s_T,a_T,r_T)\}_{m = 1}^M$    
3: repeat   
4: Update Gaussian components parameters  $\hat{\mu},\hat{\Sigma}$  using Riemannian optimization (12), where  $\lambda^{ls}$  is determined via line-search (see §3.2).   
5: until convergence   
6: repeat   
7: Update GMM weights  $\hat{\omega}$  via gradient descent on the free energy objective 10, using 14   
8: until converged   
9: end while
where  $\pi_{k + 1}(\hat{\pmb{\eta}})\stackrel {\mathrm{def}}{=}\pi (\hat{\pmb{\mu}}_{k + 1},\hat{\pmb{\Sigma}}_{k + 1},\hat{\pmb{\eta}})$ . We minimize Eq. 13 by gradient descent w.r.t.  $\pmb{\eta}$  as follows:
$$
\hat {\eta} _ {k + 1} = \hat {\eta} _ {k} - \lambda \nabla_ {\hat {\eta}} \left(\frac {W _ {2} ^ {2} \left(\pi_ {k + 1} (\boldsymbol {\eta}) , \pi_ {k}\right)}{\tau} - J \left(\pi_ {k + 1} (\hat {\boldsymbol {\eta}})\right)\right). \tag {14}
$$
The gradient of  $J(\pi)$  can be obtained analytically using a likelihood ratio estimator. For the Wasserstein term, we first compute the gradient w.r.t. the weights via the Sinkhorn algorithm [63], from which the gradient w.r.t  $\pmb{\eta}$  can be then obtained via the chain rule. Note that we have to rely on the Sinkhorn algorithm here since there is no analytic solution available for the Wasserstein distance between discrete distributions, unlike the above case of the Gaussian components. Consequently, we cannot compute the corresponding gradients.
3.2 Implementation Pipeline
To carry out the policy optimization, we proceed as in the usual on-policy RL scheme: We first roll out the current policy to collect samples of state-action-reward tuples. Then, we use the collected interaction trajectories to compute a sample-based estimate of the functional  $K_{r}(\pi)$  and its gradients w.r.t the policy parameters, as explained in § 3.1. An optimization step consists of alternating between optimizing the Gaussian parameters using Eq. 12, and updating the weights via Eq. 14. For the optimization of the Gaussian parameters we leverage Pymanopt [64] for Riemannian optimization. We extended this library by implementing the Bures-Wasserstein manifold based on the expressions provided by Han et al. [60] (see App. A.2 for details). Furthermore, we added a custom line-search routine that accounts for the 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. The details of this custom line-search are given in Algorithm 2 in App. A.5. Regarding the optimization of the GMM weights, we use POT [65], a Python library for optimal transport, from which we obtain the quantities required to compute the gradients of the Wasserstein distance w.r.t. the weights in Eq. 14.
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. Afterwards, fresh rollouts are performed with the updated policy and the aforementioned two-step alternating procedure starts over. This optimization loop is repeated until a task-specific success criterion has been fulfilled. We summarize the proposed optimization in Algorithm 1.
We tested our approach in three different robotic settings: a reaching skill, a collision-free trajectory tracking, and a multiple-goal task. All these tasks were represented in a 2D operational space. Moreover, we tested our approach on two additional versions of the collision-free trajectory tracking task that are formulated for 3D operational-space and 7D joint-space representations of the end-effector state and action. These two latter tasks were performed by a 7-DoF Franka Emika Panda robot in a virtual environment as reported in App. A.6.3. All robot motion policies were initially learned from human demonstrations collected on a Python GUI. We assumed we were given  $M$  demonstrations, each of which contained  $T_{m}$  data points for a dataset of  $T = \sum_{m} T_{m}$  total observations
$\tau = \{(s_t, a_t)\}_{t=1}^T$ . The state  $\pmb{s}$  and action  $\pmb{a}$  corresponded to the robot end-effector position  $\pmb{x} \in \mathbb{R}^d$  and velocity  $\dot{\pmb{x}} \in \mathbb{R}^d$ , with  $d$  being the operational space dimensionality. The GMM models were initially trained via classical Expectation-Maximization. The policy rollout consisted of sampling a velocity action  $\pmb{a}_t \sim \pi(\pmb{a}_t | s_t)$  using Eq. 9, and subsequently commanding the robot via a Cartesian velocity controller at a frequency of  $100\mathrm{Hz}$ . For all the experiments, we used the Robotics Toolbox for Python [66] to simulate the robotic environments.
To show the importance of accounting for the policy structure in RL settings, we compared our method against two structure-unaware baselines: PPO [45] and SAC-GMM [31]. Our baselines choice is motivated by the fact that these two RL methods are still widely used and fairly competitive [67], even in real robotic scenarios [68]. We also considered an additional baseline: PMOE [41], which is a method to train deep RL policies using a probabilistic mixture of experts via GMMs. 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.  $\boldsymbol{a} = [\Delta \omega \Delta \hat{\mu} \Delta \mathrm{vec}(\hat{\Sigma})]$ , following the same methodology as SAC-GMM [31]. The PPO and SAC implementations correspond to the code provided by Stable-Baselines3 [69], whose policies are parametrized by MLP networks. The PMOE implementation corresponds to the code given in the original paper [41]. During policy optimization, we sampled an action from the MLP policy that was then used to update the GMM parameters by adding the computed corrections to the current parameters. Later, we proceeded as described earlier, namely, the updated GMM was used to compute the velocity action via Eq. 9. For comparison purposes, we report statistical results for all the settings over 5 runs for the task success rate and solution variance. We tuned the baselines separately for each task using Optuna [70]. In addition, to assess the importance of our Riemannian formulation, we performed an ablation where we used the implicit scheme based on Euclidean gradient descent instead of the explicit optimization on the Bures-Wasserstein manifold (see App. A.6.2).
4.1 Tasks description
Reaching Task: This experiment consists of: (1) learning an initial GMM policy such that the robot end-effector reaches a target by following an L-shape trajectory from its initial position, and (2) adapting the GMM policy to reach a new target located midway and above the previously-learned L-shape trajectories. The initial policy, shown in Fig. 2-left and Fig. 13a, was learned from 12 demonstrations and encoded by a 7-component GMM. To adapt the policy, we defined a dense reward as a function of the position error between the robot end-effector and the new target. We also added a sparse penalty term that punishes rollouts leading to significantly divergent trajectories. Each optimization episode comprised 10 rollouts, each of maximum horizon length of 200 iterations. Convergence is achieved when a minimum average position error w.r.t the target – computed over an episode – is reached.
Collision-avoidance Task: 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. The initial GMM policy was learned from 10 human demonstrations and represented by a 3-component GMM, as shown in Fig. 2-middle and Fig. 13b. For policy optimization, we defined a sparse reward as a function of the position error between the robot end-effector position and the target at the end of the rollout. We also included two sparse penalty terms: the first one punishes rollouts leading to collisions with the obstacles, for which the rollout is stopped; the second term penalizes rollouts with significantly divergent trajectories. Each episode consisted of 10 rollouts, each of maximum horizon length of 150 iterations. Convergence is determined by a minimum average position error w.r.t the target computed over an episode.
Multiple-goal Task: This setting involves: (1) learning an initial GMM policy where the robot end-effector reaches two different targets (i.e., task goals) starting from the same initial position, and (2) adapting the initial policy to reach a new target located close to one of the previous task goals. The intended adapted behavior should make the robot go through the most relevant GMM components according to the new target location. The initial GMM policy was learned from 12 demonstrations and encoded by a 6-component GMM, as shown in Fig. 2-right and Fig. 13c. 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. Similar to the previous experiments, we added a sparse penalty term to penalize rollouts generating significantly divergent trajectories. An
Figure 2: The three tested 2D robotic settings: a reaching skill (left), a collision-free trajectory tracking (middle), and a multiple-goal task (right). The robot color goes from light gray to black to show the evolution of the task reproduction. Green Gaussian components (○) depict the initial GMM policy, projected on the 2D Cartesian position space. The end-effector trajectory resulting from the initial GMM policy is shown in dark blue lines (—). Red circles (●) in the collision-avoidance task represent the obstacles (middle). The different targets of the multiple-goal task (right) are depicted as red stars.
episode comprised 10 rollouts, each of maximum horizon length of 200 iterations. Again, the policy optimization converges when the average position error w.r.t the chosen target reaches a minimum threshold.
4.2 Results Analysis
The reaching task tested our method's ability to adapt a previously-learned reaching skill to a new goal, located at  $(6.0, -6.5)$  (cf. Fig. 13a-left). Achieving this required to adapt the Gaussian parameters of mainly the last four GMM components, while the others remained unchanged. We compared all methods in terms of the success rate over environment steps, where the success rate was defined as the percentage of rollouts reaching the new goal. Figure 3-left shows that our method achieved a success rate of 1 after approximately 70000 environment interactions. Despite PPO was also able to complete the task reliably, it required many more environment steps (cf. Fig. 5-left). In sharp contrast, SAC did not reach any improvement. These observations underline the importance of some kind of trust region or constraint on the policy updates, which allowed both our method and PPO to reach good success rates. The PMOE versions of PPO and SAC showed a better performance than the vanilla counterparts, which aligns with the experimental insights provided in the original PMOE paper [41] regarding the advantages of the probabilistic mixture of experts in deep RL policies. Nevertheless, our approach also outperformed both PMOE versions. This experiment showed that our method is much more sample-efficient in adapting the GMM parameters, which we attribute to the fact that our method explicitly takes the GMM structure into account in the formulation of the optimization.
In the collision-avoidance task, we tested whether our method was able to adapt a trajectory tracking skill in order to avoid collisions with newly added obstacles. 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. Fig. 2-middle). While the reaching task could be adapted by mainly varying the means of the GMM components, this task also demands to adapt the covariance of the second GMM component. Figure 3-middle shows that our method solved this task reliably after comparatively few environment interactions. Although PPO also achieved a success rate of 1, it took 6 times more environment steps than our method. SAC only reached an average success rate of 0.8, however with high variance (cf. Fig. 5-middle). Note that the PMOE versions showed a reduced solution variance and a higher success rate than their vanilla counterparts, where PMOE-PPO also outperformed PMOE-SAC. These results again show the importance of the constraints on the policy updates. The huge discrepancy in the required environment steps between our method and all the baselines further emphasizes the importance of taking the GMM structure into account in the policy optimization. The explicit consideration of this structure allows our method to identify the required changes to the GMM parameters much more efficiently. These experimental insights can also be observed in the results corresponding to the other formulations of the collision-avoidance task, namely, the 3D operational and 7D joint space representations, reported in App. A.6.3.
While the previous two tasks were accomplished by adapting mostly the Gaussian parameters of the GMM, the multiple-goal task required to adapt the GMM weights. The initial skill comprised
Figure 3: Success rate of our method (WGF) and the baselines on the reaching (left), the collision-avoidance (middle) and the multiple-goal tasks (right). The shaded area depicts the standard deviation over 5 runs.
Figure 4: Variance of the success rate over the 5 runs for our method (WGF) and the baselines on the reaching task (left), the collision avoidance task (middle) and the multiple-goal task (right). The violine plots are overlaid with box plots, quartile lines and a swarm plot, where dots indicate the success rates of individual runs. The plots show the variance at the following time steps from left to right: 80000, 90000, 95000.
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. 13c). 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. As shown in Fig. 3-right, our method again quickly achieved a success rate of 1. PPO required substantially many more environment steps, while SAC was not able to solve the task. Again, the PMOE formulation reduced the solution variance and slightly increased the success rate of both PPO and SAC.
In Fig. 4 we report the success rate variance over 5 runs at a fixed time step, which corresponded to the step at which the first method achieved a success rate of 1, thus prioritizing sample efficiency. The plots show that our method exhibits a very low solution variance. All baselines varied largely, except for the reaching task, where all SAC runs collapsed to a success rate of 0. These results show that our method, despite showing large variance at the start, was able to quickly reduce the variance and converge reliably to a good success rate. We also provide similar plots of solution variance in Fig. 6, where we report the results for both PPO and SAC using their own convergence time step.
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. 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. Moreover, the embedding of the Gaussian components of the GMM policy in the Bures-Wasserstein manifold greatly reduced the computational cost of the policy optimization. Experiments on several robotic tasks provided strong evidence of the importance of our policy-structure aware optimization against approaches that disregard the GMM structure. A possible limitation of our method is that each optimization loop involves running the Sinkhorn algorithm, which is computationally expensive. This might be improved by employing recent advances on initializing the Sinkhorn algorithm [71]. Also, we observed an intricate interplay between the optimization of the GMM weights and the Gaussian parameters, which occasionally resulted in one update hampering the other. In future work we plan to address the latter problem by using separate adaptive learning rates for weights and Gaussian parameters. Another possibility would entail to reformulate the policy optimization as a gradient flow on the space of probability measures endowed with the Wasserstein Fisher-Rao metric [72, 73, 74]. This alternative metric may allow us to leverage the Fisher-Rao geometry to optimize the GMM weights, as very recently proposed in [75] to learn isotropic Gaussian mixtures via particle-based approximations. Finally 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.
[69] Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dormann. Stable-baselines3: Reliable reinforcement learning implementations. Journal of Machine Learning Research, 22(268):1-8, 2021. URL http://jmlr.org/papers/v22/20-1364.html. 8
[70] Takuya Akiba, Shotaro Sano, Toshihiko Yanase, Takeru Ohta, and Masanori Koyama. Optuna: A next-generation hyperparameter optimization framework. In ACM SIGKDD International Conference on Knowledge Discovery and Data Mining (KDD), page 2623-2631, 2019. URL https://doi.org/10.1145/3292500.3330701.8
[71] James Thornton and Marco Cuturi. Rethinking initialization of the Sinkhorn algorithm, 2022. URL https://arxiv.org/abs/2206.07630.arXiv:2206.07630.10
[72] Lenaic Chizat, Bernhard Schmitzer, Gabriel Peyré, and François-Xavier Vialard. An interpolating distance between optimal transport and Fisher-Rao, 2015. URL https://arxiv.org/abs/1506.06430. 10
[73] Thomas O. Gallouët and Léonard Monsaingeon. A JKO splitting scheme for Kantorovich-Fisher-Rao gradient flows. SIAM Journal on Mathematical Analysis, 49(2):1100-1130, 2017. URL https://doi.org/10.1137/16M106666X.10
[74] Matthias Liero, Alexander Mielke, and Giuseppe Savare. Optimal entropy-transport problems and a new Hellinger-Kantorovich distance between positive measures. Inventiones mathematicae, 211, 03 2018. URL https://doi.org/10.1007/s00222-017-0759-8.10
[75] Yuling Yan, Kaizheng Wang, and Philippe Rigollet. Learning Gaussian mixtures using the Wasserstein-Fisher-Rao gradient flow, 2023. URL https://arxiv.org/abs/2301.01766.10
A.1 Details on Gaussian Mixture Regression (GMR)
In GMR we start from a GMM in state-action space  $\pi (\pmb {s},\pmb {a}) = \sum_{i = 1}^{N}\omega_{i}\mathcal{N}\big([s\pmb {a}]^{\top};\pmb {\mu}_{i},\pmb{\Sigma}_{i}\big)$  from which a policy, i.e. a probability distribution on the action space, can be obtained by conditioning on the state, as follows
$$
\pi (\boldsymbol {a} | \boldsymbol {s}) = \frac {\pi (\boldsymbol {s} , \boldsymbol {a})}{\int \pi (\boldsymbol {s} , \boldsymbol {a}) \mathrm {d} \boldsymbol {a}}. \tag {15}
$$
The resulting conditional distribution is another GMM on the action sapce, with state dependent parameters, given by: