SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\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(s _ {t}\right), \boldsymbol {\Sigma} _ {i} ^ {a}\right), \quad \text {w i t h} \tag {16}
$$
$$
\boldsymbol {\mu} _ {i} ^ {a} \left(s _ {t}\right) = \boldsymbol {\mu} _ {i} ^ {a} + \boldsymbol {\Sigma} _ {i} ^ {a s} \left(\boldsymbol {\Sigma} _ {i} ^ {s}\right) ^ {- 1} \left(\boldsymbol {s} _ {t} - \boldsymbol {\mu} _ {i} ^ {s}\right), \tag {17}
$$
$$
\boldsymbol {\Sigma} _ {i} ^ {a} = \boldsymbol {\Sigma} _ {i} ^ {a} - \boldsymbol {\Sigma} _ {i} ^ {a s} \left(\boldsymbol {\Sigma} _ {i} ^ {s}\right) ^ {- 1} \boldsymbol {\Sigma} _ {i} ^ {s a}, \tag {18}
$$
$$
\omega_ {i} \left(\boldsymbol {s} _ {t}\right) = \frac {\omega_ {i} \mathcal {N} \left(\boldsymbol {s} _ {t} ; \boldsymbol {\mu} _ {i} ^ {s} , \boldsymbol {\Sigma} _ {i} ^ {s}\right)}{\sum_ {k} ^ {n} \omega_ {k} \mathcal {N} \left(\boldsymbol {s} _ {t} ; \boldsymbol {\mu} _ {k} ^ {s} , \boldsymbol {\Sigma} _ {k} ^ {s}\right)}. \tag {19}
$$
Note that we have split the GMM parameters  $\pmb{\mu}_i$  and  $\Sigma_{i}$  into their state and action components according to
$$
\boldsymbol {\mu} _ {i} = \left( \begin{array}{l} \boldsymbol {\mu} _ {i} ^ {s} \\ \boldsymbol {\mu} _ {i} ^ {a} \end{array} \right), \quad \boldsymbol {\Sigma} _ {i} = \left( \begin{array}{c c} \boldsymbol {\Sigma} _ {i} ^ {s} & \boldsymbol {\Sigma} _ {i} ^ {s a} \\ \boldsymbol {\Sigma} _ {i} ^ {a s} & \boldsymbol {\Sigma} _ {i} ^ {a} \end{array} \right). \tag {20}
$$
A.2 Riemannian gradients and retractions
For completeness we give here the explicit expressions of the Riemannian gradients and the retractions used in § 3.1. As the mean vectors are assumed to lie in the Euclidean space, their Riemannian gradients actually coincide with the Euclidean gradients and no retraction is required, so Eq. 12 reduces to the well-known Euclidean gradient descent
$$
\hat {\boldsymbol {\mu}} _ {k + 1} = \hat {\boldsymbol {\mu}} _ {k} + \nabla_ {\hat {\boldsymbol {\mu}}} J (\pi_ {k}), \tag {21}
$$
where  $\nabla_{\hat{\mu}}$  denotes the Euclidean gradient w.r.t.  $\hat{\mu}$ . For the covariance matrices we use the gradient and retraction w.r.t. the Bures-Wasserstein manifold, taken from [59, 60]. The gradient is given by
$$
\operatorname {g r a d} _ {\hat {\boldsymbol {\Sigma}}} J (\pi_ {k}) = 4 \{\nabla_ {\hat {\boldsymbol {\Sigma}}} J (\pi_ {k}) \hat {\boldsymbol {\Sigma}} \} _ {S}, \tag {22}
$$
where again  $\nabla_{\hat{\Sigma}}$  denotes the Euclidean gradient w.r.t.  $\hat{\Sigma}$  and  $\{X\}_{S} = \frac{(X + X^{\top})}{2}$ .
Furthermore, the retraction is given by
$$
\mathrm {R} _ {\boldsymbol {\Sigma} _ {k}} (\hat {\boldsymbol {X}}) = \hat {\boldsymbol {\Sigma}} _ {k} + \hat {\boldsymbol {X}} + \mathcal {L} _ {\hat {\boldsymbol {X}}} [ \hat {\boldsymbol {\Sigma}} _ {k} ] \hat {\boldsymbol {X}} \mathcal {L} _ {\hat {\boldsymbol {X}}} [ \hat {\boldsymbol {\Sigma}} _ {k} ], \tag {23}
$$
where  $\mathcal{L}_{\hat{\boldsymbol{X}}}\left[\hat{\boldsymbol{\Sigma}}_k\right]$  is the Lyapunov operator, defined as the solution to the matrix linear system
$$
\mathcal {L} _ {\hat {\boldsymbol {X}}} \left[ \hat {\boldsymbol {\Sigma}} _ {k} \right] \hat {\boldsymbol {X}} + \hat {\boldsymbol {X}} \mathcal {L} _ {\hat {\boldsymbol {X}}} \left[ \hat {\boldsymbol {\Sigma}} _ {k} \right] = \hat {\boldsymbol {\Sigma}} _ {k}. \tag {24}
$$
A.3 Expressions of the free functional  $J(\pi)$  and its Euclidean gradients
For completeness sake, we provide here the explicit expression of the Euclidean gradients for the objective  $J(\pi)$  w.r.t. the parameters of the GMM, which are used in the construction of the Riemannian gradients. Using the policy gradient theorem, we obtain the gradient of Eq. 11 w.r.t to a
parameter  $\xi$  as follows
$$
\begin{array}{l} \nabla_ {\boldsymbol {\xi}} J (\pi) = \nabla_ {\boldsymbol {\xi}} \int \Pi_ {t} \mathrm {d} \boldsymbol {s} _ {0} \mathrm {d} \boldsymbol {s} _ {t} \mathrm {d} \boldsymbol {a} _ {t} \rho (\boldsymbol {s} _ {0}) \pi (\boldsymbol {a} _ {t} | \boldsymbol {s} _ {t}) p (\boldsymbol {s} _ {t + 1} | \boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}) \sum_ {t ^ {\prime} > t} \gamma^ {t} r \left(\boldsymbol {s} _ {t ^ {\prime}}, \boldsymbol {a} _ {t ^ {\prime}}\right), \tag {25} \\ = \mathbb {E} _ {\tau} \left[ \sum_ {t} \nabla_ {\boldsymbol {\xi}} \log (\pi (\boldsymbol {a} _ {t} | \boldsymbol {s} _ {t})) \sum_ {t ^ {\prime} > t} r (\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}) \right], \\ = \mathbb {E} _ {\tau} \left[ \sum_ {t} \nabla_ {\boldsymbol {\xi}} \log \left(\frac {\pi (\boldsymbol {s} _ {t} , \boldsymbol {a} _ {t})}{\int \mathrm {d} \boldsymbol {a} _ {t} \pi (\boldsymbol {s} _ {t} , \boldsymbol {a} _ {t})}\right) \sum_ {t ^ {\prime} > t} r (\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}) \right], \\ = \sum_ {t} \mathbb {E} _ {\tau} \left[ \left(\frac {\nabla_ {\boldsymbol {\xi}} \pi (\boldsymbol {s} _ {t} , \boldsymbol {a} _ {t})}{\pi (\boldsymbol {s} _ {t} , \boldsymbol {a} _ {t})} - \frac {\int \mathrm {d} \boldsymbol {a} _ {t} \nabla_ {\boldsymbol {\xi}} \pi (\boldsymbol {s} _ {t} , \boldsymbol {a} _ {t})}{\int \mathrm {d} \boldsymbol {a} _ {t} \pi (\boldsymbol {s} _ {t} , \boldsymbol {a} _ {t})}\right) \sum_ {t ^ {\prime} > t} r (\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}) \right]. \\ \end{array}
$$
In this work, we focus on GMM policies, for which the objective  $J(\pi)$  takes the form:
$$
\begin{array}{l} J (\pi) = \int \Pi_ {t} \mathrm {d} \boldsymbol {s} _ {0} \mathrm {d} \boldsymbol {s} _ {t} \mathrm {d} \boldsymbol {a} _ {t} \rho (\boldsymbol {s} _ {0}) \sum_ {i = 1} ^ {n} \omega_ {i} (\boldsymbol {s} _ {t}) \mathcal {N} (\boldsymbol {a} _ {t}; \boldsymbol {\mu} _ {i} (\boldsymbol {s} _ {t}), \boldsymbol {\Sigma} _ {i} (\boldsymbol {s} _ {t})) p (\boldsymbol {s} _ {t + 1} | \boldsymbol {s} _ {t}, a _ {t}) \sum_ {t} \gamma^ {t} r (\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}) \\ + \beta \int \mathrm {d} \boldsymbol {a} _ {t} \sum_ {i = 1} ^ {n} \omega_ {i} (\boldsymbol {s} _ {t}) \mathcal {N} (\boldsymbol {a} _ {t}; \mu_ {i} (\boldsymbol {s} _ {t}), \Sigma_ {i} (\boldsymbol {s} _ {t})) p (\boldsymbol {s} _ {t + 1} | \boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}) \\ \log \left(\sum_ {i = 1} ^ {n} \omega_ {i} \left(s _ {t}\right) \mathcal {N} \left(\boldsymbol {a} _ {t}; \mu_ {i} \left(\boldsymbol {s} _ {t}\right), \Sigma_ {i} \left(\boldsymbol {s} _ {t}\right)\right) p \left(\boldsymbol {s} _ {t + 1} \mid \boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}\right)\right). \tag {26} \\ \end{array}
$$
By inserting Eq. 26 into Eq. 25 we obtain for the individual parameters of the GMM
$$
\begin{array}{l} \nabla \boldsymbol {\mu} _ {l} J (\pi) = \mathbb {E} _ {\tau} \left[ \sum_ {t} \left(\frac {\omega_ {l} \mathcal {N} \left(\boldsymbol {s} _ {t} , \boldsymbol {a} _ {t} ; \boldsymbol {\mu} _ {l} , \boldsymbol {\Sigma} _ {l}\right) \boldsymbol {\Sigma} _ {l} ^ {- 1} \left(\left(\boldsymbol {s} _ {t} , \boldsymbol {a} _ {t}\right) - \boldsymbol {\mu} _ {l}\right)}{\sum_ {j} \omega_ {j} \mathcal {N} \left(\boldsymbol {s} _ {t} , \boldsymbol {a} _ {t} ; \boldsymbol {\mu} _ {j} , \boldsymbol {\Sigma} _ {j}\right)} \right. \right. (27) \\ \left. - \frac {\omega_ {l} \int \mathrm {d} \boldsymbol {a} \mathcal {N} (\boldsymbol {s} _ {t} , \boldsymbol {a} _ {t} ; \boldsymbol {\mu} _ {l} , \boldsymbol {\Sigma} _ {l}) \boldsymbol {\Sigma} _ {l} ^ {- 1} ((\boldsymbol {s} _ {t} , \boldsymbol {a} _ {t}) - \boldsymbol {\mu} _ {l})}{\sum_ {j} \omega_ {j} \int \mathrm {d} \boldsymbol {a} \mathcal {N} (\boldsymbol {s} _ {t} , \boldsymbol {a} _ {t} ; \boldsymbol {\mu} _ {j} , \boldsymbol {\Sigma} _ {j})}\right) \sum_ {t ^ {\prime} > t} r (\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}) \Biggr ], \\ = \mathbb {E} _ {\tau} \left[ \sum_ {t} \left(\frac {\omega_ {l} \mathcal {N} \left(\boldsymbol {s} _ {t} , \boldsymbol {a} _ {t} ; \boldsymbol {\mu} _ {l} , \boldsymbol {\Sigma} _ {l}\right) \boldsymbol {\Sigma} _ {l} ^ {- 1} \left(\left(\boldsymbol {s} _ {t} , \boldsymbol {a} _ {t}\right) - \boldsymbol {\mu} _ {l}\right)}{\sum_ {j} \omega_ {j} \mathcal {N} \left(\boldsymbol {s} _ {t} , \boldsymbol {a} _ {t} ; \boldsymbol {\mu} _ {j} , \boldsymbol {\Sigma} _ {j}\right)}\right) \right. (28) \\ \left. - \delta_ {\boldsymbol {s}} \frac {\omega_ {l} \mathcal {N} (\boldsymbol {s} _ {t} ; \boldsymbol {\mu} _ {l , s} \boldsymbol {\Sigma} _ {l , s s}) \boldsymbol {\Sigma} _ {l , s s} ^ {- 1} (\boldsymbol {s} _ {t} - \boldsymbol {\mu} _ {l , s})}{\sum_ {j} \omega_ {j} \mathcal {N} (\boldsymbol {s} _ {t} ; \boldsymbol {\mu} _ {j , s} , \boldsymbol {\Sigma} _ {j , s s})}\right) \sum_ {t ^ {\prime} > t} r (\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}) \Bigg ], \\ = \mathbb {E} _ {\tau} \left[ \sum_ {t} \left(\zeta_ {l, \boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}} \boldsymbol {\Sigma} _ {l} ^ {- 1} \left(\left(\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}\right) - \boldsymbol {\mu} _ {l}\right) - \delta_ {s} \zeta_ {l, \boldsymbol {s} _ {t}} \boldsymbol {\Sigma} _ {l, s s} ^ {- 1} \left(\boldsymbol {s} _ {t} - \boldsymbol {\mu} _ {l, s}\right)\right) \sum_ {t ^ {\prime} > t} r \left(\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}\right) \right]. \\ \end{array}
$$
Here  $\delta_s \in \{0,1\}$  indicates which terms the gradient acts on. In this case, the gradient act on the state components and it is absent for the action dimensions.
$$
\begin{array}{l} \nabla_ {\boldsymbol {\Sigma} _ {l}} J (\pi) = \mathbb {E} _ {\tau} \left[ \right. \sum_ {t} \left(- \frac {1}{2} \frac {\omega_ {l} \mathcal {N} \left(\boldsymbol {s} _ {t} , \boldsymbol {a} _ {t} ; \boldsymbol {\mu} _ {l} , \boldsymbol {\Sigma} _ {l}\right) \boldsymbol {\Sigma} _ {l} ^ {- 1} \left(1 - \left(\left(\boldsymbol {s} _ {t} , \boldsymbol {a} _ {t}\right) - \boldsymbol {\mu} _ {l}\right)\left(\left(\boldsymbol {s} _ {t} , \boldsymbol {a} _ {t}\right) - \boldsymbol {\mu} _ {l}\right) ^ {\top} \boldsymbol {\Sigma} _ {l} ^ {- 1}\right)}{\sum_ {j} \omega_ {j} \mathcal {N} \left(\boldsymbol {s} _ {t} , \boldsymbol {a} _ {t}; \boldsymbol {\mu} _ {j} , \boldsymbol {\Sigma} _ {j}\right)} \right. \tag {29} \\ \left. + \frac {1}{2} \frac {\omega_ {l} \int \mathrm {d} \boldsymbol {a} \mathcal {N} (\boldsymbol {s} _ {t} , \boldsymbol {a} _ {t} ; \boldsymbol {\mu} _ {l} , \boldsymbol {\Sigma} _ {l}) \boldsymbol {\Sigma} _ {l} ^ {- 1} \left(1 - \left(\left(\boldsymbol {s} _ {t} , \boldsymbol {a} _ {t}\right) - \boldsymbol {\mu} _ {l}\right) \left(\left(\boldsymbol {s} _ {t} , \boldsymbol {a} _ {t}\right) - \boldsymbol {\mu} _ {l}\right) ^ {\top} \boldsymbol {\Sigma} _ {l} ^ {- 1}\right)}{\sum_ {j} \omega_ {j} \int \mathrm {d} a \mathcal {N} (\boldsymbol {s} _ {t} , \boldsymbol {a} _ {t} ; \boldsymbol {\mu} _ {j} , \boldsymbol {\Sigma} _ {j})}\right) \sum_ {t ^ {\prime} > t} r (\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}) \Bigg ], \\ = \mathbb {E} _ {\tau} \left[ \sum_ {t} \left(- \frac {\zeta_ {l , \boldsymbol {s} _ {t} , \boldsymbol {a} _ {t}}}{2} \boldsymbol {\Sigma} _ {l} ^ {- 1} \left(1 - \left(\left(\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}\right) - \boldsymbol {\mu} _ {l}\right) \left(\left(\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}\right) - \boldsymbol {\mu} _ {l}\right) ^ {\top} \boldsymbol {\Sigma} _ {l} ^ {- 1}\right) \right. \right. \\ \left. \left. + \delta_ {s} \frac {\zeta_ {l , s _ {t}}}{2} \boldsymbol {\Sigma} _ {l, s} ^ {- 1} \left(1 - \left(\boldsymbol {s} _ {t} - \boldsymbol {\mu} _ {l, s}\right) \left(\boldsymbol {s} _ {t} - \boldsymbol {\mu} _ {l, s}\right) ^ {\top} \boldsymbol {\Sigma} _ {l, s} ^ {- 1}\right)\right) \sum_ {t ^ {\prime} > t} r (\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}) \right]. \\ \end{array}
$$
$$
\begin{array}{l} \nabla \omega_ {l} J (\pi) = \mathbb {E} _ {\tau} \left[ \sum_ {t} \left(\frac {\mathcal {N} \left(\boldsymbol {s} _ {t} , \boldsymbol {a} _ {t} ; \boldsymbol {\mu} _ {l} , \boldsymbol {\Sigma} _ {l}\right)}{\sum_ {j} \omega_ {j} \mathcal {N} \left(\boldsymbol {s} _ {t} , \boldsymbol {a} _ {t} ; \boldsymbol {\mu} _ {j} , \boldsymbol {\Sigma} _ {j}\right)} - \frac {\int \mathrm {d} \boldsymbol {a} \mathcal {N} \left(\boldsymbol {s} _ {t} , \boldsymbol {a} _ {t} ; \boldsymbol {\mu} _ {l} , \boldsymbol {\Sigma} _ {l}\right)}{\sum_ {j} \omega_ {j} \int \mathrm {d} \boldsymbol {a} \mathcal {N} \left(\boldsymbol {s} _ {t} , \boldsymbol {a} _ {t} ; \boldsymbol {\mu} _ {j} , \boldsymbol {\Sigma} _ {j}\right)}\right) \sum_ {t ^ {\prime} > t} r \left(\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}\right) \right] (30) \\ = \mathbb {E} _ {\tau} \left[ \sum_ {t} \frac {\left(\zeta_ {l , \boldsymbol {s} _ {t}} , \boldsymbol {a} _ {t} - \zeta_ {l , \boldsymbol {s} _ {t}}\right)}{\omega_ {l}} \sum_ {t ^ {\prime} > t} r \left(\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}\right) \right]. (31) \\ \end{array}
$$
Note that we introduced the responsibilities  $\zeta_{l,\pmb{s}_t},\pmb{a}_t$  and  $\zeta_{l,\pmb{s}_t}$ , which are defined as follows
$$
\zeta_ {l, \boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}} = \frac {\omega_ {l} \mathcal {N} \left(\boldsymbol {s} _ {t} , \boldsymbol {a} _ {t} ; \boldsymbol {\mu} _ {l} , \boldsymbol {\Sigma} _ {l}\right)}{\sum_ {j} \omega_ {j} \mathcal {N} \left(\boldsymbol {s} _ {t} , \boldsymbol {a} _ {t} ; \boldsymbol {\mu} _ {j} , \boldsymbol {\Sigma} _ {j}\right)}, \quad \text {a n d} \tag {32}
$$
$$
\zeta_ {l, \boldsymbol {s} _ {t}} = \frac {\omega_ {l} \int \mathrm {d} \boldsymbol {a} \mathcal {N} \left(\boldsymbol {s} _ {t} , \boldsymbol {a} _ {t} ; \boldsymbol {\mu} _ {l} , \boldsymbol {\Sigma} _ {l}\right)}{\sum_ {j} \omega_ {j} \int \mathrm {d} \boldsymbol {a} \mathcal {N} \left(\boldsymbol {s} _ {t} , \boldsymbol {a} _ {t} ; \boldsymbol {\mu} _ {j} , \boldsymbol {\Sigma} _ {j}\right)} = \frac {\omega_ {l} \mathcal {N} \left(\boldsymbol {s} _ {t} ; \boldsymbol {\mu} _ {l , s} , \boldsymbol {\Sigma} _ {l , s s}\right)}{\sum_ {j} \omega_ {j} \mathcal {N} \left(\boldsymbol {s} _ {t} ; \boldsymbol {\mu} _ {j , s} , \boldsymbol {\Sigma} _ {j , s s}\right)}. \tag {33}
$$
A.4 Relation between forward and backward discretization in the Bures-Wasserstein metric
In this section we outline the relation between the implicit and explicit optimization schema w.r.t. the Bures-Wasserstein metric, which is leveraged to formulate our policy optimization in  $\S 3$ . We closely follow [58]. For the sake of simplicity, we group the Gaussian parameters  $\pmb{\mu}$  and  $\pmb{\Sigma}$  into a single parameter vector  $\theta$ . 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  $\mathbb{R}^d\times \mathcal{S}_{++}^d$ .
The Riemannian gradient w.r.t the Gaussian parameters  $\theta$ ,  $\operatorname{grad}_{\theta} J(\pi(\theta))$ , satisfies by definition
$$
g _ {\boldsymbol {\theta}} \left(\operatorname {g r a d} _ {\boldsymbol {\theta}} J (\pi (\boldsymbol {\theta}), \boldsymbol {\xi}) = \nabla_ {\boldsymbol {\theta}} J (\pi (\boldsymbol {\theta})) \cdot \boldsymbol {\xi}, \right. \tag {34}
$$
where  $\nabla_{\theta}$  denotes the Euclidean gradient,  $\xi$  is an arbitrary vector on the tangent space  $\mathcal{T}_{\theta}\mathcal{M}$ , and  $g_{\theta}$  is the Riemannian metric tensor, defining the inner product on  $\mathcal{T}_{\theta}\mathcal{M}$ . The Riemannian metric  $g_{\theta}$  can be written as
$$
g _ {\boldsymbol {\theta}} (\boldsymbol {\zeta}, \boldsymbol {\xi}) = \boldsymbol {\zeta} ^ {\mathsf {T}} G _ {W} (\boldsymbol {\theta}) \boldsymbol {\xi}, \tag {35}
$$
with two arbitrary tangent vectors  $\zeta, \xi$ , and  $G_W(\theta)$  being a positive definite matrix. Moreover, note that the Wasserstein distance  $W_2^2\big(\mathcal{N}(\theta), \mathcal{N}(\theta + \Delta \theta)\big)$ , where  $\Delta \theta$  denotes a small perturbation in the Gaussian parameters  $\theta$ , can be expressed as
$$
\left. W _ {2} ^ {2} \left(\mathcal {N} (\boldsymbol {\theta}), \mathcal {N} (\boldsymbol {\theta} + \Delta \boldsymbol {\theta})\right) = \frac {1}{2} \left(\Delta \boldsymbol {\theta} ^ {\mathrm {T}}\right) G _ {W} (\boldsymbol {\theta}) \left(\Delta \boldsymbol {\theta}\right) + O \left((\Delta \boldsymbol {\theta}) ^ {2}\right), \right. \tag {36}
$$
for  $\Delta \theta \to 0$ . Similarly, we can approximate the objective evaluated at  $J(\pmb{\theta} + \Delta \pmb{\theta})$  via the Taylor theorem as
$$
J (\boldsymbol {\theta} + \Delta \boldsymbol {\theta}) = J (\boldsymbol {\theta}) + \nabla_ {\boldsymbol {\theta}} J (\boldsymbol {\theta}) \cdot \Delta \boldsymbol {\theta} + O ((\Delta \boldsymbol {\theta}) ^ {2}). \tag {37}
$$
With this, we can approximate
$$
\begin{array}{l} \boldsymbol {\theta} _ {k + 1} = \underset {\boldsymbol {\theta}} {\arg \min } \left(\frac {W _ {2} ^ {2} (\pi (\boldsymbol {\theta}) , \pi (\boldsymbol {\theta} _ {k}))}{2 \tau} - J (\pi (\boldsymbol {\theta}))\right), (38) \\ \approx \underset {\boldsymbol {\theta}} {\arg \min } \left(\frac {(\boldsymbol {\theta} - \boldsymbol {\theta} _ {\boldsymbol {k}}) ^ {\top} G _ {W} (\boldsymbol {\theta} - \boldsymbol {\theta} _ {\boldsymbol {k}})}{2 \tau} - \nabla_ {\boldsymbol {\theta}} J (\boldsymbol {\theta}) \cdot (\boldsymbol {\theta} - \boldsymbol {\theta} _ {\boldsymbol {k}})\right), (39) \\ \end{array}
$$
from which we obtain the update equation for  $\theta$  as follows
$$
\boldsymbol {\theta} _ {k + 1} = \boldsymbol {\theta} _ {k} + \tau G _ {W} \left(\boldsymbol {\theta} _ {k}\right) ^ {- 1} \nabla_ {\boldsymbol {\theta}} J \left(\pi \left(\boldsymbol {\theta} _ {k}\right)\right), \tag {40}
$$
which corresponds to the Wasserstein natural gradient with respect to the Bures-Wasserstein metric. Note that Eq. 40 in turn corresponds to an approximation of the exact Riemannian gradient descent
$$
\boldsymbol {\theta} _ {k + 1} = \operatorname {R} _ {\boldsymbol {\theta} _ {k}} (\tau \cdot \operatorname {g r a d} _ {\boldsymbol {\theta}} J (\pi (\boldsymbol {\theta} _ {k}))). \tag {41}
$$
This approximation can be obtained by considering a first-order approximation of the geodesic on the  $BW$  manifold. As the exponential map (a.k.a. the retraction) is defined via the geodesic, the retraction operator in Eq. 41 turns into a simple addition operation under a first-order approximation, leading to Eq. 40. Notice that such approximation does not guarantee that the updated parameters  $\theta$  stay on the manifold, except for the cases in which  $\theta \in \mathbb{R}^d$ . This means that the positive-definite constraints arising from the covariance matrices can be easily violated when using natural gradient approaches as they do not guarantee that the updates stay on the underlying manifold. However, it is worth noting that the use of an inexact Riemannian gradient descent is often motivated by the difficulty of computing the exponential map (or geodesic) necessary to calculate the exact Riemannian update. In our case, we leverage the retraction and Riemannian gradients of [59, 60], which allow us to apply the exact Riemannian gradient descent of 41. This avoids relying 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  $\left(\mathbb{R}^d \times \mathcal{S}_{++}^d\right)^N$ .
A.5 Additional details on the implementation
We extended the Pymanopt [64] 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. The constraint function  $c(x_0, \cdot)$  is arbitrary in general. We use the  $L^2$ -Wasserstein distance between two points on the manifold of GMMs as constraint.
Input: point  $x_0$  on the manifold, descent direction  $d$ , initial step size  $\lambda_0$ , decrement  $\alpha$ , constraint  $c(x_0, \cdot)$ , maximum allowed value for constraint  $c_{\mathrm{max}}$ , minimum step size  $\lambda_{\mathrm{min}}$
Output: step size  $s$ , updated point on manifold  $\mathcal{X}$
1:  $\mathbf{x} = \mathbf{x}_0 + \lambda_0 \cdot \mathbf{d}$
$\lambda = \lambda_0$
2: while  $c(\pmb{x_0}, \pmb{x}) > c_{\max}$  and  $\lambda > \lambda_{\min}$  do
3: decrease step size:  $\lambda = \alpha \cdot \lambda$  update point on manifold:  $x = x_0 + \lambda \cdot d$
4: end while
5: if  $\lambda < \lambda_{\mathrm{min}}$  then
6: return  $\lambda_0, x_0$
7: else
8: return  $\lambda, x$
9: end if
A.6 Additional details on experiments
A.6.1 Additional results
Fig. 5 shows the convergence curves for the two baselines as in Fig. 3 of the main paper, however, we extended the horizontal axis up to the maximum number of environment steps used for training.
Figure 5: The success rate of the two baselines on the reaching task (left), the collision-avoidance task (middle) and the multiple-goal task (right). The shaded area indicates the standard deviation over 5 runs.
Fig. 6 shows the variance of the success rate for the three methods at their time step of convergence for all three robotic tasks. Concerning SAC, which did not converge after the maximum number of
Figure 6: Variance of the success rate over the 5 runs for our method (WGF) and the two baselines on the reaching task (left), the collision avoidance task (middle) and the multiple-goal task (right). The violin plots are overlaid with box plots, quartile lines and a swarm plot, where dots indicate the success rates of individual runs. The time steps at which we determined the variance are for PPO, SAC and WGF for the three tasks from left to right: (280000, 400000, 80000), (275000, 300000, 90000), (130000, 200000, 95000).
Figure 7: The success rate of our method and an ablated version, not using the Bures-Wasserstein formulation for the reaching task (left), the collision-avoidance task (middle) and the multiple-goal task (right). The shaded area indicates the standard deviation over 5 runs.
environment steps used for training, we chose the last time step. Specifically, we chose the following time steps for PPO, SAC and WGF, respectively: reaching task (280000, 400000, 80000), collision avoidance task (275000, 300000, 90000), multiple goal task (130000, 200000, 95000). These plots show that PPO may also reach low-variance success rate over the five runs at the time step of convergence, at the cost of a prohibitively large number of steps. SAC showed huge variance in all tasks, apart from the reaching task, where all runs collapsed to a success rate of 0.
A.6.2 Ablation study
In order to assess the influence of leveraging a Riemannian optimization approach on the Bures-Wasserstein manifold, we conducted an ablation of our method by eliminating the Riemannian formulation. Instead of the explicit Euler scheme update in Eq. 12, which corresponds to Riemannian gradient descent w.r.t. the Bures-Wasserstein metric, we use the implicit Euler scheme
$$
\hat {\boldsymbol {\mu}} _ {k + 1} = \underset {\hat {\boldsymbol {\mu}}} {\arg \min } \left(\frac {W _ {2} ^ {2} \left(\pi_ {k} (\hat {\boldsymbol {\mu}}) , \pi_ {k}\right)}{2 \tau} - J \left(\pi_ {k} (\hat {\boldsymbol {\mu}})\right)\right), \tag {42}
$$
$$
\hat {\boldsymbol {\Sigma}} _ {k + 1} = \underset {\hat {\boldsymbol {\Sigma}}} {\arg \min } \left(\frac {W _ {2} ^ {2} \left(\pi_ {k} (\hat {\boldsymbol {\Sigma}}) , \pi_ {k}\right)}{2 \tau} - J \left(\pi_ {k} (\hat {\boldsymbol {\Sigma}})\right)\right). \tag {43}
$$
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. The results obtained with this non-Riemannian version of our method are shown in Fig. 7 in direct comparison to our method and Fig. 8 for an extended range.
Figure 8: Extended plot of the success rate of and ablated version of our method, not using the Bures-Wasserstein-based formulation for the reaching task (left), the collision-avoidance task (middle) and the multiple-goal task (right). The shaded area indicates the standard deviation over 5 runs.