Wasserstein Policy Optimization

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.

One of the most important problems in AI is how to learn to make sequences of good decisions (or 'actions') to optimize a desired objective. Many important applications require continuous actions. Such actions consist of smooth parameters (rather than, say, discrete symbols), each of which needs to be set appropriately and adaptively, and often quite precisely, to get the desired behaviour. Consider, for example, controlling robotic limbs or using magnets to control plasma in a nuclear fusion reactor.We present a new reinforcement learning algorithm, that can learn directly from experience in such settings, and which works especially well in “high-dimensional” problems with many continuous action components. This makes it a promising algorithm for applications such as robotic control or nuclear fusion. However, the algorithm is not specific to these tasks and we show it is effective on several continuous control problems.The new algorithm is derived from a mathematical concept called a "Wasserstein gradient flow". This has inspired earlier algorithms, but a practical and efficient algorithm had not previously been derived directly. The resulting algorithm is intuitive and elegant, and we discuss how it relates to, and improves upon, previous algorithms.