Efficient Wasserstein Natural Gradients for Reinforcement Learning
Authors: Ted Moskovitz
Michael Arbel
Ferenc Huszar
Arthur Gretton
Venue: ICLR 2021 Poster (2021)
Keywords: reinforcement learning
optimization

Abstract: A novel optimization approach is proposed for application to policy gradient methods and evolution strategies for reinforcement learning (RL). The procedure uses a computationally efficient \emph{Wasserstein natural gradient} (WNG) descent that takes advantage of the geometry induced by a Wasserstein penalty to speed optimization. This method follows the recent theme in RL of including divergence penalties in the objective to establish trust regions. Experiments on challenging tasks demonstrate improvements in both computational cost and performance over advanced baselines. 


Published as a conference paper at ICLR 2021
EFFICIENT WASSERSTEIN NATURAL GRADIENTS FOR REINFORCEMENT LEARNING
Ted Moskovitz" Arbel", Ferene H & *Arthur Gretton' Gatsby Unit, UCL Uniiversity of Cambridge
ABSTRACT
A novel optimization approach is proposed for application to policy gradient meth- ods and evolution strategies for reinforcement learning (RL). The procedure uses a computationally efficient Wasserstein natural gradient (WNG) descent that takes advantage of the geometry induced by a Wasserstein penalty to speed optimiza- tion. This method follows the recent theme in RL of including a divergence penalty in the objective to establish a trust region. Experiments on challenging tasks demonstrate improvements in both computational cost and performance over advanced baselines.
INTRODUCTION
Defining efficient optimization algorithms for reinforcement learning (RL) that are able to leverage a meaningful measure of similarity between policies a longstanding and challenging problem (Lee & Popović, 2010; Meyerson et al., 2016; Conti et al., 2018b). Many such works rely on similarity measures such as the Kullback-Leibler (KI.) divergence (Kullback & Leibler, 1951) to define procedures for updating the policy of an agent it interacts with the environment. These generally motivated by the need to maintain a small variation in the KL between successive updates in off-policy context to control the variance of the importance weights used in fthe estimation of the gradient. This includes work by Kakade (2002) and Schulman et al. (2015), who propose to use the Fisher Natural Gradient (Amari, 1997) way update policies, using local geometric information to allow larger steps in directions where policies vary less; and the work of Schulman et al. (2017), which relies on global measure of proximity using a soft KL penalty to the objective. While those methods achieve impressive performance, and the choice of the KL is well well-motivated, one can still ask if is possible include information about the behavior of policies when measuring similarity, and whether this could lead to more efficient algorithms. Pacchiano et al. (2019) provide first insight into this question, representing policies using behavioral distributionss which incorporate information about the outcome of the policies in the environment. The Wasserstein Distance (WD) (Villani, 2016) between those behavioral distributions is then used as a similarity measure between their corresponding policies. They further propose penalty to the total objective. Hence, like the KL penalty, proximity between policies is measured globally, and does not necessarily exploit the local geometry defined by the behavioral embeddings.
such behavioral similarity as a global soft
In this work, we show that substantial improvements can be achieved by taking into account the local behavior of policies. We introduce new, efficient optimization methods for RL that incorporate the local geometry defined by the behavioral distributions for both policy gradient (PG) and evolution strategies (ES) approaches. Our main contributions are as follows:
1- We leverage recent work in (Li & Montufar, 2a;;;;;; Li, 2018; Li & Zhao, 2019; Chen & Li, 2018) which introduces the notion of the Wasserstein Information Matrix to define a local behavioral similarity measure between policies. This allows us to identify the Wasserstein Natural Gradient (WNG) as key ingredient for optimization methods that rely on the local behavior of policies. To enable efficient estimation of WNG, we build on the recent work of Arbel et al. (2020), and further extend it to cases where the re- ce-paaameterizatoon trick not applicable, but only the score function of the model is available.
Denotes equal contribution. Correspondence: tedegat.sby. vcl ..a...

---

Published as a conference paper at ICLR 2021
2- This allows introduce two novel methods: Wasserstein natural policy gradients (WNPG) and Wasserstein natural evolution strategies (WNES) which use the local behavioral structure of policies through WNG and can be easily incorporated into standard RL optimization routines. When combined in addition with a global behavioral similarity such as a WD penalty, we show substantial improvement using the penalty alone without access to local information. We find that such WNG-based methods are especially useful on tasks in which initial progress is difficult.
3- Finally, we demonstrate, our knowledge, the first in-depth comparative analysis of the FNG and WNG, highlighting a clear interpretable advantage of using WNG over FNG on tasks where the optimal solution is deterministic. This scenario arises frequently in ES and in policy optimization for MDPs (Puterman, 2010). This suggests that WNG could be a powerful tool for this class of problems, especially when reaching accurate solutions quickly is crucial.
In Section 2, we present brief review of policy gradient approaches and the role of divergence measures as regularization penalties. In Section we introduce the WNG and detail its relationship with the FNG and the use of Wasserstein penalties, and in Section 4 we derive practical algorithms for applying the WNG to PG and ES. Section contains our empirical results.
BACKGROUND
Policy Gradient (PG) methods directly parametrize a policy тe, optimizing the parameter 0 using stochastic gradient ascent on the expected total discounted reward R(0). An estimate 9N of the gradient of R0) at 0₂ can be computed by differentiating a surrogate objective Co which often comes in two flavors, depending whether training is on-policy (left) off-policy (right):
To (a(aa)()) A
(()) É log To (acls))))
or (()) = E
(1)
To (ae (e)
The expectation is an empirical average over N trajectories catte-action-rewatte obtained by simulating from the environment using To, The scalar A, is an estimator of the advantage function and can be computed, for instance, using
(s), aj ri ...,, st a7. rT) of
A, = - + 7V( - V(s;)
(2)
where y E 1) 1) is a discount factor and V is the value function often learned as a parametric function via temporal difference learning (Sutton & Barto, 2018). Reusing trajectories can reduce the computational cost at the expense of increased variance of the gradient estimator (Schulman al., 2017). Indeed, performing multiple policy updates while using trajectories from an older policy Te..a means that the current policy To can drift away from the older policy. On the other hand, the objective is obtained as an expectation under Fy for which fresh trajectories are not available. Instead, the objective is estimated using importance sampling (by re-weighting the old trajectories according to importance weights Re// ). When me is too far from Te.ss the importance weight can have large variance. This can lead to drastic degradation of performance if done naïvely (Schulman et al., 2017). KL-based policy optimization (PO) aims at addressing these limitations.
KL-based PO methods ensure that the policy does not change substantially between successive updates, where change is measured by the KL divergence between the resulting action distributions. The general idea is to add either a hard KL constraint, as in TRPO (Schulman et al., 2015), a soft constraint, as in PPO (Schulman et al., 2017), to encourage proximity between policies. In the first case, TRPO recovers the FNG with step-size further adjusted using line-search enforce the hard constraint. The FNG permits larger steps in directions where policy changes the least, thus reducing the number of updates required for optimization. In the second case, the soft constraint leads to an objective of the form:
maximize, L(0) 3 [ (-1s)), (-(-)s)))]-
The KL penalty prevents the updates from deviating too far from the current policy TH thereby controlling the variance of the gradient estimator. This allows making multiple steps with the same simulated trajectories without degradation of performance. While both methods take into account the proximity between policies as measured using the KL, they do not take into account the behavior of such policies in the environment. Exploiting such information can greatly improve performance.
2

---

Published as a conference paper at ICLR 2021
Behivior-Guided Policy Optimization. Motivated by the idea that policies can differ substantially as measured by their KL divergence but still behave similarly in the environment, Pacchiano et al. (2019) recently proposed to use notion of proximity in behavior between policies for PO. Exploit- ing similarity in behavior during optimization allows to take larger steps in directions where policies behave similarly despite having a large KI. divergence. To capture a sense of global behavior, they define a behavioraal embedding map (BEM) & that maps every trajectory T to a behavior variable X = (T) belonging to some embedding space E. The behavior variable X provides simple yet meaningful representation of each the trajectory т. As a random variable, X is distributed according to a distribution qo, called the behavior distribution. Examples of Φ include simply returning the final state of a trajectory (+(r)= 8T) or its concatenated actions (Q(7) = (40.... ......... ar)). Proximity between two policies Te and Tgy is then measured using the Wasserstein distance between their be- havior distributions qe and qør. Although, the KI. could also be used in some cases, the Wasserstein distance has the advantage of being well-defined even for distributions with non-overlapping sup- port, therefore allowing more freedom in choosing the embedding & (see Section 3.1). This leads to a penalized objective that regulates behavioral proximity:
8 maximize; (()) 2 Wa go).
(4)
where / E R is a hyper-parammeter controlling the strength of the regularization. To compute the penalty, Pacchiano et al. (2019) use iterative method from Genevay et al. (2016). This procedure is highly accurate when the Wasserstein distance changes slowly between successive updates, as ensured when / is large. At the same time, larger values for 3 also mean that the policy is updated using smaller steps, which can impede convergence. An optimal trade-off between the rate of con- vergence and the precision of the estimated Wasserstein distance can achieved using an adaptive choice of B as done in the case of PPO Schulman et al. (2017). For finite value of [3, the penalty ac- counts for global proximity in behavior and doesn't 't explicitly exploit the local geometry induced by the BEM, which can further improve convergence. We introduce efficient method that explicitly exploits the local geometry induced by the BEM through the Wasserstein Natural gradient (WNG), leading to gains in performance at reduced computational cost. When global proximity is impor- tant to the task, we show that using the Wasserstein penalty in Equation (4) and optimizing it using the WNG yields more efficient updates, thus converging faster than simply optimizing Equation (4) using standard gradients.
3 THE WASSERSTEIN NATURAL GRADIENT
The Wasserstein natural gradient (WNG) (Li & Montufar, 201 (a;b)) corresponds to the steepest- ascent direction of an objective within a trust region defined by the local behavior of the Wasserstein- 2 distance (W2). The W2 between two nearby densities qo and qo+a can be approximated by computing the average cost of moving every sample X from qe to new sample X' approxi- mately distributed according to 4e+u using an optimal vector field of the form (zf so that
X' = X + (xfu(X) (see Figure 6). Optimality of Vzfa is defined as a trade-off between accu- rately moving mass from go to gota and reducing the transport cost measured by the average squared norm of Vafu
sup VOE [fa()]] u - 2 Emo [IV=1((X)I??
(5)
fu
where the optimization is over a suitable set of smooth real valued functions on E. Hence, the optimal function fo solving Equation (5) defines the optimal vector field V fz)u Proposition makes this intuition more precise and defines the Wasserstein Information Matrix.
Proposition 1 (Adapted from Defintion 3 Li & Zhao (2019)) The second-order Taylor expansion We between two nearby parametric probability distributions qe and q04u is given by
W2 (go, 4w+w)= = u T G(()u + o[I ull²)
(6)
where G(f) is the Wasserstein Information Matrix (WIM), with components in a basis (61,.., Cp)
Gj.((0) = Eng (X) (X)]
(7)
The functions S3 solve Equation (5) with chosen as ej. Moreover, for any given u, the solution fu. to Equation (5) satisfies Vo[[(Va.. = u. G(0)a.
3

---

Published as a conference paper at ICLR 2021
(a) Wy-penalty with varying f
(b) KL-penalty with varying
[c] Loss function
107
WNGD
10a
10"
Fisher NGD
10
10"1
IL-penalty
Wa-penalty
10
WNGD
10
NGD
Wy-penalty Ity
KL-penalty
SGD
SGD
10
10
10
10
10²
10
Iterations
Iterations
[d] Log-dingonal para metrization
(el diagonal parametrization
(0) Iterations up to 2 error vs A
10'
107
10
10a
10a
Derattons
30
10
10"
WNGO Fisher NGD
10
10
We-penalty
10"
10"*
aL-penalty
SGD
10
102
102 10
10?
10a
10 10
10
10
10
Iterations
Iterationn
Figure 1: Different optimization methods using an objective C(0) = Egg [xx] where qo is a v). Here la in bold is the mean vector,
gaussian of 100 dimensions with parameters parameterizes the covariance matrix Σ, which is chosen to be diagonal. Two parameterizations for the covariance matrix are considered: Σn, e (log-diagonal) and En v, (diagonal). {/ is the sum of sine functions over all dimensions. Training is up to 4000 iterations, with À = .9 and B = .1 unless they are varied. In Figure (c), a and μ refer to the std of the first component of the gaussian v EE and f1. More details about the experimental setting are provided in Appendix D.3.
When and qu+u are the behavioral embedding distributions of two policies ma and T the function fu allows transport behavior from a policy mh to a behavior as close as possible Tow with the least cost. We thus refer to fu as the behavioral transport function. The function fu determines how hard it is to change behavior locally from policy To direction u, thus providing a tool to find update directions with either maximal or minimal change in behavior.
Probing all directions in basis (e1,...,ep) Cp) of parameters allows us to construct the WIM G(0) in Equation (7) which summarizes proximity in behavior along all possible directions u using G(O)u = Eg (X)) For an objective C()), such as the expected total reward of policy, the Wasserstein natural gradient (WNG) is then defined as the direction u that locally in- creases C( + u) the most with the least change in behavior as measured by fu. Formally, the WNG is related to the usual Euclidean gradient VgL(O) by
W arg max 2g Tu u TG()u.
(8)
From Equation (8), the WNG can be expressed in closed-form in terms of G(Ø) and g as g" G-! (8)g. Hence, WNG ascent is simply performed using the update equation 0 8₂ + Ag We'll see in Section how to estimate WNG efficiently without storing or explicitly inverting the matrix G. Next, discuss the advantages of using WNG over other methods.
3.1 WHY USE THE WASSERSTEIN NATURAL GRADIENT?
To illustrate the advantages of the WNG, we consider simple setting where the objective is of the form L(0) = Eg [=)( with go being gaussian distribution. The optimal solution in this example a deterministic point mass located at the global optimum x* the function dxx). This situation arises systematically in the context of ES when using gaussian noise distribution with learnable mean and variance. Moreover, the optimal policy of a Markov Decision Processes (MDP) is neces- sarily deterministic (Puterman, 2010). Thus, despite its simplicity, this example allows to obtain closed-fome expressions for all methods while capturing a crucial property in many RL problems (deterministic optimal policies) which, will see, results in differences in performance.
4

---

Published as a conference paper at ICLR 2021
Wasserstein natural gradient Fisher natural gradient While Figure/ 1 (c) shows that both methods seem to reach the same solution, a closer inspection of the loss, as shown in Figure 1 (d) and (e) for two different parameterizations of qo, shows that the FNG is faster at first, then slows down to reach a final error of 10-1. On the other hand, WNG is slower at first then transitions suddenly to error of 10- The optimal solution being deterministic, the variance of the gaussian go needs to shrink to 0. In this case, the KL blows up, while the W2 distance remains finite. As the natural gradient methods are derived from those two divergences (Theorem 2 of Appendix B), they inherit the same behavior. This explains why, unlike the WNG, the FNG doesn't 't achieve the error of 10 Beyond this example, when the policy Ty defined only implicitly using a generative network, as in Tang& & Agrawal (2019), the FNG and KL penalty are ill-defined since Te, and Te J1 might have non-overlapping supports. However, the WNG remains well-defined (see Arbel et al. (2020)) and allows for more flexibility in representing policies, such as with behavioral embeddings.
Wasserstein penalty vs Wasserstein natural gradient The Wasserstein penalty Equation (4) en- courages global proximity between updates q0a- For small values of the penalty parameter 8, the method behaves like standard gradient descent (Figure I (a)). As 3 increases, the penalty encourages more local updates and thus incorporates more information about the local geometry defined by qa- In fact, it recovers the WNG direction (Theorem 2 of Appendix B) albeit with an infinitely small step-size which is detrimental to convergence of the algorithm. To avoid slowing-down, there is an intricate balance between the step-size and penalty 3 that needs to be maintained (Schulman et al., 2017). All of these issues avoided when directly using the WNG, as shown in Figure I (a), which performs the best and tolerates the widest range of step-sizes Figure (f). Moreover, when using the log-diagonal parameterization in Figure (d.a), the WNGD (in red) achieves an error of le-8, while W2-penalty achieves a larger error of order le-0 for various values of the 3. When using the diagonal parameterization instead, as shown in Figure (e), both methods achieve similar error of lc-6. This discrepancy in performance highlights the robustness of WNG to parameterization of the model.
Combining WNG and Wasserstein penalty. The global proximity encouraged by a W2 penalty can be useful on its own, for instance, to explicitly guarantee policy improvement as in (Pacchiano et al., 2019, Theorem 1) However, this requires estimating the W22 at every iteration, which can be costly. Using WNG instead of the usual gradient can yield more efficient updates, thus reducing the number of time W₂ needs to be estimated. The speed-up can be understood as performing second- order optimization on the W2 penalty since the WNG arises precisely from a second- -odder expansion of the W2 distance, shown in Section 3 (See also Example in Arbel et al. (2020)).
4 POLICY OPTIMIZATION USING BEHAVIORAL GEOMETRY
We now present practical algorithms to exploit the behavioral geometry induced by the embeddings 4. We begin by describing how to efficiently estimate the WNG.
Efficient estimation of the WNG can be performed using kernel methods, as shown in Arbel et al. (2020) in the case where the re-parametrization trick is applicable. This is the case, if for instance, the behavioral variable is the concatenation of actions X = [ag...., arl and if actions are sampled from gaussian with mean and variance parameterized by neural network, as is often done in practice for real-valued actions. Then X can be expressed as X = (2) where Bg is a known function and Z is input sample consisting in the concatenation of states {so, ...; 8T] and the gaussian noise used to generate the actions. However, the proposed algorithm is not readily applicable if for instance the behavioral variable X is a function of the reward.
We now introduce a procedure that extends the previous method to more general cases, including those where only the score V, log go is available without an explicit re-parametrization trick. The core idea is to approximate the functions fe, defining G(fx) in Equation (7) using linear combi- nations of user-fee basis functions (h1 (1), hz))):
fe, (x)
Σ o (x), m=1
The number M controls the computational cost of the estimation and is typically chosen on the order of M = 10. The basis can be chosen to be data- dependennn using kernel methods. More precisely,
5

---

Published as a conference paper at ICLR 2021
Algorithm 1: Wasserstein Natural Policy Gradient
1: Input Initial policy Teo
2: for iteration & = 1,2,... do
3: Obtain N rollouts {T) of length T using policy no.
4:
Compute loss [(0)) in forward pass
5:
Compute gradient d in the backward pass on [(0x)
6:
Compute Behavioral embeddings (Xn ([
7:
Compute WNG 0% using Algorithm with samples {Xn) J and gradient estimate Ôk-
8: Update policy using: 0x+1 = 02 Ag
9:
end for
we use the same approach as in Arbel et al. (2020), where we first subsample M data-points Ym from batch of N variables Xn, and M indices im from {1,..., d} where d is the dimension of Xx. Then, each basis can of the form (x) = 8r K(Ym; where K a positive semi-definite kernel, such as the gaussian kernel K ((2,)) y) exp( -|ا= This choice of basis allows us to provide guarantees for functions fj in terms of the batch size N and the number of basis points M (Arbel al., 2020, Theorem 7). Plugging-in each S3 in the transport cost problem Equation (5) yields quadratic problem of dimension M in the coefficients a²:
maximize,s 2J 2J
La
where is a square matrix of size M x M independent of the index j and is a Jacobian matrix of shape M p with rows given by Jm.. depending on the applicability of the re-parametrization trick or the availability of the score
FglE [hm (X)]. There two expressions for J,
Jm.. PBgg / (X)V (Z)
J Jm.. F Vo log go (X)hm (X)]
(10)
Computing J can be done efficiently for moderate size M by first computing surrogate vector of V of size M whose Jacobian recovers using automatic differentiation software:
V Rago [l (X,)).
Vm 1e [log qo (x,) hmm (X,))
(11)
The optimal coefficients a are then simply expressed as a = LtJ. Plugging-in in the optimal func- tions in the expression of the Wasserstein Information Matrix (Equation (7)), yields low rank approximation of G of the form G = L'J. By adding a small diagonal perturbation matrix el, it is possible efficiently compute (G + el)- using generalized Woodbury matrix identity which yields an estimator for the Wasserstein Natural gradient
6W 10-0000 J (JJJ + +cL) J3).
(12)
The pseudo-inverse is only computed for a matrix of size M. Using the dacobian-vector product, Equation (12) can be computed without storing large matrices G shown in Algorithm
Wasserstein Natural Policy Gradient (WNPG). It is possible to incorporate local information about the behavior of a policy in standard algorithms for policy gradient as summarized Algo- rithm 1. In its simplest form, one first needs to compute the gradient ğk of the objective L(Ox) using, for instance, the REINFORCE estimator computed using N trajectories used compute the BEMs which fed input, along with the gradient OR to get an estimate of the WNG Finally, the policy can be updated in the direction of g) Algorithm I can also be used in combination with an explicit W2 penalty control non-local changes in behavior of the policy thus ensuring a policy improvement property as in (Pacchiano et al., 2019, Theorem 5.1). In that case, WNG enhances convergence by acting as a second-order optimizer, as discussed Sec- tion 3.1. The standard gradient Ou in Algorithm is then simply replaced by the one computed in (Pacchiano et al., 2019, Algorithm 3). Section we show that this combination, which we call behavior-guide WNPG (BG-WNPG) leads to the best overall performance.
The trajectories are then
Wasserstein Natural Evolution Strategies (WNES). ES treats the total reward observed on a trajectory under policy ne black-box function L()) (Salimans al., 2017; Mania et al., 2018;

---

Published as a conference paper at ICLR 2021
Algorithm 2: Wasserstein Natural Evolution Strategics
1: Input Initial policy Re. a > 0, <
2: for iteration & = 1,2, do
3: Sample 61...., 6n (0,
4: Perform rollouts {Ta) of length T using the perturbed parameters {0n O2 + Ten J
and compute behavioral embeddings (X,, ()(>))
5:
Compute gradient estimate of C(0") using Equation (13) and trajectories (T")
6: Compute Jacobian matrix J appearing in Algorithm 3 using Equation (14).
7: Compute WNG D3 using Algorithm 3, with samples {X and computed OJ and J.
8: Update policy using Equation (15).
9: end for
Choromanski et al., 2020). Evaluating it under N policies whose parameters 0t are gaussian pertur- bations centered around 0g and with variance can give an estimate of the gradient of L(Os):
N Σ No
9
(0" -0))
(13)
n=1
Instead of directly updating the policy using Equation (13), it is possible to encourage either prox- imity or diversity in behavior using the embeddings Xx, = ((TT)) of the trajectories Tn generated for each perturbed policy "æ." Those embeddings can be used as input Algorithm 3 (see appendix), along with Equation (13) to estimate the 9% which captures similarity in behavior. The algorithm remains unchanged except for the estimation of the Jacobian J of Equation (10) which becomes
1
Jm..
hm (X,) 0z).
(14)
No
The policy parameter can then be updated using an interpolation between 9R and the WNG 9h i.e.,
Adg x (1 8) 50
(15)
with 5 that can also be negative. Positive values for 5 encourage proximity in behavior, the limit case being 8 = where full WNG step is taken. Negative values encourage repulsion and therefore need to compensated by ÿN to ensure overall policy improvement. Algorithm summarizes the whole procedure, which can be easily adapted from existing ES implementations by calling variant Algorithm 3. particular, it can also be used along with an explicit W2 penalty, which case the proposed algorithm in Pacchiano al. (2019) is used estimate the standard gradient ğN of the penalized loss. Then the policy is updated using Equation (15) instead of Gr. We refer to this approach behavior-guided WNES (BG-WNES).
EXPERIMENTS
We now test the performance of our estimators for both policy gradients (PG) and evolution strate- gies (ES) against their associated baseline methods. We show that in addition to an improved com- putational efficiency, our approach can effectively utilize the geometry induced by Wasserstein penalty to improve performance, particularly when the optimization problem is ill-conditioned. Fur- ther experimental details can be found in the appendix, and our code available online!.
Policy Gradients. We first apply WNPG and BG- WNPG to challenging tasks from OpenAI Gym (Brockman et al., 2016) and Roboschool (RS). We compare performance against behavior guided policy gradients (BGPG), (Pacchiano et al., 2019), PPO with clipped surrogate objective (Schulman et al., 2017) (PPO (Clip)), and PG with no trust region (None). From Figure 2, we can see that BGPG outperforms the corresponding KL-based method (PPO) and vanilla PG, also demonstrated in the work of Pacchiano et al. (2019). Our method (WNPG) matches or exceeds final performance of BGPG all tasks. Moreover, combining both (BG-WNPG) produces the largest gains all
'https:// http://tt.ttbcccmmmmcccossst/ttt

---

Published as a conference paper at ICLR 2021
(RS) Walker2d
(RS) Halfcheetah
(RS) Ant
2000
750
2000
Reward
500
1000
250
0.25
0.50
0.75
1.00
0.00
0.25
0.50
1.00
0.0
2.0
(RS) Reacher
Inveree@DoublePenddlumm
Hopper
3000
7500
Reward
2000
:20
2500
1000
50000
100000
150000
200000
0.00
0.25
0.50
0.55
1.00
0.00
0.25
0.50
0.55
1.00
Timesteps
Timesteps
led
Timesteps
None
PPO (Clip)
BGPG
WNPG
BG-WNPG
Figure 2: WNG-based algorithms provide large gains on tasks where initial progress is difficult. The performance mean + standard deviation is plotted versus time steps for 5 random seeds on each task.
environments. Final mean rewards are reported in Table 1. is also important to note that WNG- based methods appear offer the biggest advantage tasks where initial progress is difficult. To investigate this further, we computed the hessian matrix at the end of training for each task and measured the ratios of its largest eigenvalue to each successive eigenvalue (Figure 3). Larger ratios indicate ill-conditioning, and it is significant that WNG methods produce the greatest improvement on the environments with the poorest conditioning. This is consistent with the findings in Arbel et al. (2020) that showed WNG to perform most favorably compared to other methods when the optimization problem is ill-conditioned, and implies useful heuristic for gauging when WNG- based methods are most useful for a given problem.
Evolution Strategies To test our estimator for WNES, as well as BG-WNES, we applied our approach the environment introduced by Pacchiano et al. (2019), designed to test the ability of gehavior-guides learning succeed despite deceptive rewards. During the task, the
50
(RS) Walker2d
Halfcheatah
40
Ant
log (Og(o//1
30
(RS) Rucche
agent receives a penalty proportional to its dis- tance from goal, but wall is placed directly in the agent's path (Figure 7). This barrier induces local maximum in the objective-a naïve agent will simply walk directly towards the goal and get stuck the barrier. The idea is that the behavioral repulsion fostered by ap- plying positive coefficient to the Wasserstein penalty (5 > 0) will encourage the agent to seek novel policies, helping
20
Hopper
200
400
600
800
1000
Figure 3: Condition numbers for different tasks.
eventually circumvent the wall. As Pacchiano al. (2019), we test two agents, a simple point and quadruped. We then compare our method with vanilla ES as described by Salimans et (2017), ES with gradient norm clipping, BGES (Pacchiano et al., 2019), and NSR-ES (Conti et al., 2018a). In Figure 4, prove over the baselines for both agents. To test that the improvement shown by BG- WNES wasn't simply case of additional "repulsion" supplied by the WNG to BGES, also tested BGES with an increased / = 0.75, compared to the default of 0.5. This resulted a decrease in performance, attesting to the unique benefit provided by the WNES estimator.
see that WNES and BG-WNES im-
Computational Efficiency We define the computational efficiency of an algorithm as the rate with which it accumulates reward relative its runtime. To test the computational efficiency of
8

---

Published as a conference paper at ICLR 2021
Point
Quadruped
Reward
E5-10
5000
H63-53
-3000
-000
DGHNNNS
0.75)
.000
100000
325000
300000
400000
Timesteps (a)
Timesteps
(b)
Figure 4: WNES methods more reliably overcome local maxima. Results obtained on the point (a) and quadruped (b) tasks. The mean + standard deviation is plotted across 5 random seeds.
-44
100
(CP)
gradecpp
Reword
403
BGPG
BG-WNPG
BiGs op 0.79)
112
1100
BG-WNES
(R5) Walkerdd (15) Hallcheetshh
[R5] Ant
(R5) Reacher
DoxclePendslum
Hogper
Paint
Quddraged
Figure 5: WNG methods improve computational efficiency. Average reward per minute is plotted for both gradient tasks (left) and ES tasks (right) for the runs depicted above.
approach, we plotted the total reward divided by wall clock time obtained by each agent for each task (Fig, 5). Methods using WNG estimator were the most efficient on each task for both PG and ES agents. On several environments used for the policy gradient tasks, the added cost of BG-WNPG reduced its efficiency, despite having the highest absolute performance.
CONCLUSION
Explicit regularization using divergence measures between policy representations has been a com- mon theme in recent work on policy optimization for RL. While prior works have previously fo- cused on the KL divergence, Pacchiano et al. (2019) showed that Wasserstein regularizer over behavioral distributions provides powerful alternative framework. Both approaches implicitly de- fine a form of natural gradient, depending on which divergence measure chosen. Through introduction of WNPG and WNES, demonstrate that directly estimating the natural gradient of the un-regularized objective can deliver greater performance at lower computational cost. These algorithms represent novel extensions of previous work the WNG to problems where the parameterization trick not available, as well as to black-box methods like ES. Moreover, using the WNG in conjunction with a WD penalty allows the WNG take advantage of the local geometry induced by the regularization, further improving performance. We also provide a novel compar- ison between the WNG and FNG, showing that the former has significant advantages on certain problems. We believe this framework opens up a number of avenues for future work. Developing a principled way identify useful behavioral embeddings for a given RL task would allow to get the highest benefit form WNPG and WNES. From theoretical perspective, it would be useful to characterize the convergence boost granted by the combination of explicit regularization and the corresponding natural gradient approach.
The authors would like to thank Jack Parker-Holder for sharing his code for BGPG and BGES, well colleagues at Gatsby for useful discussions.
Acknowledgments

---

Published as a conference paper at ICLR 2021
REFERENCES
Shun-ichi Amari. Neural learning in structured parameter spaces- natural riemannian gradient. In M. C. Mozer, M. Jordan, and T. Petsche (eds.), Advances Neural Information Processing Systems 9, pp. 127-133. MIT Press, 1997.
M Arbel, Gretton, W Li, and G Montufar. Kernelized wasserstein natural gradient. Interna- fional Conference Learning Representations, 2020. URL
n
Barumiid=Hk12211e
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym, 2016.
Yifan Chen and Wuchen Li. Natural gradient in Wasserstein statistical manifold. arXiv: 08880000888 [cs, math), May 2018. URL http:://axtat org/ ats. 08380. arXiv: 1805.08380.
Krzysztof Choromanski, Aldo Pacchiano, Jack Parker-Holder, Yunhao Tang, Deepali Jain, Yuxiang Yang. Atil Iscen, Jasmine Hsu, and Vikas Sindhwani. Provably robust blackbox optimization for reinforcement learning. In Conference on Robot Learning, pp. 683-696, 2020.
Edoardo Conti, Vashisht Madhavan, Felipe Petroski Such, Joel Lehman, Kenneth Stanley, and Jeff Clune. Improving exploration in evolution strategies for deep reinforcement learning via popu- lation of novelty-secking agents. In Advances in Neural Information Processing Systems 31, pp- 5027-5038. Curran Associates, Inc., 2018a.
Edoardo Conti, Vashisht Madhavan, Felipe Petroski Such, Joel Lehman, Kenneth Stanley, and Jeff Clune. Improving exploration in evolution strategies for deep reinforcement learning via a pop- ulation of novelty-seeking agents. In Advances in neural information processing systems, pp- 5027-5038, 2018b.
Prafulla Dhariwal, Christopher Hesse, Oleg Klimov, Alex Nichol, Matthias Plappert, Alec Radford, John Schulman, Szymon Sidor, Yuhuai Wu, and Peter Zhokhov. Openai baselines. https: // ///github.com com/ openai /basel ines, 2017.
Aude Genevay, Marco Cuturi, Gabriel Peyré, and Francis Bach. Stochastic optimization for large-scale optimal transport. In D. D. Lee, M. Sugiyama, U. V. Luxburg, I. Guyon, and R. Garnett (eds.), Advances in Neural Information Processing Systems 29, pp. 3440- 3448. Curran Associates, Inc., 2016. URL http: / /ttpppaaersaae/////ee nips cc/paper/ cechhhsstt-cte stochastic-oppimiaat..... pdf.
for-large scale seeale-optinne imallmma-- transport...
Sham M Kakade. A natural policy gradient. In Advances in neural information processing systems, pp. 1531-1538, 2002.
Solomon Kullback and Richard A Leibler. On information and sufficiency. The annals of mathe- matical statisties, 22(1):79-88 86, 1951.
Seong Jae Lee and Zoran Popović. Learning behavior styles with inverse reinforcement learning. ACM transactions on graphics (TOG), 29(4):1-7, 2010.
Wuchen Li. Geometry of probability simplex via optimal transport. arXiv: 1803.06600 [math), March 2018. URL http::://www..oo////eeea// /t///aabs//80...// /1003. 06360. arXiv: 1803.06360.
Wuchen Li and Guido Montufar. Natural gradient via optimal transport. arXiv:1803.. 07033 [cs,
math), March 2018a. URL http: atwpv..xivvvv/ wttp://..rrr/bb//////0/ 0703333333303.....33 arXiv: 1803.07033.
Wuchen Li and Guido Montufar. Ricci curvature for parametric statisties via optimal transport. arXiv:1807.07095 (cs. math, stat], July 2018b. URL http:://axxxv..rg/bbbb. 07095. arXiv: 1807.07095.
Wuchen and Jiaxi Zhao. Wasserstein information matrix. 1xXiv1190.111244 [cs, math, stat], November 2019. URL http: /tap:iv////iv...// wowploorg/aaoao htgg/bbs//99/ 99.0. 11248. arXiv: 1910.11248
10

---

Published as a conference paper at ICLR 2021
Horia Mania, Aurelia Guy, and Benjamin Recht. Simple random search of static linear policies is competitive for reinforcement learning. Advances in Neural Information Processing Systems, pp. 1800-1809, 2018.
Elliot Meyerson, Joel Lehman, and Risto Miikkulainen. Learning behavior characterizations for novelty search. In Proceedings of the Genetic and Evolutionary Computation Conference 2016, pp. 156, 2016.
Aldo Pacchiano, Jack Parker-Holder Yunhao Tang, Anna Choromanska, Krzysztof Choromanski, and Michael I Jordan. Learning score behaviors for guided policy optimization. arXiv preprint arXiv:1906.04349 2019.
Martin Puterman. Markov decision processes: discrete stochastic dynamic programming. John Wiley & Sons, 20100
Tim Salimans, Jonathan Ho, Xi Chen, Szymon Sidor, and Πlya Sutskever. Evolution Strategies as Scalable Alternative to Reinforcement Learning. arXiv e-prints, art. arXiv:1703.0386/ March
2017.
John Schulman, Sergey Levine, Philipp Moritz, Michael I. Jordan, and Pieter Abbeel. Trust region policy optimization. CoRR, abs/1502.05477. 2015. URL http://arrret /ttp:i///xxiv..a/ wttp://...r//b//1022222222 05477.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: 1707.06777, 2017.
Richard Sutton and Andrew G. Barto. Reinforcement Learning: An Introduction. The MIT Press, second edition, 2018. URL html.
Yunhao Tang and Shipra Agrawal. Implicit Policy for Reinforcementteeneee arXiv: [cs, stat), February 2019. URL http://www...eee.ee.eeeeeeeeoooe 1806.06798.
2nd.
http://v..rggb/////b/b/oo 117.a///s...66.06..... 06798. arXiv:
Cedric Villani. OPTIMAL TRANSPORT: old and new. SPRINGER-VERLAC BERLIN AN, 2016.

---

Published as a conference paper at ICLR 2021
Environment
BG-WNPG (ours)
WNPG (ours)
BGPG
PPO (Clip)
None
(RS) Walker2d
739.51+81.10
683.38*
652.53
516.47
529.60
(RS) Halfcheetah
1817.68 79.64 79.64
1655.04
1668.75*
1412.57
1334.27
(RS) Ant
2072.56 85.24
1844.14
1735.98
1627.01
1566.63
(RS) Reacher
4.37
2.13*
0.46
-0.37
-6.46
DoublePendulum
7254.98 419.36
6740.96*
6526.48
4401.24
3964.51
Hopper
2439.56 = 394.39
2235.23*
1934.43
1890.12
1816.58
Table 1: Final average return over 5 trials for the PG experiments depicted in Figure 2. = values denote one standard deviation across trials. The value for the best-performing method is listed in bold, while denotes the second best-performing method. BG-WNPG reaches the highest performance on all tasks. WNPG beats the best-performing baseline (BGPG) on all tasks except HallCheetah, where the difference is small.
A BACKGROUND
A.1 POLICY OPTIMIZATION
agent interacting with an environment form a system that can be described by a state variable belonging state space S. In the Markov Decision Process (MDP) setting, the agent interact with the environment by taking an action from a set of possible actions A given the current state state s' according
of the system. As consequence, the system moves transition function ((s" a, s) which describes the probability of moving to state s' given the previous state s and action a. The agent also receives partial reward which can be expressed a possibly randomized function of the new state s', r = r(s'). The agent has access to a set of possible policies (a(a)s) parametrized by 0 E R² and that generates an action a given a current state s. Thus, each policy can be seen as probability distribution conditioned state s. Using the same policy induces whole trajectory of catte-action-rewarde
probability
(s,, a, T Tr)120 which can be viewed as sample from trajectory distribution Pn defined over the space of possible trajectories for given random trajectory 7 induced by a policy ne. the agent receives a total discounted reward R(T) := E 7f (r=(r(s)) with discount factor 0 < 1. This allows to define the value function as the expected total reward conditioned on a particular initial state s:
Hence,
(8() Eggler
Σ 7 r(814))
(16)
When the gradient of the score function V log To (a(s)) is available, the policy gradient theorem allows us express the gradient R(0):
VoR(O) Er
V log (og
(17)
where the expectation is taken over trajectories 7 under F. and Ag(s, a) represents the advantage function which can be expressed in terms of the value function Vo(s) in terms of
As a,) = E es [r(s(41) + VV Vø(st).
The agent seeks an optimal policy Tg> that maximizes the expected total reward under the trajectory distribution: R(0) [R(T))
B WASSERSTEIN NATURAL GRADIENT
Connection the Fisher natural gradient and proximal methods. Both WNG and FNG are obtained from a proximity measure between probability distributions:
Proposition 2 Let D(0,6 8¹) be either the KL-divergence KL( Tgy or the Wasserstein-2 distance between the behavioral distributions W2(go, gor) and let gD be either the FNG gf WNG g".
12

---

Published as a conference paper at ICLR 2021
X
(
X
90
au+p
Ø1
π++U
Figure 6: visualization of the behavioral transport function.
then
D lim arg max 3 9k
3-) - [(0x) -10)
(0
(18)
+-++00
Equation (18) simply states that the both WNG and FNG arise limit cases of penalized objectives provided the strength of the penalty 3 diverges to infinity and the step-size is shrank proportion- ally 8-1. An additional global rescaling by of the total objective prevents it from collapsing to O. Intuitively, performing Taylor expansion of Equation (18) recovers an equation similar to Equation (8). Equation (18) shows that using penalty that encourages global proximity between successive policies, it is possible to recover the local geometry of policies (captured by the local ) by increasing the strength of the penalty using appropriate re- scalin. This also informally shows why both natural gradients are said to be invariant to re-parametrizatioo (Arbel al., 2020, Proposition 1), since both KL and W2 remains unchanged if go is parameterized in a different way.
C ALGORITHM FOR ESTIMATING WNG
Algorithm 3: Efficient Wasserstein Natural Gradient
1: Input mini-batch of samples {X,, b distributed according to ge. gradient direction I, basis functions f,, hus regularization parameter e.
2: Output Wasserstein Natural gradient g
3: Compute matrix C of shape M x Nd using Cm.(n.n) +- Ozhi (Xx).
4: Compute similarity matrix L 1. CCCT.
5: Compute surrogate vector V using Equation (11).
6: for iteration= 1,2, ...MM do
7: Use automatic differentiation on Vm. to compute Jacobian matrix J in Equation (10).
8: end for
9: Compute matrix D of shape M M using D +- JJJ +el.
10: Compute a vector of size M using b < Jg.
11: Solve linear system of size M b solve (D,
12: Return 9 + [ JT
D ADDITIONAL EXPERIMENTAL DETAILS
D. POLICY GRADIENT TASKS
We conserve all baseline and shared hyperparameters used by Pacchiano et al. (2019). More precisely, for each task
hyperparameter sweep over learning rates in the set [le-5, 5e-5, 1e-4, 3e-4), and used the concatenation-of-actin behavioral embedding &(7) = (og, (1,-.-, ar) with the base network implementation the same Dhariwal al. (2017).
13

---

Published as a conference paper at ICLR 2021
Figure 7: A visualization of the quadruped task. The agent receives receives reward the closer is to the goal (green). A na've agent will get stuck in the local maximum at the wall if it attempts move directly the goal.
The WNG hyperparameters also left the same in Arbel et al. (2020). Specifically, the number of basis points was set as M 5, the reduction factor was bounded in the range [0.25, 0.75]. and [1e-10, le5].
D.2 EVOLUTION STRATEGIES TASKS
As with the policy gradient tasks, we conserved all baseline and shared hyperparameters used by Pacchiano et al. (2019). Specifically, for the point task, we set the learning rate be 7) 0.1, the standard deviation of the noise to be a = 0.01, the rollout length H was 50 time steps, and the behavioral embedding function to be the last state (DT) sи- For the quadruped task we set
? 0.02, 0.02, H 400, and ④(7) Ei r: Ei ei (reward-to-go encoding; Pacchiano et al. (2019) for more details). Both tasks used 1000-dimensional random features and embeddings from the previous policies to compute the WD.
For WNG, the same hyperparameters were used as in the policy gradient tasks.
D.3 EXPERIMENTAL SETTING OF FIGURE
The Objective E IR 100
We consider a function (z) is the sum of sinc functions over all dimensions of
100
sin(x:) 1
((x)
Σ
(19)
Such function is highly non-convex and admits multiple bad local minima with the global minimum of (xr) reached for x² = 0. However, we do not make use of this information during optimization. To alleviate the non-convexity of this loss, we consider gaussian relaxation objective [(0) obtained
by taking the expectation of (xx) over the 100 dimensional vector w.r.t. to gaussian qa with parameter vector 0. Thus the objective function to be optimized is a function of
(()) = Egu [v]()]
(20)
The parameter vector 0 is of the form 0 (u, v). where y is the mean of the gaussian go and is a vector in IR 100 parameterizing the covariance matrix E of the gaussian We will later consider two parameterizations for the covariance matrix.
The minimal value of C(0) is reached when the gaussian qa is degenerate with Σ = 0 and mean μ x* 0. Hence, the mean parameter of the global minimum of L(0) recover the global optimum of vh.
Parameterization of the gaussian We choose the covariance matrix of the gaussian to be diagonal and consider two parameterizations for the covariance matrix Σ: diagonal and log diiagonal. For the diagonal parameterization the Covariance Σn en and for the log-diagonal set En; exp(2n:).
14

---

Published as a conference paper at ICLR 2021
Optimization methods We consider different optimization methods using the same objective (()). For the penalty methods, we use the closed form expressions for the both the Wasserstein distance and KL which are available explicitly in the case of gaussians.
For the Natural gradient methods (WNG) and (FNG), we use the closed form expressions which are also available in the gaussian case. We denote them V L(O) for (WNG) and V L()) for FNG and express them in terms of the euclidean/standard gradient (()):
Diagonal parameterization:
WNG:
V. =+
C v
(21)
FNG:
25 == EV -V
(22)
Log-diagonal parameterization:
WNG::
V. 210)= =10)=
C(0) [(0) = v2((0)
(23)
FNG:
V L(0) = 5 V.C(O),
C(0) = EV n(())
(24)
Training details are varied.
Training up to 4000 gradient iterations, with A -9 and .1 unless they
15


Reviews:
[{"review_id": "wETg1hY3KOM", "reviewer": "AnonReviewer5", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper introduces methods for reinforcement learning based on a Wasserstein Natural Gradients (WNG), an approach to measuring similarity between policies. The methods are based on Policy Gradients and Evolutionary Strategies and add policy similarity term based on WNG instead of KL constraint as in TRPO.\n\nThe paper is well written and easy to follow (the conclusion section is missing though). Although the method is interesting, I think that the current experimental evaluation has significant flaws: the method does not demonstrate the state-of-the-art performance and significantly improves over the baselines on a small number of tasks. I believe that the paper will significantly benefit from comparisons  with stronger baselines such as PPO. Moreover, the experiments performed in the paper (figure 2) demonstrate that the approach only marginally outperforms the baselines on the harder HalfCheetah and Hopper tasks that raises concerns regarding the generality of the approach.\n\nThe paper proposed an interesting approach to policy constraints in RL but the experimental evaluation is not sufficient. ", "extra_scores": {}, "extra_text": {}}, {"review_id": "mOAU452mo3", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "### Summary\nThis paper proposes to use natural gradient instead of standard gradient to optimize a regularized objective with the regularization being the Wasserstein distance between the so-called behaviour distributions for the previous policy and new policy. It then combines this Wasserstein gradient descent with Policy Gradient and Evolutionary Strategies.  Experiments conducted in OpenAI and Roboschool show some promising results for this combination.\n\n### Strong points: \n-\tClarity: The paper is well structured\n-\tEmpirical significance: The empirical results seem promising where it shows that Wasserstein gradient descent could be a good choice to constraint the behaviour changes efficiently without sacrificing too much computational cost. I personally like that the authors used Roboschool (RS) as a benchmark for continuous control tasks as it is, as opposed to Mujuco, a free simulator software (thus it is more accessible and reproducible for everyone).   \n\n### Weak points: \n-\tNovelty: the work however has low technical novelty where it combines several known results into a new framework. In particular, the idea of constraining behaviour policy changes and the efficient way to estimate Wasserstein gradient descent are all known and off-the-shelf. Adopting Wasserstein gradient descent to RL constraint update seems straightforward that does not require any significant technical challenge. The interpretation of the framework also seems straightforward, e.g., it is of course that updating along the Wasserstein natural gradient would incorporate the local geometry of parameterization and help overcome some ill-conditioning issues where KL has. \n-\tEmpirical significance: The empirical results though promising are not strong given that this is mostly an empirical work. In particular, the present work presents the experiments for PG case in only 4 environments which I think insufficient to make a reliable conclusion about its empirical significance. \n\n###  Questions for the authors \n-\tIn Section 2: “Reusing trajectories can reduce the computational cost but drastically increases the variance of the gradient estimator”. Could the authors elaborate on why reusing trajectories drastically increases the variance of the gradient estimator?\n-\tFig. 1 (c): Have all algorithms been initialized at the same initial point? Also, according to Fig.1 (c) that it seems that FNGD has a ‘right’ convergence when it converges to the point where \\sigma=0 and \\mu = midpoint, why in the last paragraph of page 4, the authors conclude that “FNG remains far away from optimum”? What do I miss here? \n-\tWhat is the difference between W2-penality and WNGD in Figure 1? \n-\tOn page 5, “The Wasserstein penalty Equation (4) encourages global proximity between updates qθk”. What does globality here refer to while Eq (6) holds only locally?\n\n### Minor comments \n-\tThe second term of Eq. (5): Shouldn’t it be f_u instead of f there?\n-\tEq. (8): \\argmax_{u}\n-\tThe second last sentence at the end of page 4: cite -> \\cite{sth}\n-\tIt seems that \\mu and \\sigma in Fig. 1 (c) have not defined explicitly anywhere. In the capture of Fig.1., it writes \\theta = (mu,v), so there is a chance of inconsistent notations here? \n-\tOn page 5, “To avoid slowing-down, there is an intricate balance between the step-size and penalty β that needs to be maintained Schulman et al.\n(2017)”: \\cite -> \\citep  \n\n###  My initial recommendation\n\nGiven the weak and strong points above, I vote for rejecting for this current form.\n\n### My finial recommendation \n\nAfter the discussion and revision, the authors have presented more convincingly and more clearly the empirical significance and applicability of their method. I highly recommend the authors to highlight the lastest discussion in the final paper, especially the ill-conditioned argument, as it is highly relevant to the practitioners. I think this paper can be interesting for a moderate number of readers, especially the use of the open-sourced Roboschool could also increase its reproduciability. I agree to increas my score to 6. ", "extra_scores": {}, "extra_text": {}}, {"review_id": "6xkuPclfEPI", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "\nAmari's Natural Gradient has been very successfully applied for policy optimization, e.g. in a recent line of work by Schulman et al. These benefit of using these natural gradients that restrict the change in KL divergence between successive policies was more stable and faster convergence. However two distributions may differ a lot in terms of KL divergence but because of the dynamics of the MDP they may still have almost the same behavior, therefore recent work has focused on using the Wasserstein distance to measure the divergence between successive policies which naturally leads to \"Wasserstein Natural Descent\". In this paper the authors build upon recent work on Kernelized Wasserstein NGD by making it more widely applicable and scalable. Moreover they present a good empirical comparison between KL-NGD and Wasserstein-NGD on a combination of pedagogical toy problems and some standard RL benchmarks from OpenAI gym. Overall this paper will be a good contribution to the conference and I recommend acceptance.\n\n\n**Corrections and suggestion for improving presentation**\n\n1. [Citations] Page 2, third paragraph from bottom cites Schulman 2015 instead of Schulman 2017 for PPO. Page 4 last paragraph has a citation missing and Li and Zhao \"Wasserstein Information Matrix\" paper is cited twice. \n\n2. I think it will be better to write down that equation (5) defines $f_u$. Also shouldn't the $1/2$ factor in equation (5) be actually $\\beta / 2$ ? \n\n3. On page 3 it is said that \"the penalty only accounts for global proximity in behavior .... \" in reference to equation (4), but PPO is not implemented with a single value of $\\beta$, i.e. the strength of the penalty typically varies as optimization goes on. Why can't the same be done for the weight of the $W_2$ penalty ?\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper explores the Wasserstein natural gradient in the context of reinforcement learning. R5 rated the paper marginally below the acceptance threshold, but is not very confident about the correctness of his/her assessment. His/her main criticism was the experimental evaluation. This concern was shared by a confident R1. R1 found the paper well structured and that it contains encouraging empirical results, but low technical novelty and (initially) insufficient experiments. His/her initial recommendation was reject, but following an extensive discussion and improvements of the manuscript by the authors, he/she was more convinced about the empirical significance and applicability of the method, and raised his/her score to 6, indicating that the interpretation and presentation improved but that the paper might be interesting only to a moderate number of readers. A confident R2 found this paper very good, although only providing a short review. Two other unfinished or not sufficiently confident reports were not taken into account. Weighing the reports by contents, confidence, and participation in the discussion, the paper scores marginally above the acceptance threshold. In view of the authors' responses, I am discounting R5's criticism about lack of comparison with the PPO baseline. I personally consider the paper very well written, that it presents a natural and potentially useful application of the Wasserstein natural gradient to the context of reinforcement learning, and enjoyed the discussion of behavioral geometry. I am recommending a borderline accept. However, I also appreciate the concern of the referees about the limited technical innovation and how some of the strengths of the method could be presented more convincingly. Please take these comments carefully into consideration when preparing the final version of the paper. 
Decision: Accept (Poster)