Wasserstein Gradient Flows for Optimizing Gaussian Mixture Policies

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 \emph{de facto} paradigm to adapt robot policies as a function of task-specific objectives. Most commonly-used motion policies carry particular structures that are often overlooked in policy optimization algorithms. We instead propose to leverage the structure of probabilistic policies by casting the policy optimization as an optimal transport problem. Specifically, we focus on robot motion policies that build on Gaussian mixture models (GMMs) and formulate the policy optimization as a Wassertein gradient flow over the GMMs space. This naturally allows us to constrain the policy updates via the $L^2$-Wasserstein distance between GMMs to enhance the stability of the policy optimization process. Furthermore, we leverage the geometry of the Bures-Wasserstein manifold to optimize the Gaussian distributions of the GMM policy via Riemannian optimization. We evaluate our approach 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.

Paper

References (77)

Scroll for more · 38 remaining

Similar papers

Peer review

Reviewer vzrR7/10 · confidence 3/52023-06-27

Summary

The paper proposes a new approach to adapt robot motion policies for different task conditions. It suggests leveraging the structure of probabilistic policies, specifically Gaussian mixture models (GMMs), and formulating policy optimization as an optimal transport problem. By using the L2-Wasserstein distance between GMMs, the policy updates can be constrained to improve the stability of the optimization process. Additionally, the Bures-Wasserstein manifold geometry is utilized for optimizing the Gaussian distributions of the GMM policy through Riemannian optimization. The proposed method is evaluated on various robotic scenarios and demonstrates better performance in terms of task success rate and low-variance solutions compared to common policy optimization baselines.

Strengths

Novel Approach: The paper introduces a new approach to adapt robot motion policies by leveraging the structure of probabilistic policies and formulating policy optimization as an optimal transport problem. This novel perspective can provide insights into enhancing the adaptability and performance of robot motion policies. Utilization of Gaussian Mixture Models (GMMs): The paper focuses on GMMs, a widely used representation for modeling complex motion policies. By exploiting the specific structure of GMMs, the proposed method offers a tailored solution for policy optimization, which can potentially lead to more effective and efficient adaptation of robot motion policies. Consideration of Stability: The incorporation of the L2-Wasserstein distance between GMMs as a constraint in policy updates aims to enhance the stability of the optimization process. This consideration addresses a common challenge in policy optimization algorithms and can contribute to more reliable and consistent results. Riemannian Optimization: The paper leverages the geometry of the Bures-Wasserstein manifold for optimizing the Gaussian distributions of the GMM policy. This utilization of Riemannian optimization techniques showcases a sophisticated mathematical framework to refine the policy parameters, potentially leading to improved performance and convergence properties. Experimental Evaluation: The proposed method is thoroughly evaluated on common robotic settings, including reaching motions, collision-avoidance behaviors, and multi-goal tasks. The results demonstrate that the proposed approach outperforms common policy optimization baselines in terms of task success rate and low-variance solutions. This empirical validation strengthens the credibility and practical relevance of the proposed method.

Weaknesses

I noticed that this does not discuss much other GMM-based methods, such as the PMOE[1]. Including a discussion and comparison with existing GMM-based methods would greatly enhance the comprehensiveness and value of your work. Furthermore, I suggest you consider conducting additional experiments to compare your proposed approach with these methods. Because comparing only with non-GMM methods, like PPO and SAC, is not fair. [1] Ren J, Li Y, Ding Z, et al. Probabilistic mixture-of-experts for efficient deep reinforcement learning[J]. arXiv preprint arXiv:2104.09122, 2021.

Questions

N/A

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

N/A

Reviewer uu8b5/10 · confidence 3/52023-06-30

Summary

This paper proposes to formulate policy optimization as a Wasserstein gradient flow over the Gaussian Mixture Model (GMM) space, which enhances the stability of policy optimization processes. In the proposed GMM policy updates, the mean and variance of gaussians are optimized through Riemann gradient descent via the Bures-Wasserstein metric, and the weight of gaussians are optimized through the implicit Euler scheme. The paper then demonstrates the effectiveness of their proposed approach by conducting experiments on various robotic tasks in simulation.

Strengths

Overall the paper is well-written and well-presented. Discovering more principled approaches to optimize GMM policies to model more complex policy distributions is a meaningful research direction.

Weaknesses

My main concern lies in the quality of experiments and empirical evaluations. Specifically, - The three robotic tasks illustrated in the main experiments are rather simple tasks with low-dimensional observation and action spaces. It would be interesting to demonstrate whether the proposed approach can generalize to higher dimensional tasks that exhibit multimodality solutions. - More in-depth analysis on why the proposed approach outperforms the baselines (PPO+GMM & SAC+GMM) can be included. Specifically, it would be meaningful to qualitatively compare policy distributions between the proposed approach and the baselines, particularly for the observations where the proposed approach is better than the baselines in modeling policy distributions. This would significantly enrich the insights provided by the paper.

Questions

See "weaknesses".

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

Limitations are mostly addressed. Another limitation to include is regarding the current empirical evaluations on low-dimensional robotic tasks. Showing the potential for the proposed approach to scale up to higher dimensional tasks with multimodality natures would significantly enhance the paper.

Reviewer UKYq4/10 · confidence 4/52023-07-01

Summary

The authors investigate Wasserstein gradient flows (WGF) for a Gaussian mixture model policy. The WGF is a principled natural gradient method for updating the parameters since it follows the metric space defined by the Wasserstein-2 divergence. This approach is compared to deep RL methods on some planar control tasks.

Strengths

This was a very nicely written paper with also excelled presentation. I commend the authors for this. I also think WGF is an interesting research direction for RL.

Weaknesses

**Ignores natural gradients in RL** I was very surprised to see this paper almost completely ignores the work on natural gradient methods for RL. While TRPO is cited, it is only very vaguely with no mention of the KL natural gradient it uses. Here are some key references: [1] Amari, S. I. (1998). Natural gradient works efficiently in learning. Neural computation [2] A natural policy gradient, Sham M Kakade - Advances in neural information processing, 2001 [3] Natural actor-critic J Peters, S Schaal - Neurocomputing, 2008 Related to this is the linear programming view of RL, which also uses KL regularization (but uses EM rather than NG) [4] Relative entropy policy search, J Peters, K Mulling, Y Altun AAAI 2010 [5] A unified view of entropy-regularized markov decision processes, G Neu, A Jonsson, V Gómez, 2017 and mirror descent as well [6] Mirror Descent Policy Optimization, Manan Tomar, Lior Shani, Yonathan Efroni, Mohammad Ghavamzadeh, 2020 These methods are basically Equation 5 of this paper, but with the KL rather than the W_2 divergence. The interesting question for me would be how these two divergences compare in the context of policy optimization. This seems to be central in the recent paper by Moskovitz et al but completely absent in this work. I think there has been a lot of work on doing NG methods for GMMs, e.g. Handling the Positive-Definite Constraint in the Bayesian Learning Rule, Lin et al 2020. This criticism naturally extends to the choice of baselines in the experiments. There should definitely be a natural policy gradient and EM-based method (for example MPO, advantage weighted regression or actor critic) to compare the quality of the update. I believe this is quite a big issue as it would require a major overhaul of the paper and experiments. The value of the contribution of the paper drops significantly if you cannot place WGF relative to prior work on natural gradients in RL. **Limited relevancy of GMR policies for NeurIPS** This paper focuses exclusively on GMM / GMR policies, which would also be called 'locally weighted regression' (LWR) back in the day. These models are rarely used because they scale very poorly to high dimensions, because you need either need hierarchy, dimensionality reduction or many many mixtures (e.g. [7]). [7] Locally Weighted Projection Regression, Vijayakumar et al ICML 2000. The cited works for GMR used them for learning from demonstrations for low-dimensional tasks. As a consequence, I feel like this work would be better suited for venues such as CoRL or ICRA where these methods may be still used, unless this method can be extended to use WGF on models more relevant to the Neurips community. **Misc** Typo: Equation 7 is missing a square root on the right-most squared Bures metric term.

Questions

1) For the 2D experiments, if the policies are initialized from demonstrations, why is this success rate 0? Are the baselines also initialized with the demonstrations? Also, consider using the *-from demonstrations variant of RL algorithms (e.g. SACfD). 2) I was a bit confused by the action space for the deep RL baselines. It seems the SAC-GMM paper focuses more on a hierarchical RL set up, where the SAC policy defines a GMM policy that runs for several timesteps. There is no reason why you could not have just replace the PPO and SAC policies with the GMR policy, since PPO just needs samples and log probabilities and SAC just needs to do the reparameterization trick. The issue with SAC would be that the max entropy regularization would be too powerful for GMR policies (since the SAC policy is usually clamped to upper-bound the entropy) so you would have to use the KL-regularied version of SAC, e.g. see Marino et al. Iterative amortized policy optimization, Marino et al NeurIPS 2021.

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

1 poor

Limitations

The focus on finetuning from demonstration is understated in the title and abstract. It's not clear from the experiments if this method + policy can be used to train a policy from scratch.

Reviewer zWwA7/10 · confidence 5/52023-07-06

Summary

This paper proposes an algorithm for RL in continuous action spaces based on the Wasserstein Gradient Flow formulation of Richemond and Maginnis. By specializing to the case of Gaussian Mixture policies, a simpler formulation is obtained, where the Gaussian part of the parameters (mean and covariance for each mixture component) are learned by optimizing a closed-form quadratic loss term rather than having to compute gradients of the Wasserstein_2 distance numerically.

Strengths

Just like in Zhang et al. although taking another route, Wasserstein_2 distance terms simplify to L2 distances thanks to the Gaussian assumption. In itself the algorithm represents a straightforward application of the principles described in Richemond & Maginnis or Zhang eet al., all the way to the use of the Sinkhorn algorithm, although the optimization part, framed as a splitting operation over two Riemannian gradient descent problems, is nicely executed. The paper is well written and easy to follow, even if several typos still remain ('Wassertein', 'WFG' instead of WGF...). Empirical results show that the Wasserstein Gradient Flow outperforms non optimal-transport motivated baselines.

Weaknesses

However, we do have several concerns regarding the scope and significance of these results. These span three axes. First, the PPO and SAC comparison baselines, while standard and still competitive, are 2017 or 2018 algorithms and would deserve updating. Second, only three robotics environments are considered, and it would be good to include more, in particular standard control domains. Finally, the Future Work section of the paper mentions multiple conceptual possibilities for improving and modifying Algorithm 1, as details such as the initialization of the Sinkhorn algorithm used for optimizing the mixture parameter can be absolutely critical in practice. We encourage the authors to implement some of those ideas. For these reasons, and being cognizant of the associated compute requirements, we feel the paper would much benefit from a revamped and extended empirical section implementing these ideas and testing them across a wider set of environments.

Questions

What can the authors do in order to improve empirical evaluation ? I would be willing to raise my score if some of the above concerns were addressed.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

Reviewer vzrR2023-08-11

Thank you for adding experiments to solve my questions within the limited time. I think I can raise my score based on your descriptions, but where are the detailed updated experiment results?

Authorsrebuttal2023-08-11

PDF location

Dear reviewer vzrR, Thank you for having read out rebuttal in such a timely manner. The figures of the experiments are attached to the PDF in the main reply (above) to all the reviewers and AC. The link is [here](https://openreview.net/attachment?id=QuMwbM0knB&name=pdf). Note that we will of course add the experiments details to the Appendix of the revised version of the paper as we were not allowed to add text to the attached PDF.

Reviewer vzrR2023-08-11

I appreciate the authors' rapid reply and the newly added experiments for answering my doubts. I don't think I have any other questions and I raised my score.

Authorsrebuttal2023-08-11

Thank you very much for your feedback and availability during the discussion phase, we believe your suggestions definitely improved our original paper.

Reviewer UKYq2023-08-11

Rebuttal response

Thanks for replying to my review. **Natural Gradients.** In the main rebuttal the authors state 'As the core of our paper is on Wasserstein gradient flows, we omitted the literature on KL divergence.', however Richemond et al. [30] and Moscovitz et al. [36] both make comparisons to natural gradient with the KL central to their paper's motivation on using the Wasserstein divergence, so it feels like a reasonable request that this submission make a similar acknowledgement and positioning. While the authors do cite Moscovitz et al., they do not explain what a natural gradient is in the main paper, how it compares to the WGF, nor do they mention the KL divergence once. The appendix also does not explicitly mention the KL divergence or natural gradient. I don't think it's fair to expect a reader to make the connection between forward and backward discretization and the natural gradient. It appears the authors have a lot to say about NG with KL vs WGF, so I'm surprised this discussion isn't included in the paper. In this field it is not uncommon that the mathematically more rigorous approach under-performs in practice against a more approximate method, so it would be good to see empirically that the WGF is the better approach. Currently, the complete omission of NG with the KL looks suspicious. **GMR policies.** My comments about GMMs specifically referenced the (linear) GMR policies used in the paper, not GMMs in general, so this weakness still stands. **Q1.** Why are policies pre-trained on demonstrations irrelevant to the task? What is the benefit of this? **Replay Buffers.** Soft actor critic is a deep off-policy reinforcement learning algorithm and therefore a replay buffer is central to its implementation. What have the authors implemented as 'SAC' if no replay buffers are used? The authors say in the main text that they use the SB3 Implementation of SAC, which has a replay buffer. **Q2.** If you read section IV.C of SAC-GMM [25], the GMM policy is executed for $N=32$ timesteps and the GMM 'represents a dynamical system controlling the motion in the trajectory space'. The paper doesn't call their method hierarchical but uses the term 'hybrid' instead to the same end. Could the author's clarify what algorithm they have implemented?

Authorsrebuttal2023-08-11

Natural gradients

We thank the reviewer for having read our rebuttal and for the timely feedback. Below our reply. 1. We agree with the reviewer that our Introduction section may be improved by revising previous works using the KL divergence in policy optimization, along similar lines as [30, 32, 36], so that our motivation of using the Wasserstein distance is more clear in the context of "regularized objectives" in policy optimization. We plan to include the following short revision in Section 1 (from line 49) of the revised version of the paper (including some minor text changes): "A well-established technique for enhancing the stability of policy optimization involves introducing regularization to the objective function using the Kullback-Leibler (KL) divergence [40, R1, R2]. This regularization mechanism aims at maintaining small changes between successive policy updates. Policies similarity can also be quantified via the Wasserstein distance, as recently proposed in [30, 32, 36]. Unlike the KL divergence, the Wasserstein distance enjoys powerful geometrical, computational, and statistical features [26, 27]. In this paper we exploit such properties in a Wasserstein-regularized objective for GMM-based policies optimization. This allows us not only to see the policy iteration as a Wasserstein gradient flow, as in [30], but also to leverage the Riemannian geometry associated to the GMM space to formulate exact Riemannian gradient descent updates." Moreover, we will provide further details in Appendix A.4 so that the connections among our implicit scheme, the approximation with NGs, and the retraction operation in Riemannian manifods become clearer. We believe that the minor changes proposed above complement and improve the motivation and position of our method w.r.t KL-regularized approaches. [R1] Towards an Understanding of Default Policies in Multitask Policy Optimization. AISTATS, 2022. [R2] Information Asymmetry in KL-regularized RL. ICLR, 2019. 2. We would like to point out that we did not deliberately omit an empirical comparison against NG- or KL-based approaches, and we did not intend our paper to "look suspicious", as this reviewer suggests. In this regard, we would like to highlight that: - Our original benchmark experiments included PPO, which optimizes a KL-regularized objective [44, Eq. 8]. - During our rebuttal, we added a GMM-based formulation of PPO, i.e, PMOE-PPO. - Our ablation study reported in App. A.6.2 is aimed at showing experimental evidence on the difference between our Riemannian approach and the implicit scheme. We would like to emphasize that the natural gradient is an approximation of such an implicit scheme, as discussed in our Appendix and in other related works. Although we did not compare against an NG-based approximation of our method, we still believe that the results reported in our paper and the rebuttal provide empirical evidence of our approach performance against methods that include KL divergence (as soft constraints), and an implicit schema (related to the inexact NG). So, we find the reviewer's argument "the complete omission of NG and the KL looks suspicious" inaccurate and unfair. 3. Note that a method that formulates the Wasserstein NG for multivariate GMMs in policy optimization has not been proposed yet (to the best of our knowledge). As pointed out in our previous answer, Moskovitz et al [36] work is the only work where Wasserstein NG is used in policy optimization, but this does not propose a formulation for GMM policies. Therefore, in order to carry out such comparison we would need to develop a new policy optimization approach that builds on the use of the Wasserstein NG developed in [R3], and that employs similar numerical estimates as proposed in [36], perhaps independently for each component of the GMM. Still, such an approach would disregard the positive-definite constraint of the Riemannian manifold arising from the covariance matrices (for which approximated projections would be then necessary). As a conclusion, we agree with the reviewer that such an approach would be relevant, but it is non-existent in the current literature. We believe this could be an interesting future research direction, although it remains an approximation of our mathematically rigorous method. [R3] Optimal transport natural gradient for statistical manifolds with continuous sample space, Information Geometry, 2020.

Authorsrebuttal2023-08-11

GMM/GMR policies

We thank the reviewer for having read our rebuttal and for the timely feedback. Below you can find our reply. We would like to point out that the despite the GMM conditional distribution results in a "linear" combination of Gaussians, this does not mean that the resulting policy is unable to encode highly-nonlinear motion skills. This mainly depends on the GMM parameterization (e.g., using full covariance matrices, the number of Gaussian components). Therefore, we do not see this as a weakness, and the wide use of GMMs for robot motion policies in a plethora of real settings (e.g., [9, 12, 13-25] just to cite a few) is a good indicator of this model's capabilities.

Authorsrebuttal2023-08-11

Further clarifications on questions Q1 and Q2

**Q1:** Note that training robot policies from demonstrations and then adapting them using RL follows a similar and well-known pre-training pipeline where behavior cloning is used to warm-start RL methods. However, this can serve other purposes. For example, our reaching task involved to adapt the initial policy so that the robot reaches a new target midway the original demonstrations. This meant that the RL algorithm did not need to adapt all the GMM components but a subset of them. So, the original policy (learned from demonstrations) is still useful to the resulting policy optimized by RL. **Replay buffer** The reviewer is correct, SAC uses a replay buffer for training and so does the SB3 implementation we used. What we meant with our original answer was that the demonstrations were only used to train the initial GMM policy, and not used later as part of the SAC's replay buffer. We apologize for the confusion that our first answer raised. **Q2:** We implemented precisely the approach proposed in [25], which trains an initial GMM policy, that is later adapted using SAC (as described in Section IV. from [25]). SAC-GMM adapts the underlying GMM by applying residual changes to the GMM parameters (See Eq. (2) in [25]). Note that the performance reported in our paper and the rebuttal correspond to baselines whose hyperparameters (incl. the number of interactions N in SAC-GMM) were tuned using Optuna [57].

Reviewer UKYq2023-08-19

Thanks for getting back to me. Sorry for the delay in my response. I will lump all the responses in this comment. **Natural Gradients.** Why cite two paper two papers from 2019 - 2022 when natural gradients have been used in RL since 2001? I included a list of relevant publications in my initial review. Also reference [40] refers to 'The Wasserstein Space' by Panaretos et al which doesn't look correct. PPO is also not an ideal NG-based baseline. The popular version of PPO (that I think stable baselines implement) uses objective clipping in lieu of KL regularization. While it has a similar regularization effect, I don't think its a fair natural gradient baseline. Another aspect is that PPO was designed for complex neural network policies whereas this work has a more structured policy with fewer parameters. A fair NG baseline would be similar to the moment update equations in Equation 12, but derived using the natural gradient of the KL instead. For example Equations 16 in [A]. I appreciate that Section A.4 covers an ablation in this spirit, but since the Appendix doesn't actually discuss natural gradients explicitly to hard to know exactly how the methods differ, and whether the hyperparameters have been tuned for this alternative update (since ablations aren't usually tuned). [A] Fast and Simple Natural-Gradient Variational Inference with Mixture of Exponential-family Approximations, Lin et al. ICML 2019 **GMR policies.** The paper has experiments with a linear model. The challenge is to demonstrate that the method can train richer policies that are maybe GMR models with network feature spaces. Currently its unclear if the method can be used this way. **Q1.** Is the implication that the GMR policy cannot be trained from random initialization? I am familiar with using behavioral cloning to initialize a policy (for the same task) but I am not familiar with with initializing a policy to solve a different task. **Q2.** I find it strange that these GMM baselines don't use GMMs in the same way that the WGF method does, specifically the passage > 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 = [\Delta\omega, \Delta\mu, \Delta\text{vec}(\Sigma)]$, following the same methodology as SAC-GMM [25 ]. This means they are not fair baselines since the policy implementation is different and the contribution of this paper is the optimizer rather than the policy. It looks like the PMOE-* baselines suggested by another review are better baselines as they are closer to the GMR policy. However, these policies parameterize GMMs with neural networks. Can the authors clarify if the PMOE-* baselines use the same GMR policy class with demonstration pre-training as the WGF implementaiton, or the NN-parameterized GMM policies? I'm not sure it's fair to have baselines with more complex policy classes, since it will take longer to train all the parameters and converge and this may appear like poor policy optimization.

Authorsrebuttal2023-08-21

Reply to reviewer UKYq

Thank you very much for your valuable comments and time for discussing our paper, we truly appreciate your time and engagement in the rebuttal phase. **Natural gradients revision**: We apologize as we introduced a typo in our suggested amendments for Section 1. This reviewer is correct, the reference [40] was not the intended inclusion. Our intention was to incorporate reference [44] instead, and we are grateful for the reviewer's vigilance in highlighting this discrepancy. Moreover, we are committed to extend the proposed revision to comprehensively integrate the early KL divergence- and NG-based works as suggested by this reviewer. **PPO**: Note that the implementation provided by Stable Baselines considers both regularization methods: the clipped objective and a KL penalty. In fact, the PPO documentation provided by Stable Baselines advises to *"limit the KL divergence between updates, because the clipping is not enough to prevent large updates"*. In light of this guidance, our benchmarked PPO did include the KL regularization. For completeness, we will add the "target_kl" values found via Optuna for all the experiments reported in our paper. On a different note, we would like to clarify that *all* our experiments (including the ablation) used Optuna to tune the hyperparameters. **Comparison against NG-based approaches**: We truly appreciate the insightful suggestions made by this reviewer regarding the connections between regularized policy optimization and natural gradients. As acknowledged in this reply above, we are committed to expanding upon our revision in Sec 1. Nonetheless, it is pertinent to highlight that the experimental comparisons advocated by the reviewer demand the conception of two novel algorithms. To the best of our knowledge, no such formulations have been introduced within the corpus of policy optimization literature. Let us emphasize the novel nature of these proposed comparisons: (1) although Eqs. 16 in [A] may be used in combination with a natural gradient approximation w.r.t the KL divergence for GMMs, this approach for GMM policy optimization remains new; (2) an approach that uses Wasserstein natural gradients for GMM policy optimization would be also a brand new method. We agree that such NG alternatives are interesting, but we would like to emphasize that these methods do not exist and their development may even deserve a separate publication dedicated to natural gradient methods for GMM policy optimization. Finally, on a theoretical front, NG would remain an approximation of our mathematically rigorous method because: (1) numerical estimates for the Fisher information matrix are known to break the invariance-to-parametrization properties, and (2) it approximates the exact Riemannian gradient descent in Eq. (40), thus it does not guarantee that the updates stay on the underlying manifold. **GMM/GMR policies**: We believe that we have shown the potential of our novel Wasserstein policy optimization for GMMs in different experimental settings: three distinct robotic tasks with planar robots and two simulated robotic settings with GMMs trained on 6D and 14D state-actions spaces (covering both task- and joint-space policies). We believe that applications that combine our method with perception backbones that output state embeddings is an interesting extension of our paper, and we will discuss this as future work in Sec 5. **Q1**: Our approach also works with a randomly-initiliazed GMR policy without any change in the mathematical formulation or the optimization loop. However, note that initializing robot policies via imitation learning may be exploited by the subsequent policy optimization to avoid updating a subset of the GMM components (as described in Sec. 4.2, for the reaching task). **Q2**: Note that plenty of RL works use state of the art methods such as PPO or SAC overlooking the specific structure of the policy that is being optimized. Therefore, our reasoning to compare against vanilla PPO followed what a naive user may do: use an out-of-the-box RL algorithm such as PPO and adapt it to the policy of interest. Such methodology is actually very similar to what was proposed by SAC-GMM, therefore we believe such baselines represent natural implementations of standard RL methods for a specific policy structure, in our case, GMM-based policies. Nevertheless, a very important outcome from our discussion with the reviewers is the addition of two new baselines (PMOE-PPO and PMOE-SAC), which were indeed designed for GMM-based policies. Regarding the policy representation for PMOE-PPO and PMOE-SAC, we used shallow NNs so that a single layer represented the GMM parameters. We agree with the reviewer that using deep NNs would have impacted negatively the PMOE convergence rates, that is the reason why our implementation built on shallow architectures for a fair comparison. We will include a table of the PMOE hyparameters in the appendix for completeness.

Reviewer zWwA2023-08-13

Re : Rebuttal and further experiments

Thanks for these clarifications. Having run additional experiments is appreciated and the performance uplift of using the WGF method is now more tangible. Based on these efforts and results I am indeed raising my score.

Authorsrebuttal2023-08-15

Thank you!

We truly appreciate your positive feedback on the additional work we made to improve our paper according to your suggestions. Once more, we are very grateful for your time and feedback provided during the review/rebuttal processes.

Authorsrebuttal2023-08-21

Thanks!

The discussion period is near-finished and therefore we do not expect to be able to engage in further discussions. We wanted to express our gratitude for the helpful reviews that improved our paper on a multitude of fronts. The bulk of these changes are summarized in our "Rebuttal discussion" and some further clarifications have been made in our reply to reviewer UKYq. As far as we can tell, these changes should have addressed all reviewer concerns. We are grateful for the improvements that our discussion led to, the positive feedback from the reviewers regarding our new experiments and rebuttal's feedback, and for the interesting ideas for future work.

Program Chairsdecision2023-09-21

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC