Title: DIFFUSING POLICIES : TOWARDS WASSERSTEIN POLICY GRADIENT FLOWS

Abstract: 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.

1. INTRODUCTION AND SETTING:

Deep reinforcement learning algorithms have enjoyed tremendous practical success at scale Mnih et al. (2015; 2016) . Separately, theoretical and practical success through smoothing has also been achieved by generative adversarial networks with the introduction of Wasserstein GANs Arjovsky et al. (2017) . 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 Kullback-Leibler divergence. Simultaneously, gradient flows relax and generalize to continuous time the notion of gradient steps. An important mathematical result due to Jordan et al. (1998) 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 Williams & Peng. (1991) ; Mnih et al. (2016) 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:

• Deterministic rewards It is already known Sabes & Jordan. (1996) that entropic regularization of policy gradients leads to a limit energy-based policy that probabilistically matches the rewards distribution. We investigate the dynamics of that convergence. Our contributions are as follows:

1. We interpret the mathematical concept of gradient flow as a continuous-time version of policy iteration.

2. We show that the choice of a Wasserstein-2 trust region in such a setting leads to solving the Fokker-Planck equation in (infinite dimensional) policy space, leading to the concept of policy transport. This shows optimal policies are arrived at via diffusion and advection. This also justifies empirical practices such as adding Gaussian noise to gradients.

2. GRADIENT FLOWS:

2.1 ENTROPY-REGULARISED REWARDS Let r(a) be the reward obtained by taking action a. The expected reward with respect to a policy π is:

EQUATION

Shannon entropy is often added as a regularization term to improve exploration and avoid early convergence to suboptimal policies. This gives us the entropy-regularised reward, which is a free energy functional, named by analogy with a similar quantity in statistical mechanics1 :

J(π) = A r(a)dπ(a) -β A log π(a)dπ(a) = K r (π) -βH(π) (2) 2.2 POLICY ITERATION AS GRADIENT FLOW

We are interested in the process of policy iteration, that is, finding a sequence of policies (π n ) converging towards the optimal policy π * . In this section we follow closely the exposition by Santambrogio. (2015) . Policy iteration is often implemented using gradient ascent according to

EQUATION

Rearranging gives the explict Euler method

EQUATION

In this article we are more interested in the implicit Euler method which simply replaces ∇J(

π τ k ) with ∇J(π τ k+1 ) π k+1 -π k τ -∇J(π k+1 ) = 0 (5)

This is a policy iteration method. If integrated and interpreted as an L 2 regularized iterative problem, it is strictly equivalent to finding a solution to the proximal problem:

EQUATION

Rather than just the L 2 distance between policies for constraining and regularization, one can envision the more general case of any policy distance d:

EQUATION

2.3 WASSERSTEIN PROXIMAL MAPPING: POLICY TRANSPORT d can be chosen as any general metric distance or divergence. For instance, should we choose d = √ D KL , the square root of the Kullback-Leibler divergence, in the proximal mapping above, we'd get a policy iteration procedure close in spirit to Schulman's trust region policy optimization Schulman et al. (2015a.) .

The case of interest in this paper is when d = W 2 is the second Wasserstein distance, that gives the optimal transport cost for cost function c(x, y) = 1 2 |x -y| 2 . We therefore do iterative minimization in the Wasserstein-2 space W 2 :

EQUATION

A gradient flow is obtained in the small step limit τ → 0. This proximal mapping is an example of Moreau envelope, and remains in a convex optimization setting when J is convex. The Wasserstein distance W 2 is defined for pairs of measures (µ, ν) seen as marginals of a coupling γ by:

∀(µ, ν) ∈ P 2 , W 2 2 (µ, ν) = inf γ∈Γ(µ,ν) |x -y| 2 dγ(x, y) = inf X∼µ,Y ∼ν E|X -Y | 2 (9)

The optimal coupling γ * in the infimum above is also called the optimal transport plan.

Also note that later, we will reformulate this Monge- Kantorovich problem Kantorovich. (1942) as an equivalent linear problem of inner product minimization in L 2 (P 2 ):

EQUATION

3 DERIVING THE FOKKER-PLANCK EQUATION Jordan, Kinderlehrer and Otto showed in a seminal result that several partial differential equations can be interpreted as steepest descent, or gradient flows, of functionals in Wasserstein space W 2 (Jordan et al., 1998) . Here we compute the PDE associated with the entropy-regularized rewards functional J and its steepest descent within W 2 .

3.1. THE EULER CONTINUITY EQUATION:

If we take the limit of small steps size τ in the implicit Euler method of equation 5 we get the Cauchy problem

EQUATION

π(0) = π 0 (12) which describes a gradient flow.

Just like gradient flows are the continuous-time analogue of discrete gradient descent steps, the Euler continuity equation is the continuous-time analogue of the discrete Euler methods seen earlier.

A classic result in Wasserstein space analysis is that because optimal transport acts on probability measures, it must satisfy conservation of mass. Hence all absolutely continuous curves, or flows of measures, (π t ) in W 2 (P) are solutions of the Euler continuity equation. The Euler continuity equation can be seen as the formal continuous-time limit of the Euler implicit method described above

EQUATION

where v t is a suitable vector velocity field (∇• is the divergence operator). In the case we are looking at:

EQUATION

where δJ δπ (π) is the first variation density defined via the Gateaux derivative as δJ δπ

EQUATION

for every perturbation ξ = π -π.

Substituting v t into 13 gives us a partial differential equation necessarily of the form:

EQUATION

This is proven rigorously in Jordan et al. (1998) .

3.2. THE FOKKER-PLANCK EQUATION:

It remains to compute the first variation density δJ δπ (π) for entropy regularised policy gradients. First, the linear part K r , or potential energy has first variation given naturally by

K r (π) = A rdπ ⇒ δK r δπ (π) = r (17)

Second, the entropy part H is a special case H = U t log t of the general internal energy density functional:

EQUATION

In the case of entropy, f (t) = t log t, f (t) = 1 + log t, and therefore δH δπ

EQUATION

)

Finally we require the gradient of this first variation density, given by:

EQUATION

The gradients δ/δπ are functional, whereas the gradients ∇ are action-gradients ∇ a with respect to actions a ∈ A.

Plugging this into 16 gives us the partial differential equation associated with steepest descent within W 2 for entropy-regularized rewards:

∂ t π = -∇ • (π∇r) + β∆π (21) 4 INTERPRETATION 4.1 CONVERGING TO THE OPTIMAL POLICY

The entropy-regularized rewards J is convex in the policy π, which means there is a single optimal policy. The optimal policy will be achieved as long as each step improves the policy and this will be the case as long as the steps taken as are small enough.

Given that we converge to the optimal policy, we know that at optimality ∂ t π = 0. Using equation 16 then gives us a necessary condition for the optimal policy π * :

-∇ • (π∇ δJ δπ (π) = 0 (22)

By setting δJ δπ (π) = 0 and substituting in 20 we get

∇r = β ∇π * π * = β∇ log π * (23)

which gives us the optimal policy π * ∝ e r/β (24) This is the Gibbs measure of the rewards density per action -also seen as an energy-based policy, in line with the static result of Sabes & Jordan. (1996) . The unregularized case β = 0 appears degenerate.

The gradient flow associated with the pure entropy functional βH is the heat equation ∂ t π = β∆π.

Here the Laplacian ∆ is in action space. This is one of the key messages of the derivations we have done in this part: the Wasserstein gradient flow turns the entropy into the Laplacian operator2 .

For the full, entropy-regularized reward J(π), there is an extra transport term generated by the rewards, and the PDE is therefore the Fokker-Planck equation. This means that taking policy gradient ascent steps in W 2 (according to equation 8), is equivalent to solving the Fokker-Planck equation for the policy π with potential field equal to the gradient of rewards r. Alternately, we can say, the policy undergoes diffusion and advection -it concentrates around actions with high reward.

4.3. BROWNIAN MOTION AND NOISY GRADIENTS:

A partial differential equation for diffuse measures also admits a particle interpretation. The result above can also be written, through Ito's lemma Revuz & Yor. (1999) , as the stochastic diffusion version of the Fokker-Planck equation Jordan et al. (1998) :

EQUATION

)

with Π t a finite dimensional discretization of policy π t , and B t a Brownian motion of same dimensionality. In that case, the density of solutions verifies equation 21 -formally, one replaces increments of the Brownian motion dB t by its infinitesimal generator, the Laplacian 1 2 ∆. This stochastic differential equation can also be seen as a Cauchy problem of on-policy rewards maximization, this time with added isotropic Gaussian policy noise √ 2βdB t .

Discretizing again -but in the time variable rather than the action space -one writes, with an explicit method:

EQUATION

This is just noisy stochastic action-gradients ascent on the rewards field. This shows a link between entropy-regularization and noisy gradients. It also suggests a possible technique for implementation.

The key issue to overcome is to generate gradient noise in parameter space that is equivalent to isotropic Gaussian policy noise.

4.4. RELATIONSHIP BETWEEN WASSERSTEIN AND KULLBACK-LEIBLER DISTANCES:

We note the Kullback-Leibler and Wasserstein-2 problems are related by the Talagrand pinequalities Gozlan & Léonard. (2010) , which for p = 2 and under some conditions, ensure that for some constant C and given a suitable reference measure ν, ∀µ ∈ P(X), W 2 2 (µ, ν) ≤ C • D KL (µ, ν). This justifies the square root in d = √ D KL in the proximal mappings studied earlier (equation 8), but more would be beyond the scope of this article.

4.5. ENERGY-BASED POLICIES AS GIBBS MEASURES:

Since the first variation process δJ(π) δπ is known, and we derive it explicitly earlier, we can use a variational argument specific to W 2 (and invalid in W 1 !). We admit (Santambrogio. (2015) ) that the solution of the minimization problem has to satisfy:

- δJ δπ (π) + φ W2 τ = constant (27)

with φ W2 a Kantorovich potential function for transport cost c(x, y) = 1 2 |x -y| 2 . One useful way to think of the Kantorovich potential is that it is a function, whose gradient field generates the difference between the optimal transport plan T and the identity, according to the equation T (x) = x -∇φ(x).

It is well known Revuz & Yor. (1999) that the Gibbs distribution is the invariant measure of the stochastic differential equation above. Therefore we expect it to play a role of primary importance.

In fact, the solution of the W 2 gradient flow with discrete steps is known explicitly Santambrogio. (2015) if we know the successive Kantorovich transport potentials associated: 

In practice, deriving the W 2 optimal transport as well as its cost at each gradient flow step is numerically instable and computationally expensive. Furthermore, numerical estimators for the gradients of Wasserstein distances have been found to be biased; alternatives such as the Cramer distance behave better in practice but not in theory Bellemare et al. (2017b) . (To our knowledge, the gradient flow of the Cramer distance is not known, and no results exist that relate it to the entropy and Fisher information functionals). In appendix, we use fast approximate algorithms in their small parameter regime. We show that another Gibbs measure, the two-dimensional kernel e -c/ , where c is the ground transport cost, and an auxiliary regularization strength, arises naturally in this context, and leads to taking successive Kullback-Leibler steepest descent steps in the coupling, in a spirit close to trust region policy optimization, but using optimal transport and the Sinkhorn algorithm.

5. RELATED WORK:

Optimal transport, and the study of Wasserstein distances, is a very active research field. Foundational developments are found in Villani's reference opus Villani. (2008) . The gradient flows perspective is presented in Ambrosio's book Ambrosio et al. (2006) for a complete theoretical treatment, and in Santambrogio Santambrogio. (2015) for a more applied view including a presentation of the Jordan-Kinderlehrer-Otto result. A classic reference for connecting Brownian motions and partial differential equations is Revuz-Yor Revuz & Yor. (1999) . Efficient algorithms for regularized optimal transport were first explored by Cuturi Cuturi. (2013) , and then Peyré Peyré (2015) who showed the equivalence to steepest descent of KL with respect to the smoothed Gibbs ground cost, and its formulation as a convex problem. Carlier Carlier et al. (2015) gives proofs of Γ-convergence of W 2, to W 2 . Léonard Léonard. (2014) makes the connection with the Schrodinger problem Schrodinger.

(1931) and concentration of measure.

In the context of neural networks, partial differential equations and convex analysis methods are covered by Chaudhari Chaudhari et al. (2017) . The Monge-Kantorovich duality in the W 1 case, and Wasserstein representation gradients, are applied to generative adversarial networks by Arjovsky Arjovsky et al. (2017) . The W 2 connection with generative models is studied by Bousquet Bousquet et al. (2017) . Similarly, Genevay et al Genevay et al. (2017) define Minimum Kantorovich Estimators in order to formulate a wide array of machine learning problems in a Wasserstein framework.

6. DISCUSSION AND FURTHER WORK:

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 Bellemare et al. (2017a) .

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.

A JUSTIFYING STEEPEST DESCENT OF RELATIVE ENTROPY : ENTROPIC REGULARIZATION OF W 2 ITSELF Here we show that while taking discrete proximal steps according to distance d = W 2 is ill-advised, we can perform optimal transport with respect to the entropic regularization of the Wasserstein distance itself. This is a second layer of entropic regularization.

If we let be a small positive real number, then we define W 2 2, as a regularization of the W 2 minimization in equation 9:

EQUATION

with the entropy H(γ) now extended to two-dimensional coupling space Γ(µ, ν) as, in the discretized case, with cardA = q,

EQUATION

and by an analogue of continuity named Γ-convergence Carlier et al. (2015) ,

W 2 2, (µ, ν) → →0 W 2 2, (µ, ν). In fact, W 2 2, converges exponentially Cominetti & Martin. (1994) fast towards W 2 2 as → 0. W 2 2,

is not a distance anymore, but rather a divergence. Yet it enjoys better numerical properties, and enables closed-form solutions for various problems. We do not have a linear Monge-Kantorovich minimization program anymore, but rather a strictly convex program. With c(x, y) = 1 2 |x -y| 2 :

EQUATION

This change from a linear to a convex problem makes the solution set better conditioned numerically; the solution does not have to lie on a vertex of a convex polytope (by analogy with the simplex algorithm, see Nesterov et al. (1994) ) anymore, and therefore, is more robust to initial conditions. In practice, cannot be chosen too small or these stability properties are lost. A certain amount of smoothing is to be tolerated, which is acceptable in the reinforcement learning context, due to the inherent uncertainty on the rewards distribution. Returning to our optimization problem, moving the inner product bracket inside the H part turns the expression into a single KL divergence. This yields the equivalent problem, as detailed in Peyré ( 2015)

EQUATION

At this stage, the link with earlier sections becomes intuitively very clear, since the reference measure e -c/ is for c W2 (x, y) = 1 2 |x -y| 2 simply the heat kernel e -|x-y| 2 2

. It is therefore not surprising the evolution gradient flows considered earlier were linked to the heat equation. Performing JKO stepping from from

EQUATION

instead of equation 8. Combining both definitions therefore gives the problem

EQUATION

to be solved in 2-d coupling space. With this entropic smoothing, we can now re-cast the optimal transport problem as a Kullback-Leibler problem, trading a single optimal transport proximal step for several, 'fast' KL steps. This is done next section using iterative convex projection algorithms.

B DERIVATION OF THE SINKHORN ALGORITHM B.0.1 THE BREGMAN ALGORITHM:

This algorithm is used in convex optimization for iterative projections. The method generalizes the computation of the projection on the intersection of convex sets. Assume we give ourselves a convex function Ψ and that we consider the associated Bregman divergence D Ψ Amari. ( 2016) defined by

EQUATION

We look to minimize this Bregman divergence D Ψ on the intersection of convex sets C = ∩ C i .In our case of interest there are two such sets C 1 and C 2 . For y a given point, or function, we solve for

EQUATION

or equivalently with φ 1 and φ 2 (π) playing the role of indicator barrier functions

EQUATION

In the case where Ψ = H, we get ∇Ψ = log, ∇Ψ * = exp through Legendre transform gradient bijection, and

D Ψ = H Ψ = D KL (•|Ψ).

The Bregman algorithm Bregman. (1967) simply consists in solving the problem 36 by iteratively performing projection on each of the sets C i in a cyclical manner, therefore building the sequence

EQUATION

with [n] the modulo operator ensuring cyclicality of the projections. Therefore, any problem that can be cast under the convex Bregman form 36 can be solved by taking many steepest descent steps. We now proceed to explicit the P DΨ C [n] operators, which in our case are P KL C [n] KL proximal steps, and integrate them into an efficient practical algorithm.

B.0.2 THE SINKHORN ALGORITHM:

We start with the need to minimize the convex form 1 ij p ij log p ij + p ij c ij , subject to marginal constraints that the discretized measure µ is transported by p onto ν. From a matrix perspective this translates into the two following constraints: that the sum in column of P being equal to vector µ, and the sum across lines of P equal to ν:

EQUATION

The matrix set U (µ, ν) is convex. We can form the Lagrangian of this optimization problem in p ij using vector Lagrangian multipliers α, β,

EQUATION

A necessary condition for optimality is then

EQUATION

Hence we have shown that the optimal coupling is a diagonal scaling of the ground cost's Gibbs that makes u tangent to the simplex Frogner et al. (2015) . Under this form, the algorithm is compatible with function approximation, where policy µ is a function of a parameter and reads µ θ . We note that another possibility to create this compatibility would be to unroll a fixed number of iterations of the algorithm, as they are effectively matrix and vector operations, as has already been done with generative adversarial networks and deep Q-networks. We hypothesize that learning with a Wasserstein loss, in a continuous action state setting, will help agents pick actions that are semantically close to the optimum action, therefore increasing policy quality, and reducing the 'unnaturalness' of policy mistakes. It is our hope that a Wasserstein loss, by implying relevant semantic directions in action space, will speed up convergence and training of reinforcement learning agents. In practice, we are still limited by our fundamental assumption that the MDP and the statewise rewards density a → r(s, a) are known. Possibilities such as bootstrapping the rewards density distribution exist, and will be explored practically in further work.

:

In this article we follow the convention of convex analysis and optimal transport, that is, entropy H is taken to be convex, rather than that of information theory with H preceded by a negative sign and concave.

This is part of what's known as Otto calculus.

Figures and Tables:

kernel. With the two positive vectors u, v defined as diag(u) = e -p * ij = u i exp(-c ij / )v j . Once this is derived, the iterative convex projections framework of the Bregman algorithm enables us to derive the full Sinkhorn method.If one recasts the entropic transport problemmin γ∈Γ N i,j=1 c i,j γ i,j + H(γ|(µ i ν j ) ij )(42)

Peer review (iclr2018):
 In this paper the authors studied policy gradient with change of policies limited by a trust region of Wasserstein distance in the multi-armed bandit setting. They show that in the small steps limit, the policy dynamics are governed by the heat equation (Fokker-Planck equation). This theoretical result helps us understand both the convergence property and the probability matching property in policy gradient using concepts in diffusion and advection from the heat equation. To the best of my knowledge, this line of research was dated back to the paper by Jordan et al in 1998, where they showed that the continuous control policy transport follows the Fokker-Planck equation. In general I found this line of research very interesting as it connects the convergence of proximal policy optimization to optimal transport, and I appreciate seeing recent developments on this line of work. In terms of theoretical contributions, I see that this paper contains some novel ideas in connecting gradient flow with Wasserstein distance regularization to the Fokker-Planck equation. Furthermore its interpretation on the Brownian diffusion processes justifies the link between entropy-regularization and noisy gradients (with isotropic Gaussian noise regularization for exploration). I also think this paper is well-written and mathematically sound. While I understand the knowledge of this paper based on standard knowledge in PDE of diffusion processes and Ito calculus, I am not experienced enough in this field to judge whether these contributions are significant enough for a standalone contribution, as the problem setting is limited to multi-armed bandits. My major critic to this paper is its practical value. Besides the proposed Sinkhorn-Knopp based algorithm in the Appendix that finds the optimal policy as fixed point of (44), I am unsure how these results lead to more effective policy gradient algorithms (with lower variance in gradient estimators, or with quasi-monotonic performance improvement etc.). There are also no experiments in this paper (for example to compare the standard policy gradient algorithm with the one that solves the Fokker-Planck equation) to demonstrate the effectiveness of the theoretical findings.