Wasserstein Policy Optimization

TL;DR

Wasserstein Policy Optimization (WPO) introduces a new actor-critic reinforcement learning algorithm for continuous action spaces by deriving a policy update from Wasserstein gradient flows and projecting it onto a parametric policy class. The resulting update, $\theta_{t+1} = \theta_t + \mathcal{F}_{\theta\theta}^{-1} \mathbb{E}_{\pi}[\nabla_\theta \nabla_{\mathbf{a}} \log \pi(\mathbf{a}|\mathbf{s}) \nabla_{\mathbf{a}} Q^{\pi}(\mathbf{s},\mathbf{a})]$, combines action-gradient information with a parametric Fisher preconditioning, and is made practical via a diagonal Fisher approximation, Gaussian policies, and KL regularization to stabilize learning. Empirically, WPO is robust across the DeepMind Control Suite, scales effectively to very high-dimensional action spaces, and performs competitively with, or surpasses, strong baselines like MPO, DDPG, and SAC, including challenging fusion control tasks. The work broadens the toolkit for continuous-control RL by offering a general, off-policy-friendly framework that preserves the benefits of both stochastic and deterministic policy gradients, with promising avenues for non-Gaussian policies and alternative Wasserstein formulations.

Abstract

We introduce Wasserstein Policy Optimization (WPO), an actor-critic algorithm for reinforcement learning in continuous action spaces. WPO can be derived as an approximation to Wasserstein gradient flow over the space of all policies projected into a finite-dimensional parameter space (e.g., the weights of a neural network), leading to a simple and completely general closed-form update. The resulting algorithm combines many properties of deterministic and classic policy gradient methods. Like deterministic policy gradients, it exploits knowledge of the gradient of the action-value function with respect to the action. Like classic policy gradients, it can be applied to stochastic policies with arbitrary distributions over actions -- without using the reparameterization trick. We show results on the DeepMind Control Suite and a magnetic confinement fusion task which compare favorably with state-of-the-art continuous control methods.

Figures (12)

• Figure 1: Conceptual illustration of how WPO combines elements of stochastic and deterministic policy gradient methods. Left: "classic" policy gradient. Samples are taken from a stochastic policy. Each sample contributes a scalar $Q^\pi(\mathbf{s},\mathbf{a})$ factor to the gradient. Middle: deterministic policy gradient (DPG). A deterministic action is chosen and the policy gradient depends on the gradient of $Q^\pi(\mathbf{s},\mathbf{a})$. Right: Wasserstein policy optimization (WPO). Samples are taken from a stochastic policy, as in classic policy gradient, but depend on the gradient of $Q^\pi$ with respect to the action, as in DPG.
• Figure 2: Concrete WPO updates for a single-variate normal policy for two different action-value functions. In the left and middle plots we consider $Q(a) = -a^2/2$, with an obvious optimum at $a=0$, and a policy with $\mu = \sigma = 1$. The left plot shows the gradient on the mean at several sampled actions. These are averaged to produce the update, moving the mean towards the optimal action. The expected WPO update, as shown in the middle plot, is then $\Delta\mu = - \mu$ and $\Delta\sigma = -\sigma$. Both the mean and variance will decrease (or, more generally, move probability mass to the optimal action), as shown. Conversely, in the right plot we consider $Q(a) = a^2 / 2$ then $\Delta\mu = \mu$ and $\Delta\sigma = \sigma$, and both the mean and variance will increase, as shown. In all cases, the expected value of the resultant policy is increased.
• Figure 3: Concrete WPO learning for a one dimensional mixture of Gaussians policy for the non-concave action-value function $Q(a) = -\frac{1}{100}a^4+a^2$. The left plot shows the action-value function and mixture of Gaussians policy. In the middle plot we show the evolution of the policy under a standard policy gradient update, both with samples from the policy and the change in the means of each mixture component. On the left we show the same evolution for WPO. WPO converges faster, is more stable around the optimum, and converges to both optima if the policy is initialized symmetrically.
• Figure 4: Results from selected DeepMind Control Suite tasks. Full results are in Fig. \ref{['fig:control_suite_all']} in the appendix.
• Figure 5: Plots of reward from various agents on combined Humanoid Stand environments. Left to right: 1, 3 and 5 replicated environments (21, 65 and 105 action dimension). Solid line denotes the mean and the shaded region highlights the minimum and maximum over 5 seeds. As the number of replicas grows, WPO is able to learn faster than other methods by a larger margin.