Policy gradients methods often achieve better performance when the change in policy is limited to a small Kullback-Leibler divergence. We derive policy gradients where the change in policy is limited to a small Wasserstein distance (or trust region). This is done in the discrete and continuous multi-armed bandit settings with entropy regularisation. We show that in the small steps limit with respect to the Wasserstein distance $W_2$, policy dynamics are governed by the heat equation, following the Jordan-Kinderlehrer-Otto result. This means that policies undergo diffusion and advection, concentrating near actions with high reward. This helps elucidate the nature of convergence in the probability matching setup, and provides justification for empirical practices such as Gaussian policy priors and additive gradient noise. Deep reinforcement learning algorithms have enjoyed tremendous practical success at scale BID17 BID0 . Separately, theoretical and practical success through smoothing has also been achieved by generative adversarial networks with the introduction of Wasserstein GANs BID2 . In both instances, a smooth relaxation of the original problem has been key to further theoretical understanding. In this work, we take the view of policy gradients iteration through the lens of converging towards a function of the rewards field r(s, a) for a given state s. This view uses optimal transport metrized by the second Wasserstein distance rather than the standard KullbackLeibler divergence. Simultaneously, gradient flows relax and generalize to continuous time the notion of gradient steps. An important mathematical result due to BID14 shows that in that setting, continuous control policy transport is smooth; this achieved by the heat flow following the Fokker-Planck equation, which also admits a stochastic diffusion representation, and sheds light on qualitative convergence towards the optimal policy. This is to our knowledge the first time that the connection between variational optimal transport and reinforcement learning is made.Policy gradient methods BID28 ; BID18 look to directly maximize the functional of expected reward under a certain policy π. π(a|s) is the probability of taking action a in state s under policy π. A policy can hence be identified to a probability measure π ∈ P, the space of all policies. In what follows, functionals are applications from P → R. Out of a desire for simplification, we focus on formal derivations, and skip over regularity and integrability questions.We investigate policy gradients with entropy regularisation in the following setting:• Bandits, or reinforcement learning with 1-step returns• Continuous action space We have used tools of quadratic optimal transport in order to provide a theoretical framework for entropy-regularized reinforcement learning, under the strongly restrictive assumption of maximising one-step returns. There, we equate policy gradient ascent in Wasserstein trust regions with the heat equation using the JKO result. We show advection and diffusion of policies towards the optimal policy. This optimal policy is the Gibbs measure of rewards, and is also the stationary distribution of the heat PDE. Recast as a stochastic Brownian diffusion, this helps explain recent methods used empirically by practitioners -in particular it sheds some light on the success of noisy gradient methods. It also provides a speculative mechanism besides the central limit theorem for why Gaussian distributions seem to arise in practice in distributional reinforcement learning BID3 .Our contribution largely consists in highlighting the connection between the functional of reinforcement learning and these mathematical methods inspired by statistical thermodynamics, in particular the Jordan-Kinderlehrer-Otto result. While we have aimed to keep proofs in this paper as simple and intuitive as possible, an extension to the n-step returns (multi-step) case is the most urgent and obvious line of further research. Finally , exploring efficient numerical methods for heat equation flows compatible with function approximation, are directions that will also be considered in future research.
Linking Wasserstein-trust region entropic policy gradients, and the heat equation. The paper explores the connections between reinforcement learning and the theory of quadratic optimal transport The authors studied policy gradient with change of policies limited by a trust region of Wasserstein distance in the multi-armed bandit setting, showing that in the small steps limit, the policy dynamics are governed by the heat equation (Fokker-Planck equation).