SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Hanna Ziesche* and Leonel Rozo*
Bosch Center for Artificial Intelligence (BCAI)
Renningen, Germany
name.surname@de.bosch.com
Robots often rely on a repertoire of previously-learned motion policies for performing tasks of diverse complexities. When facing unseen task conditions or when new task requirements arise, robots must adapt their motion policies accordingly. In this context, policy optimization is the de facto paradigm to adapt robot policies as a function of task-specific objectives. Most commonly-used motion policies carry particular structures that are often overlooked in policy optimization algorithms. We instead propose to leverage the structure of probabilistic policies by casting the policy optimization as an optimal transport problem. Specifically, we focus on robot motion policies that build on Gaussian mixture models (GMMs) and formulate the policy optimization as a Wasserstein gradient flow over the GMMs space. This naturally allows us to constrain the policy updates via the  $L^2$ -Wasserstein distance between GMMs to enhance the stability of the policy optimization process. Furthermore, we leverage the geometry of the Bures-Wasserstein manifold to optimize the Gaussian distributions of the GMM policy via Riemannian optimization. We evaluate our approach on common robotic settings: Reaching motions, collision-avoidance behaviors, and multi-goal tasks. Our results show that our method outperforms common policy optimization baselines in terms of task success rate and low-variance solutions.
One of the main premises about autonomous robots is their ability to successfully perform a large range of tasks in unstructured environments. This demands robots to adapt their task models according to environment changes or new task requirements, and consequently to adjust their actions to successfully perform under unseen conditions [1]. In general, robotic tasks, such as picking or inserting an object, are usually executed by composing previously-learned skills [2], each represented by a motion policy. Therefore, in order to successfully perform under new settings, the robot should adapt its motion policies according to the new task requirements and environment conditions.
Research on methods for robot motion policy adaptation is vast [3, 4], with approaches mainly building on black-box optimizers [5, 6], end-to-end deep reinforcement learning [7, 8], and policy search [9]. Regardless of the optimization method, most approaches disregard the intrinsic policy structure in the adaptation strategy. However, several motion policy models (e.g., dynamic movement primitives (DMP)[10], Gaussian mixture models (GMM) [11], probabilistic movement primitives (ProMPs) [12, 13], and neural networks [14], among others), carry specific physical or probabilistic structures that should not be ignored. First, these policy models are often learned from demonstrations in a starting learning phase [2], thus the policy structure already encapsulates relevant prior information about the skill. Second, structure-unaware adaptation strategies optimize the policy parameters disregarding
Figure 1: Illustration our policy structure-aware adaptation of GMM policies. Left: A robot manipulator tracks a reference trajectory (—) extracted from a GMM policy (○) initially learned from demonstrations, while being required to reach for a new target (★). Center: Our policy optimization provides policy updates that follow a Wasserstein gradient flow on the manifold of GMM policies  $\mathrm{GMM}_d$ . Right: The final GMM policy where some Gaussian components were adapted in order to retrieve a new reference trajectory that allows the robot to reach for the new target.
the intrinsic structural characteristics of the policy model (e.g., a DMP represents a second-order dynamical system). In this context, we hypothesize that the policy structure may be leveraged to better control the adaptation strategy via policy structure-aware gradients and trust regions.
Our main idea is to design a policy optimization strategy that explicitly builds on a particular policy structure. Specifically, we focus on GMM policies, which have been widely used to learn motion skills from human demonstrations [11, 15, 16, 17, 18, 19]. GMMs provide a simple but expressive enough representation for learning a large variety of robot skills such as: Stable dynamic motions [20, 21, 22], collaborative behaviors [23, 24], assembly strategies [25, 26], and contact-rich manipulation [27, 28], among others. Often, skills learned from demonstrations need to be refined — due to imperfect data — or adapted to comply with new task requirements. Existing adaptation strategies for GMM policies either build a kernel method on top of the original skill model [29], or leverage reinforcement learning (RL) to adapt the policy itself [30, 31]. However, none of these techniques explicitly considered the intrinsic probabilistic structure of the Gaussian mixture policy.
Unlike the aforementioned approaches, we propose a policy optimization technique that explicitly considers the underlying GMM structure. To do so, we exploit optimal transport theory [32, 33], which allows us to view the set of GMM policies as a particular space of probability distributions  $\mathrm{GMM}_d$ . Specifically, we leverage the idea of Chen et al. [34] and Delon and Desolneux [35] to view a GMM as a set of discrete measures (Dirac masses) on the space of Gaussian distributions  $\mathcal{G}(\mathbb{R}^d)$ , which is endowed with a Wasserstein distance (see § 2). This allows us to formulate the policy optimization as a Wasserstein gradient flow (WGF) over the space of GMMs (as illustrated in Fig. 1 and explained in §3), where the policy updates are naturally guaranteed to be GMMs. Moreover, we take advantage of the geometry of the Bures-Wasserstein manifold to optimize the Gaussian distributions of a GMM policy via Riemannian optimization. We evaluate our approach over a set of different GMM policies featuring common robot skills: Reaching motions, collision-avoidance behaviors, and multi-goal tasks (see § 4). Our results show that our method outperforms common policy optimization baselines in terms of task success rate while providing low-variance solutions.
Related Work: Richemond and Maginnis [36] pioneered the idea of understanding policy optimization through the lens of optimal transport. 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. They observed that the resulting policy optimization resembles the gradient flow of the Fokker-Planck equation (JKO scheme) [37]. In a similar spirit, Zhang et al. [38] 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. 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. Although Zhang et al. [38] 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. Recently, Mokrov et al. [39] tackled the computational burden of particle methods by leveraging input-convex neural networks to approximate the WGFs computation. They reformulated the well-known JKO optimization [37] over probability measures by an optimization over convex functions. Yet, this work remains a general solution for WGF computation and it did not address its use for policy optimization problems.
Aside from optimal transport approaches, Arenz et al. [30] proposed a trust-region variational inference for GMMs to approximate multimodal distributions. Although not originally designed
for policy optimization, the authors elucidated a connection to learn GMMs of policy parameters in black-box RL. However, their method cannot directly be applied to our GMM policy adaptation setting, nor does it consider the GMM structure from an optimal transport perspective. Nematollahi et al. [31] proposed SAC-GMM, a hybrid model that employs the well-known SAC algorithm [40] to refine dynamic skills encoded by GMMs. The SAC policy was designed to learn residuals on a single vectorized stack of GMM parameters, thus fully disregarding the GMM structure and the geometric constraints of its parameters. Ren et al. [41] introduced PMOE, a method to train deep RL policies using a probabilistic mixture of experts via GMMs. They addressed the non-differentiability problem caused by the optimization of categorical distribution parameters associated with the mixture weights via gradient estimators. However, PMOE does not take an optimal transport perspective on the policy optimization problem, nor does it account for the geometry arising from the GMM parameters.
Concerning approaches that explicitly account for the geometry induced by the policy, we may broadly categorize them into two groups: exact models and natural gradient methods. The former category refers to policy optimization methods that fully exploit the geometric structure of the policy, as proposed in this paper. In this case, the state of the art is scarce but recent works such as the Riemannian proximal policy optimization for GMMs proposed in [42], showed the importance of leveraging the geometry induced by the GMM parameters in the policy optimization via Riemannian gradients, similarly to our method. Their policy optimization was regularized by a Wasserstein distance to control the exploration-exploitation trade-off. However, their method did not formulate the policy optimization as an optimal transport problem, i.e., the policy updates did not follow a Wasserstein gradient flow, as in our approach, but it employed instead a classical non-convex optimization procedure.
The second category encloses approaches that leverage the well-established technique for enhancing the stability of policy optimization, namely, introducing regularization to the objective function using the Kullback-Leibler (KL) divergence, which induces the so-called natural gradient [43, 44, 45, 46]. In this context, the regularization provides a global measure of the policies similarity, while the natural gradient perspective aims at locally controlling the changes between successive policy updates. In other words, natural gradient approaches leverage the local geometry of the policy parameters space to compute controlled updates. Note that policies similarity can also be quantified via the Wasserstein distance, as proposed in this paper and in recent works [36, 38, 47, 48]. Unlike the KL divergence, the Wasserstein distance enjoys powerful geometrical, computational, and statistical features [32, 33]. For example, Moskovitz et al. [48] recently employed the Wasserstein natural gradient to exploit the local geometry induced by the Wasserstein regularization of behavioral policy optimization [47], but neither the policy nor the behavioral embeddings had particular structures that could be leveraged in the policy optimization.
In this paper we exploit the geometric properties arising from both a Wasserstein-regularized objective and a Gaussian mixture policy. This allows us not only to see the policy iteration as a Wasserstein gradient flow, as in [36], but also to leverage the Riemannian geometry associated to the GMM space  $\mathrm{GMM}_d$  to formulate exact Riemannian gradient descent updates [49], instead of relying on inexact natural gradients approximations. To achieve this, our method leverages the geometry induced by the structure of the space of GMM policies via the Bures-Wasserstein manifold, which naturally guarantees that policy updates stay on  $\mathrm{GMM}_d$ . To the best of our knowledge, our optimization of GMM policies based on Wasserstein gradient flows and the Bures-Wasserstein manifold is the first of its kind in the domain of policy adaptation.
2.1 Wasserstein gradient flows
In Euclidean space a gradient flow is a smooth curve  $x: \mathbb{R} \to \mathbb{R}^d$  that satisfies the partial differential equation (PDE)  $\dot{x}(t) = -\nabla L(x(t))$  for a given loss function  $L: \mathbb{R}^d \to \mathbb{R}$  and starting point  $x_0$  at  $t = 0$  [32, 50]. A solution can be found straightforwardly by forward discretization, leading to the well-known explicit Euler update scheme  $x_{k+1}^{\tau} = x_k^{\tau} - \lambda \nabla L(x_k^{\tau})$ , where  $\lambda$  denotes the learning rate and  $x^{\tau}$  indicates a discretization of the curve  $x(t)$  with discretization parameter  $\tau$ . Alternatively, we can use a backward discretization, which leads to the following implicit Euler scheme
$$
x _ {k + 1} ^ {\tau} = \arg \min  _ {x} \left(\frac {\| x - x _ {k} ^ {\tau} \| ^ {2}}{2 \tau} + L (x)\right). \tag {1}
$$
Eq. 1 is sometimes referred to as Minimizing Movement Scheme and can be used as an alternative characterization of a gradient flow.
This characterization is particularly interesting when we need to extend the concept of gradient flows to (non-Euclidean) general metric settings, since there is no notion of  $\nabla L$  in these cases [32, 51]. Eq. 1 does not involve any gradients and can be expressed using only metric quantities. In this work, we are particularly interested in gradient flows in the  $L^2$ -Wasserstein space, defined as the set of probability measures  $\mathbb{P}(\mathcal{X})$  on a separable Banach space  $\mathcal{X}$  [52] and endowed with the  $L^2$ -Wasserstein distance  $W_2$  defined as
$$
W _ {2} (\mu , \nu) = \left(\inf  _ {\gamma \in \Pi (\mu , \nu)} \int_ {\mathcal {X} \times \mathcal {X}} \| x _ {1} - x _ {2} \| ^ {2} \mathrm {d} \gamma (x _ {1}, x _ {2})\right) ^ {\frac {1}{2}}, \tag {2}
$$
where  $\mu, \nu \in \mathbb{P}(\mathcal{X})$  and  $\gamma \in \mathbb{P}(\mathcal{X}^2)$  is defined to have the two marginals  $\mu$  and  $\nu$ .
A Generalized Minimizing Movement scheme characterizing gradient flows in the Wasserstein space can be written in analogy to Eq. 1 as:
$$
\pi_ {k + 1} ^ {\tau} = \arg \min  _ {\pi} \left(\frac {W _ {2} ^ {2} \left(\pi , \pi_ {k} ^ {\tau}\right)}{2 \tau} + L (\pi)\right), \tag {3}
$$
where  $L$  is a functional to be minimized on the Wasserstein space and  $\pi_k\in \mathbb{P}(X)$ . In the following, we will omit the superscript  $\tau$  for notational convenience.
2.2 Reinforcement Learning as Wasserstein Gradient Flows
Our view of the policy structure-aware optimization builds on the approach outlined in [36], which in turn is based on the JKO scheme [37]. They proposed a formulation of 1-step RL problems in terms of Wasserstein gradient flows. In particular, they studied the evolution of a policy  $\pi$  under the influence of a free energy functional  $J$  of the form:
$$
J (\pi) = K _ {r} (\pi) + \beta \mathcal {H} (\pi) = \int_ {\mathcal {A}} \mathrm {d} \pi (\boldsymbol {a} | \boldsymbol {s}) r (\boldsymbol {s}, \boldsymbol {a}) - \beta \int_ {\mathcal {A}} \mathrm {d} \pi (\boldsymbol {a} | \boldsymbol {s}) \log (\pi (\boldsymbol {a} | \boldsymbol {s})) \tag {4}
$$
where  $K_{r}(\pi)$  denotes the inner energy of the system, here determined by the reward  $r(s, a)$ , and  $\mathcal{H}(\pi)$  is the entropy of the policy  $\pi(a|s)$ , with  $s$  and  $a$  denoting the state and action, respectively. Thus, Eq. 4 can be recognized as the usual objective in 1-step RL settings with entropy regularization. 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. Richemond and Maginnis [36] exploited the result of Jordan et al. [37], which stated that this evolution can be interpreted as the gradient flow of the functional  $J$  in Wasserstein space. This flow is characterized by the following minimizing movement scheme
$$
\pi_ {k + 1} = \arg \min  _ {\pi} \left(\frac {W _ {2} ^ {2} (\pi , \pi_ {k})}{2 \tau} - J (\pi)\right), \tag {5}
$$
which naturally provides iterative updates for the policy  $\pi$ . While Richemond and Maginnis [36] considered a 1-step bandit setting, we extend this approach to full multi-step RL problems and consequently learn policies for long-horizon tasks.
2.3 The  $L^2$ -Wasserstein distance between Gaussian Mixture Models (GMMs)
We consider policies  $\pi(\boldsymbol{x})$  that build on a GMM structure, i.e.,  $\pi(\boldsymbol{x}) = \sum_{i=1}^{N} \omega_i \mathcal{N}(\boldsymbol{x}; \boldsymbol{\mu}_i, \boldsymbol{\Sigma}_i)$ , where  $\mathcal{N}$  denotes a multivariate Gaussian distribution with mean  $\boldsymbol{\mu}_i$  and covariance matrix  $\boldsymbol{\Sigma}_i$ , and  $\omega_i$  are the weights of the  $N$  individual Gaussian components, which are subject to  $\sum_i \omega_i = 1$ . In the following, we will write  $\hat{\boldsymbol{\mu}}, \hat{\boldsymbol{\Sigma}}$  and  $\hat{\boldsymbol{\omega}}$  to denote the stacked means, covariance matrices and weights of the  $N$  components. Therefore, we do not consider WGFs on the full manifold of probability distributions (Wasserstein space)  $\mathbb{P}(\mathbb{R}^d)$  but rather focus on WGFs evolving on the submanifold of GMMs, that is  $\mathrm{GMM}_d \subset \mathbb{P}(\mathbb{R}^d)$ . Following [34, 35], we can approximately describe this submanifold as a discrete distribution over the space of Gaussian distributions equipped with the Wasserstein metric. This in turn can be identified with the Bures-Wasserstein manifold which is the product manifold  $\mathbb{R}^d \times S_{++}^d$ ,
where  $S_{++}^{d}$  denotes the Riemannian manifold of  $d$ -dimensional symmetric positive definite matrices. The corresponding approximated Wasserstein distance between two GMMs  $\pi_1, \pi_2$  is given by
$$
W _ {2} ^ {2} \left(\pi_ {1} (\boldsymbol {x}), \pi_ {2} (\boldsymbol {x})\right) = \min  _ {\boldsymbol {P} \in U \left(\boldsymbol {\omega} _ {1}, \boldsymbol {\omega} _ {2}\right)} \sum_ {i, j} ^ {N} \boldsymbol {P} _ {i j} W _ {2} ^ {2} \left(\mathcal {N} _ {1} (\boldsymbol {x}; \boldsymbol {\mu} _ {i}, \boldsymbol {\Sigma} _ {i}), \mathcal {N} _ {2} (\boldsymbol {x}; \boldsymbol {\mu} _ {j}, \boldsymbol {\Sigma} _ {j})\right), \tag {6}
$$
where  $U(\omega_1, \omega_2) = \{P \in \mathbb{R}_+^{N \times N} | P\mathbf{1}_N = \omega_1, P^{\mathrm{T}}\mathbf{1}_N = \omega_2\}$  with  $\mathbf{1}_N$  denoting an  $N$ -dimensional vector of ones. The Wasserstein distance between two Gaussian distributions in Eq. 6 can be computed analytically as follows
$$
\left. \right. W _ {2} ^ {2} \left(\mathcal {N} _ {1} (\boldsymbol {x}; \boldsymbol {\mu} _ {i}, \boldsymbol {\Sigma} _ {i}), \mathcal {N} _ {2} (\boldsymbol {x}; \boldsymbol {\mu} _ {j}, \boldsymbol {\Sigma} _ {j})\right) = \| \boldsymbol {\mu} _ {i} - \boldsymbol {\mu} _ {j} \| ^ {2} + \operatorname {t r} \left[ \boldsymbol {\Sigma} _ {i} + \boldsymbol {\Sigma} _ {j} - 2 \left(\boldsymbol {\Sigma} _ {i} ^ {1 / 2} \boldsymbol {\Sigma} _ {j} \boldsymbol {\Sigma} _ {i} ^ {1 / 2}\right) ^ {1 / 2} \right]. \tag {7}
$$
2.4 Learning GMM policies from demonstrations
A popular approach in RL — particularly in robotics — 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. In this work we choose to represent our policy via a GMM. We assume that demonstrations are provided as a set of trajectories  $\tau$  of state-action pairs  $\tau = \{(s_0, a_0), (s_1, a_1), \ldots, (s_T, a_T)\}$ . To initialize our policy, we first use the Expectation-Maximization (EM) algorithm to fit a GMM, in the state-action space, to the demonstrations. This results in a mixture distribution  $\pi(s, a) = \sum_{i=1}^{N} \omega_i \mathcal{N}([s, a]^{\top}; \pmb{\mu}_i, \pmb{\Sigma}_i)$  from which a policy can be obtained by conditioning on the state  $s$ , as follows
$$
\pi (\boldsymbol {a} | \boldsymbol {s}) = \frac {\pi (\boldsymbol {s} , \boldsymbol {a})}{\int \pi (\boldsymbol {s} , \boldsymbol {a}) \mathrm {d} \boldsymbol {a}}. \tag {8}
$$
In the context of GMMs, this is also known as Gaussian Mixture Regression (GMR) [53]. The resulting conditional distribution is another GMM on action space with state-dependent parameters,
$$
\pi \left(\boldsymbol {a} _ {t} \mid \boldsymbol {s} _ {t}\right) = \sum_ {i = 1} ^ {N} \omega_ {i} \left(\boldsymbol {s} _ {t}\right) \mathcal {N} \left(\boldsymbol {a} _ {t}; \boldsymbol {\mu} _ {i} ^ {a} \left(\boldsymbol {s} _ {t}\right), \boldsymbol {\Sigma} _ {i} ^ {a}\right). \tag {9}
$$
Details on computation of Eq. 9 from the original GMM are given in App. A.1.
3 Wasserstein Gradient Flows for GMM Policy Optimization
In this work, we focus on multi-step RL tasks for policy adaptation. We consider a finite-horizon Markov Decision Process (MDP) with continuous state and action spaces  $S \in \mathbb{R}^n$  and  $A \in \mathbb{R}^m$ , transition and reward functions  $p(\boldsymbol{s}_{t+1} | \boldsymbol{s}_t, \boldsymbol{a}_t)$  and  $r(\boldsymbol{s}_t, \boldsymbol{a}_t)$ , initial state distribution  $\rho(\boldsymbol{s}_0)$  and a discount factor  $\gamma$ . Further, we assume to have an initial policy  $\pi(\boldsymbol{a}_t | \boldsymbol{s}_t)$ , which is to be adapted by optimizing some objective function  $K_r(\pi)$ . 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. To promote exploration and avoid premature convergence to suboptimal policies, we leverage maximum entropy RL [54] by adding an entropy term  $\mathcal{H}(\pi)$  to the objective. Thus, the overall objective has the form of a free energy functional (resembling Eq. 4) and can be written as
$$
J (\pi) = K _ {r} (\pi) + \beta \mathcal {H} (\pi), \tag {10}
$$
where  $\beta$  is a hyperparameter and  $K_{r}(\pi)$  corresponds to the usual cumulative return
$$
K _ {r} (\pi) = \mathbb {E} _ {\tau} \left[ \sum_ {t} r \left(\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}\right) \right] = \int \Pi_ {t} \mathrm {d} \boldsymbol {s} _ {0} \mathrm {d} \boldsymbol {s} _ {t} \mathrm {d} \boldsymbol {a} _ {t} \rho \left(\boldsymbol {s} _ {0}\right) \pi \left(\boldsymbol {a} _ {t} \mid \boldsymbol {s} _ {t}\right) p \left(\boldsymbol {s} _ {t + 1} \mid \boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}\right) \sum_ {t} \gamma^ {t} r \left(\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}\right). \tag {11}
$$
The evolution of the policy  $\pi (\pmb{a}_t|\pmb{s}_t)$  over the course of the optimization can be described as a flow of a probability distribution in Wasserstein space. This formulation comes with three major benefits: (i) We directly leverage the Wasserstein metric properties for describing the evolution of probability distributions; (ii) We exploit the  $L^2$ -Wasserstein distance to constrain the policy updates, which is important to guarantee stability in policy optimization [55, 45, 56]; (iii) By constraining to specific
submanifolds of the Wasserstein space, in this case GMMs, we can impose additional structural properties on the policy optimization.
Since our objective in Eq. 10 has the form of the free energy functional studied by [36, 37], we can leverage the iterative updates scheme of Eq. 5 to formulate the evolution of our policy iteration under the flow generated by Eq. 10. 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  $\mathrm{GMM}_d$ . We refer the interested reader to App. A.3, where we provide the explicit form of  $J(\pi)$  of Eq. 10 for the GMM case.
3.1 Policy optimization