Wasserstein Policy Optimization

We thank the reviewer for their comments and hope to convince them of the merits of our approach.

1. With regards to the performance of WPO relative to other baseline methods: while Figs. 3 and 6 show that WPO is competitive over the DM control suite, our main performance claims are highlighted in Figs. 4 and 5 that show that, as conjectured, WPO especially shines in domains with relatively high-dimensional action spaces. We would also dispute the characterization that “the alternative approaches consistently outperform the proposed WPO approach”. As we state in the paper, no single approach consistently outperforms the others across all tasks. Even on DeepMind Control Suite, there are tasks for which WPO performs the best, such as Cartpole - Three Poles, and many others where it matches state-of-the art. We also think the reviewer may misunderstand the main claim of the paper. We are not claiming that “training with a Wasserstein gradient flow facilitates learning” in all cases. We are claiming that WPO is novel, fills in a much needed gap in the space of actor-critic algorithms (how to generalize DPG-like updates to arbitrary stochastic policies without reparameterization) and, with only minimal effort, can be adapted into a performant deep RL algorithm.
2. In response to the issues raised in “Experimental Designs and Analyses”: we wanted to show that the WPO gradient (Eq. 6) can be dropped into existing deep RL methods with minimal modification, and to do a fair comparison we wanted the baselines to be trained in as similar a fashion as possible. Our results use the same hyperparameters per algorithm across all tasks. Many deep RL algorithms are highly sensitive to small details in how they are trained. Thus we may exactly match the absolute best performance of each baseline method. However, where possible, we made sure to compare against numbers in the literature to make sure our comparison is fair - for instance, we both ran the open-source Dopamine implementation of SAC ourselves and added baseline numbers from the literature. Since submission, we implemented SAC in the same framework that was used for the other experiments, and will include those results in the final camera-ready paper. We are also running DDPG with 1-step updates instead of 5-step updates for comparison and can include those results in the appendix, although they do not look appreciably better than what we have already included. In general, while it may be possible to tune the hyperparameters to get better baseline performance, it may also be possible to tune WPO hyperparameters to get better performance, so we believe the comparison provided in the paper is quite fair.
3. Re:TD3 and other possible baseline methods, we are primarily interested in the WPO update itself (Eq. 6), and compare like-for-like with algorithms like MPO. The algorithm we arrive at in this paper which we call “WPO” is one of an enormous number of possible algorithms that use the update in Eq 6. TD3 is an extension of DDPG that has some additional algorithmic ideas like target policy smoothing that can also be combined with other algorithms. It would be lovely to see future work decomposing these separable ideas and combining them in novel ways, to understand the individual contributions of all these ideas separately, rather than just comparing complete agents with somewhat-arbitrary combinations. Maybe WPO with TD3-like tricks will outperform TD3, maybe not, but that is out of scope of this particular paper.
4. When you say “the proposed solution is straightforward in the sense that it does not yield downstream scientific challenges that would contribute to the advancement of the field” - we do not understand if this is a criticism or a compliment. A solution which is straightforward and does not create future challenges is a good thing, is it not?
5. “Why is ∇_θπ(s_t) a Jacobian but not the gradient?” - ∇_θπ(s_t) is only a Jacobian in the case of DPG, because in this case π(s_t) is a deterministic vector-valued mapping from state vector to action vector. In the case of WPO or standard policy gradient, π(a_t|s_t) is a scalar-valued probability density and thus ∇_alogπ(a_t|s_t) is a vector of the same shape as π(s_t) in DPG.
6. The DMC results are based on proprioceptive control, not pixel-based control.

We hope that this resolves any lingering issues with the paper. We want to emphasize again that our goal was not to show that we have the right combination of tricks and hyperparameter settings to achieve optimal performance across all possible tasks. We wanted to show that an open question in continuous control - how to train arbitrary stochastic policies with updates that exploit the action-value gradient - has a clear and elegant answer using approximations to Wasserstein gradient flows, and make a first foray into translating that into a practical deep RL algorithm. We hope you will agree that that is a significant contribution worth publishing at ICML.

We thank the reviewer for the comments and helpful suggestions.

1. With regards to the state-occupancy term $d^\pi(s)$ in the definition of the gradient of the value function, we agree that including a correction in the off-policy case could potentially improve performance, but want to emphasize that virtually no methods in off-policy RL actually use a correction of this form in practice. Therefore, our algorithm conforms to standard practice and we are confident that the experimental comparisons in the paper are fair. One reason that this term is almost always excluded is that estimating it in continuous state spaces is extremely challenging. It’s also worth highlighting that it has been argued in previous work that excluding this term actually can improve the efficiency of the updates, and it has been shown to be theoretically a sound thing to do (see, e.g., https://arxiv.org/pdf/2302.11381). We will discuss this explicitly in the paper.
2. We have fixed the typo in the equation on line 685
3. Unfortunately, we do not control the rights to the fusion simulator used in our experiments, and so have no control over open-sourcing it. However, the organization which does own the rights intends to release an open-source version of the simulation code in the long term. We will do our best to encourage them to release this as soon as they can.
4. The use of the reparameterization trick becomes much more complicated with distributions with discrete latent variables, like mixtures of Gaussians. The exact SVG(0) update, and the efficiency of this, will be sensitive to the choice of reparameterization. It is difficult to find a practical problem for which mixture-of-Gaussian action distributions are necessary, since Gaussian action distributions work well in practice for most environments investigated here. However, we will extend Section 4 to include a toy example of a mixture of Gaussians policy on a complex loss landscape showing the qualitative difference in dynamics between WPO and standard policy gradient, and Gaussian vs. mixture-of-Gaussian action distributions.
5. The discussion of the variance of the linear action-value function case in Section 4 can be improved. First, we will clarify that the WPO gradient is only constant when the variance of the policy is fixed (e.g. only the mean is updated). Second, we will clarify that it is only standard policy gradient for which the variance is finite, while SVG(0) with the appropriate parameterization is identical to WPO, as you noted. Third, we will highlight that the agreement between WPO and SVG(0) is only true for specific choices of parameterization for SVG(0), and provide a concrete example where this identity is broken. For instance, replacing a Gaussian action distribution with an exponential distribution π(a) = exp(a/β)/β (with the natural reparameterization a = β * η) leads to different updates for WPO and SVG(0), even in expectation: the expected natural WPO update with respect to the scale parameters β will then simply be E[ ∇_a Q(a) ], while the SVG(0) update will be E[ ∇_a Q(a) diag(a/β) ], with η a standard exponential with p(η = x) = exp(-x). (If we consider the gradient with respect to the network parameters, both these updates just get multiplied with the Jacobian ∇_θ β, as per standard backprop.)
6. We agree that using WPO in combination with diffusion models is a promising future avenue to investigate. Thank you for the suggestion.

We would like to thank the reviewer for their thoughtful comments and favorable review. The reviewers made a few suggestions for improving the paper which we will address:

1. We have re-run experiments on the high-variance environments with more seeds to smooth out the learning curves and will add the new results to the paper. The new results look qualitatively similar to those in the original submission, with some amount of noise.
2. We appreciate the suggestion for more challenging high-dimensional tasks to investigate such as Bi-DexHands. The state-of-the-art RL methods for bidextrous manipulation often involve many additional steps beyond blank-slate actor-critic methods, such as imitation learning from human data (e.g. ALOHA Unleashed, https://arxiv.org/pdf/2410.13126). This adds an additional layer of complexity for comparison against baselines that goes beyond the scope of what we have the space and time to investigate in this paper, however we will be sure to investigate these sorts of tasks going forward.

With regards to the lack of uniformly superior performance that the reviewer points out, we refer to the rebuttal to reviewer A8Ro for a longer discussion. We hope that this addresses all of the issues the reviewer has and thank them again.

Thank you for your helpful comments. We appreciate the overall positive evaluation of the paper, and wanted to address the two points raised in “Claims and Evidence”:

1. To go from the idealized WPO update to the practical algorithm presented in the paper, two approximations were made: the Fisher information matrix was approximated by the FIM for a diagonal Gaussian distribution, and an entropy regularization term was added.

First, there is no reason that the “bare” WPO update (i.e. without the Fisher term) couldn’t be used for general distributions over actions. As noted in the paper, the variance in the update would grow as the policy becomes more deterministic, but with gradient clipping this might be manageable. Beyond that, there are many classes of distributions for which the exact Fisher information matrix is known (e.g. exponential families) and, going even beyond that, numerous practical approximations to natural gradient descent exist which could be applied here, e.g. KFAC (Grosse and Martens 2015). Ultimately, there are going to be a spectrum of approaches with the exact update at one end and the “bare” update at the other which can be explored in future work.

Second, regularization of the policy update through penalties or trust region methods is a standard technique, and the approach used here could be extended to more complex distributions over action spaces so long as an appropriate way to estimate the KL divergence between action distributions exists. For instance, a sampling-based approximation could be used for generic distributions, or an upper bound or other approximation to the KL could be used in cases where a closed form solution does not exist. In practice, the exact measure of divergence between action distributions is probably not that important, as long as it is simple to compute and is well behaved.

2. Understanding the precise cause of WPO underperforming other algorithms in certain environments is challenging. Note that we only tried tuning a small number of parameters, primarily the weights of the KL regularization, while other methods have benefited from years of extensive hyperparameter tuning on many of these environments to achieve optimal performance. Note also that, while WPO is not always the top performing algorithm, it does not fail dramatically on any environments in the way that, say, DDPG does on Cartpole - Swingup Sparse, which illustrates the overall robustness of WPO. We agree that it would be quite valuable to investigate the cases where WPO underperforms in more detail, but given the large number of moving parts and possible causes, we feel this is best left to future work. Furthermore, while Figs. 3 and 6 show that WPO is competitive over the DM control suite, our main performance claims are highlighted in Figs. 4 and 5 that show that, as conjectured, WPO especially shines in domains with relatively high-dimensional action spaces.

We also appreciate the suggestion to compare against Zhang et al (2018). The algorithm implemented in Zhang (2018) is a combination of a DDPG-like gradient with the reparameterization trick, and thus closely resembles SVG(0) and SAC, albeit with a slightly different regularization. The significance of WPO lies mainly in its generality, going beyond methods like SVG(0), SAC and Zhang (2018) that rely on reparameterization. Additionally, we believe that by including SAC and DDPG in the experiments already, we have good coverage of methods with the same basic form of the policy gradient update. If we investigate the effect of regularization in more detail in future work we will be sure to include an analysis of Zhang (2018). Once again, we are pleased that you liked the paper, and appreciate all of your helpful comments.