E FFICIENT W ASSERSTEIN N ATURAL G RADIENTS FOR R EINFORCEMENT L EARNING
A BSTRACT
A novel optimization approach is proposed for application to policy gradient methods and evolution strategies for reinforcement learning (RL).
The procedure uses a computationally efﬁcient 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 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
I NTRODUCTION
Deﬁning efﬁcient optimization algorithms for reinforcement learning (RL) that are able to leverage a meaningful measure of similarity between policies is a longstanding and challenging problem (Lee & Popovi ´ c, 2010; Meyerson et al., 2016;
Conti et al., 2018b).
Many such works rely on similarity measures such as the Kullback-Leibler (KL) divergence (Kullback & Leibler, 1951) to deﬁne procedures for updating the policy of an agent as it interacts with the environment.
These are generally motivated by the need to maintain a small variation in the KL between successive updates in an 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) as a way to 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 a 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-motivated, one can still ask if it is possible to include information about the behavior of policies when measuring similarity, and whether this could lead to more efﬁcient algorithms.
Pacchiano et al.
(2019) provide a ﬁrst insight into this question, representing policies using behavioral distributions 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 to use such behavioral similarity as a global soft penalty to the total objective.
Hence, like the KL penalty, proximity between policies is measured globally, and does not necessarily exploit the local geometry deﬁned by the behavioral embeddings.
In this work, we show that substantial improvements can be achieved by taking into account the local behavior of policies.
We introduce new, efﬁcient optimization methods for RL that incorporate the local geometry deﬁned 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, 2018a;b; Li, 2018; Li & Zhao, 2019; Chen & Li, 2018) which introduces the notion of the Wasserstein Information Matrix to deﬁne a local behavioral similarity measure between policies.
This allows us to identify the Wasserstein Natural Gradient (WNG) as a key ingredient for optimization methods that rely on the local behavior of policies To
enable efﬁcient estimation of WNG, we build on the recent work of Arbel et al.
, and further extend it to cases where the re-parameterization trick is not applicable but only the score function of the model is available.
2- This allows us to 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 over using the penalty alone without access to local information.
We ﬁnd that such
WNG-based methods are especially useful on tasks in which initial progress is difﬁcult.
3- Finally, we demonstrate, to our knowledge, the ﬁrst 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 a brief review of policy gradient approaches and the role of divergence measures as regularization penalties.
In Section 3 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 5 contains our empirical results.
B ACKGROUND
Policy Gradient (PG) methods directly parametrize a policy π θ , optimizing the parameter θ using stochastic gradient ascent on the expected total discounted reward R ( θ ) .
An estimate ˆ g k of the gradient of R ( θ ) at θ k can be computed by differentiating a surrogate objective L θ which often comes in two ﬂavors, depending on whether training is on-policy (left) or off-policy (right):
[Equation]
The expectation ˆ E is an empirical average over N trajectories τ i
= ( s i 1 , a i 1 , r i 1 , ..., s iT , a iT , r iT ) of state-action-rewards obtained by simulating from the environment using π θ k .
The scalar ˆ A t is an estimator of the advantage function and can be computed, for instance, using
[Equation]
where γ ∈ [0 , 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 et al., 2017).
Indeed, performing multiple policy updates while using trajectories from an older policy π θ old means that the current policy π θ can drift away from the older policy.
On the other hand, the objective is obtained as an expectation under π θ 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
π θ /π θ old ).
When π θ is too far from π θ old , the importance weight can have a large variance.
This can lead to a 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), or a soft constraint, as in PPO (Schulman et al., 2017), to encourage proximity between policies.
In the ﬁrst case, TRPO recovers the FNG with a step-size further adjusted using line-search to 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:
[Equation]
The KL penalty prevents the updates from deviating too far from the current policy π θ k , 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.
Behavior-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 a notion of proximity in behavior between policies for PO.
Exploiting similarity in behavior during optimization allows to take larger steps in directions where policies behave similarly despite having a large KL divergence.
To capture a sense of global behavior, they deﬁne a behavioral embedding map (BEM) Φ that maps every trajectory τ to a behavior variable
X =
Φ( τ ) belonging to some embedding space E .
The behavior variable X provides a simple yet meaningful representation of each the trajectory τ .
As a random variable, X is distributed according to a distribution q θ , called the behavior distribution .
Examples of Φ include simply returning the ﬁnal state of a trajectory ( Φ( τ )
=
s T ) or its concatenated actions ( Φ( τ ) =
[ a 0 , . .
. , a T ] ).
Proximity between two policies π θ and π θ (cid:48) is then measured using the Wasserstein distance between their behavior distributions q θ and q θ (cid:48) .
Although, the KL could also be used in some cases, the Wasserstein distance has the advantage of being well-deﬁned even for distributions with non-overlapping support, therefore allowing more freedom in choosing the embedding Φ (see Section 3.1).
This leads to a penalized objective that regulates behavioral proximity:
[Equation]
where β ∈ R is a hyper-parameter controlling the strength of the regularization.
To compute the penalty, Pacchiano et al.
(2019) use an iterative method from Genevay et al.
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 β also mean that the policy is updated using smaller steps, which can impede convergence.
An optimal trade-off between the rate of convergence and the precision of the estimated Wasserstein distance can be achieved using an adaptive choice of β as done in the case of PPO Schulman et al.
For a ﬁnite value of β , the penalty accounts for global proximity in behavior and doesn’t explicitly exploit the local geometry induced by the BEM, which can further improve convergence.
We introduce an efﬁcient method that explicitly exploits the local geometry induced by the BEM through the Wasserstein Natural gradient (WNG), leading to gains in performance at a reduced computational cost.
When global proximity is important to the task, we show that using the Wasserstein penalty in Equation (4) and optimizing it using the WNG yields more efﬁcient updates, thus converging faster than simply optimizing Equation (4) using standard gradients.
T HE W ASSERSTEIN N ATURAL G RADIENT
The Wasserstein natural gradient (WNG) (Li & Montufar, 2018a;b) corresponds to the steepestascent direction of an objective within a trust region deﬁned by the local behavior of the Wasserstein2 distance ( W 2 ).
The W 2 between two nearby densities q θ and q θ + u can be approximated by computing the average cost of moving every sample X from q θ to a new sample X (cid:48) approximately distributed according to q θ + u using an optimal vector ﬁeld of the form ∇ x f u ( x ) so that
X (cid:48)
=
X + ∇ x f u ( X ) (see Figure 6).
Optimality of ∇ x f u is deﬁned as a trade-off between accurately moving mass from q θ to q θ + u and reducing the transport cost measured by the average squared norm of ∇ x f u
[Equation]
[Equation]
where the optimization is over a suitable set of smooth real valued functions on E .
Hence, the optimal function f u solving Equation (5) deﬁnes the optimal vector ﬁeld ∇ x f u ( x ) .
Propositionmakes this intuition more precise and deﬁnes the Wasserstein Information Matrix.
Proposition 1 (Adapted from Deﬁntion 3 Li & Zhao (2019))
The second-order Taylor expansion of W 2 between two nearby parametric probability distributions q θ and q θ + u is given by
[Equation]
) where G ( θ ) is the Wasserstein Information Matrix (WIM), with components in a basis ( e 1 , ..., e p )
[Equation]
The functions f j solve Equation (5) with u chosen as e j .
Moreover, for any given u , the solution f u to Equation (5) satisﬁes E
θ [ (cid:107)∇ x f u ( X ) (cid:107) 2 ] =
u (cid:62) G
( θ ) u .
When q θ and q θ + u are the behavioral embedding distributions of two policies π θ and π θ + u , the function f u allows to transport behavior from a policy π θ to a behavior as close as possible to π θ + u with the least cost.
We thus refer to f u as the behavioral transport function .
The function f u determines how hard it is to change behavior locally from policy π θ in a direction u , thus providing a tool to ﬁnd update directions u with either maximal or minimal change in behavior.
Probing all directions in a basis ( e 1 , ..., e p ) of parameters allows us to construct the WIM G ( θ ) in Equation (7) which summarizes proximity in behavior along all possible directions u using u (cid:62) G
( θ ) u
= E q θ
[ (cid:107)∇ x f u ( X ) (cid:107) 2 ] .
For an objective L ( θ ) , such as the expected total reward of a policy, the Wasserstein natural gradient (WNG) is then deﬁned as the direction u that locally increases L ( θ + u ) the most with the least change in behavior as measured by f u .
Formally, the WNG is related to the usual Euclidean gradient g = ∇ θ L ( θ ) by
[Equation]
From Equation (8), the WNG can be expressed in closed-form in terms of G ( θ ) and g as g W =
G − 1 ( θ ) g .
Hence, WNG ascent is simply performed using the update equation θ k +
=
θ k + λg Wk .
We’ll see in Section 4 how to estimate WNG efﬁciently without storing or explicitly inverting the matrix G .
Next, we discuss the advantages of using WNG over other methods.
3. W HY USE THE W ASSERSTEIN N ATURAL G RADIENT ?
To illustrate the advantages of the WNG, we consider a simple setting where the objective is of the form L ( θ ) =
E q θ
[ ψ ( x )] , with q θ being a gaussian distribution.
The optimal solution in this example is a deterministic point mass located at the global optimum x (cid:63) of the function ψ ( x ) .
This situation arises systematically in the context of ES when using a gaussian noise distribution with learnable mean and variance.
Moreover, the optimal policy of a Markov Decision Processes (MDP) is necessarily deterministic (Puterman, 2010).
Thus, despite its simplicity, this example allows us to obtain closed-form expressions for all methods while capturing a crucial property in many RL problems (deterministic optimal policies) which, as we will see, results in differences in performance.
Wasserstein natural gradient vs 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 q θ , shows that the FNG is faster at ﬁrst, then slows down to reach a ﬁnal error of 10 − 4 .
On the other hand, WNG is slower at ﬁrst then transitions suddenly to an error of 10 − 8 .
The optimal solution being deterministic, the variance of the gaussian q θ needs to shrink to 0 .
In this case, the KL blows up, while the W 2 distance remains ﬁnite.
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 achieve the error of 10 − 8 .
Beyond this example, when the policy π θ is deﬁned only implicitly using a generative network, as in Tang & Agrawal (2019), the FNG and KL penalty are ill-deﬁned since π θ k and π θ k +1 might have non-overlapping supports.
However, the WNG remains well-deﬁned (see Arbel et al.
and allows for more ﬂexibility in representing policies, such as with behavioral embeddings.
Wasserstein
penalty vs
Wasserstein natural gradient The Wasserstein penalty Equation (4) encourages global proximity between updates q θ k .
For small values of the penalty parameter β , the method behaves like standard gradient descent (Figure 1 (a)).
As β increases, the penalty encourages more local updates and thus incorporates more information about the local geometry deﬁned by q θ .
In fact, it recovers the WNG direction (Theorem 2 of Appendix B) albeit with an inﬁnitely 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 β that needs to be maintained (Schulman et al., 2017).
All of these issues are avoided when directly using the WNG, as shown in Figure 1 (a), which performs the best and tolerates the widest range of step-sizes Figure 1 (f).
Moreover, when using the log-diagonal parameterization as in Figure 1 (d,a), the WNGD (in red) achieves an error of 1e-8, while W 2 -penalty achieves a larger error of order 1e-0 for various values of the β .
When using the diagonal parameterization instead, as shown in Figure 1 (e), both methods achieve a similar error of 1e-6.
This discrepancy in performance highlights the robustness of WNG to parameterization of the model.
Combining WNG and a Wasserstein penalty.
The global proximity encouraged by a W 2 penalty can be useful on its own, for instance, to explicitly guarantee policy improvement as in (Pacchiano et al., 2019, Theorem 5.1).
However, this requires estimating the W 2 at every iteration, which can be costly.
Using WNG instead of the usual gradient can yield more efﬁcient updates, thus reducing the number of time W 2 needs to be estimated.
The speed-up can be understood as performing secondorder optimization on the W 2 penalty since the WNG arises precisely from a second-order expansion of the W 2 distance, as shown in Section 3 (See also Example 2 in Arbel et al.
P OLICY O PTIMIZATION USING B EHAVIORAL G EOMETRY
We now present practical algorithms to exploit the behavioral geometry induced by the embeddings Φ .
We begin by describing how to efﬁciently estimate the WNG.
Efﬁcient 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
= [ a 0 , ..., a T ] and if actions are sampled from a gaussian with mean and variance parameterized by a neural network, as is often done in practice for real-valued actions.
Then X can be expressed as X =
B θ ( Z ) where B θ is a known function and Z is an input sample consisting in the concatenation of states [ s 0 , ..., s T ] 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 ∇ θ log q θ is available without an explicit re-parametrization trick.
The core idea is to approximate the functions f e j deﬁning G ( θ k ) in Equation (7) using a linear combinations of user-speciﬁed basis functions ( h 1 ( x ) , ..., h M ( x )) :
[Equation]
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-dependent using kernel methods.
More precisely,
1: Input Initial policy π θ2: for iteration k = 1 , 2 , ... do 3:
Obtain N rollouts { τ } Nn =1 of length T using policy π θ k 4:
Compute loss L ( θ k ) in a forward pass 5:
Compute gradient ˆ g k in the backward pass on L ( θ k ) 6:
Compute Behavioral embeddings { X n = Φ( τ n ) }
Nn = 7:
Compute WNG ˆ g Wk using Algorithm 3 with samples { X n } Nn =1 and gradient estimate ˆ g k .
Update policy using: θ k +
=
θ k
+ λ ˆ g Wk .
9: end for we use the same approach as in Arbel et al.
(2020), where we ﬁrst subsample M data-points Y m from a batch of N variables X n and M indices i m from { 1 , ..., d } where d is the dimension of X n .
Then, each basis can of the form h m ( x ) =
∂
i m K
( Y m , x ) where K is a positive semi-deﬁnite kernel, such as the gaussian kernel K ( x, y ) =
exp( − (cid:107) x − y (cid:107) 2 σ) .
This choice of basis allows us to provide guarantees for functions f j in terms of the batch size N and the number of basis points M (Arbel et al., 2020, Theorem 7).
Plugging-in each ˆ f j in the transport cost problem Equation (5) yields a quadratic problem of dimension M in the coefﬁcients α j :
[Equation]
where L is a square matrix of size M × M independent of the index j and J is a Jacobian matrix of shape M × p with rows given by J m,.
= ∇ θ E q θk [ h m ( X )] .
There are two expressions for J , depending on the applicability of the re-parametrization trick or the availability of the score
J m,.
= ˆ
[Equation]
Computing J can be done efﬁciently for moderate size M by ﬁrst computing a surrogate vector of
V of size M whose Jacobian recovers J using automatic differentiation software:
[Equation]
The optimal coefﬁcients α j are then simply expressed as α = L † J .
Plugging-in the optimal functions in the expression of the Wasserstein Information Matrix (Equation (7)), yields a low rank approximation of G of the form ˆ G =
J (cid:62)
L † J .
By adding a small diagonal perturbation matrix (cid:15)I , it is possible efﬁciently compute ( ˆ G + (cid:15)I ) −
ˆ
g using a generalized Woodbury matrix identity which yields an estimator for the Wasserstein Natural gradient
[Equation]
The pseudo-inverse is only computed for a matrix of size M .
Using the Jacobian-vector product,
Equation (12) can be computed without storing large matrices G as shown in Algorithm 3.
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 in Algorithm 1.
In its simplest form, one ﬁrst needs to compute the gradient ˆ g k of the objective L ( θ k ) using, for instance, the REINFORCE estimator computed using N trajectories τ n .
The trajectories are then used to compute the BEMs which are fed as input, along with the gradient ˆ g k to get an estimate of the WNG g Wk .
Finally, the policy can be updated in the direction of g Wk .
Algorithm 1 can also be used in combination with an explicit W 2 penalty to 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 in Section 3.1.
The standard gradient ˆ g k in Algorithm 1 is then simply replaced by the one computed in (Pacchiano et al., 2019, Algorithm 3).
In Section 5, we show that this combination, which we call behavior-guided WNPG (BG-WNPG), leads to the best overall performance.
Wasserstein Natural Evolution Strategies (WNES).
ES treats the total reward observed on a trajectory under policy π θ as a black-box function L ( θ ) (Salimans et al., 2017; Mania et al., 2018;
1: Input Initial policy π θ 0 , α > 0 , δ ≤2: for iteration k = 1 , 2 , ... do 3:
Sample (cid:15) 1 , . .
, (cid:15) n ∼ N (0 , I ) .
Perform rollouts { τ n } Nn − 1 of length T using the perturbed parameters { (cid:101) θ n = θ k + σ(cid:15) n } Nn = and compute behavioral embeddings { X n = Φ( τ n ) } Nn = 5:
Compute gradient estimate of L ( (cid:101) θ n ) using Equation (13) and trajectories { τ n } Nn =1 .
Compute Jacobian matrix J appearing in Algorithm 3 using Equation (14).
Compute WNG ˆ g Wk using Algorithm 3, with samples { X n } Ni =1 and computed ˆ g k and J .
Update policy using Equation (15).
9: end for
Choromanski et al., 2020).
Evaluating it under N policies whose parameters (cid:101) θ n are gaussian perturbations centered around θ k and with variance σ can give an estimate of the gradient of L ( θ k ) :
[Equation]
Instead of directly updating the policy using Equation (13), it is possible to encourage either proximity or diversity in behavior using the embeddings X n = Φ( τ n ) of the trajectories τ n generated for each perturbed policy π (cid:101) θ n .
Those embeddings can be used as input to Algorithm 3 (see appendix), along with Equation (13) to estimate the ˆ g Wk , which captures similarity in behavior.
The algorithm remains unchanged except for the estimation of the Jacobian J of Equation (10) which becomes
[Equation]
The policy parameter can then be updated using an interpolation between ˆ g k and the WNG ˆ g Wk , i.e.,
[Equation]
with δ ≤ 1 that can also be negative.
Positive values for δ encourage proximity in behavior, the limit case being δ = 1 where a full WNG step is taken.
Negative values encourage repulsion and therefore need to compensated by ˆ g k to ensure overall policy improvement.
Algorithm 2 summarizes the whole procedure, which can be easily adapted from existing ES implementations by calling a variant of Algorithm 3.
In particular, it can also be used along with an explicit W 2 penalty, in which case the proposed algorithm in Pacchiano et al.
(2019) is used to estimate the standard gradient ˆ g k of the penalized loss.
Then the policy is updated using Equation (15) instead of ˆ g k .
We refer to this approach as behavior-guided WNES (BG-WNES).
E XPERIMENTS
We now test the performance of our estimators for both policy gradients (PG) and evolution strategies (ES) against their associated baseline methods.
We show that in addition to an improved computational efﬁciency, our approach can effectively utilize the geometry induced by a Wasserstein penalty to improve performance, particularly when the optimization problem is ill-conditioned.
Further experimental details can be found in the appendix, and our code is available online 1 .
Policy Gradients.
We ﬁrst 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, as also demonstrated in the work of Pacchiano et al.
Our method (WNPG) matches or exceeds ﬁnal performance of BGPG on all tasks.
Moreover, combining both (BG-WNPG) produces the largest gains on all [Figure] None PPO (Clip)
[Figure]
environments.
Final mean rewards are reported in Table 1.
It is also important to note that WNGbased methods appear to offer the biggest advantage on tasks where initial progress is difﬁcult.
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 signiﬁcant that WNG methods produce the greatest improvement on the environments with the poorest conditioning.
This is consistent with the ﬁndings 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 a useful heuristic for gauging when WNGbased methods are most useful for a given problem.
[Figure]
Evolution Strategies To test our estimator for WNES, as well as BG-WNES, we applied our approach to the environment introduced by
Pacchiano et al.
(2019), designed to test the ability of behavior-guided learning to succeed despite deceptive rewards.
During the task, the agent receives a penalty proportional to its distance from a goal, but a wall is placed directly in the agent’s path (Figure 7).
This barrier induces a local maximum in the objective—a na¨ıve agent will simply walk directly towards the goal and get stuck at the barrier.
The idea is that the behavioral repulsion fostered by applying a positive coefﬁcient to the Wasserstein penalty ( β > 0 ) will encourage the agent to seek novel policies, helping it to eventually circumvent the wall.
As in Pacchiano et al.
(2019), we test two agents, a simple point and a quadruped.
We then compare our method with vanilla ES as described by Salimans et al.
(2017), ES with gradient norm clipping, BGES (Pacchiano et al., 2019), and NSR-ES (Conti et al., 2018a).
In Figure 4, we can see that WNES and BG-WNES improve over the baselines for both agents.
To test that the improvement shown by BG-WNES wasn’t simply a case of additional “repulsion” supplied by the WNG to BGES, we also tested BGES with an increased β = 0 .
75 , compared to the default of 0 . 5 .
This resulted in a decrease in performance, attesting to the unique beneﬁt provided by the WNES estimator.
Computational Efﬁciency We deﬁne the computational efﬁciency of an algorithm as the rate with which it accumulates reward relative to its runtime.
To test the computational efﬁciency of our [Figure] [Figure] approach, we plotted the total reward divided by wall clock time obtained by each agent for each task (Fig. 5).
Methods using a WNG estimator were the most efﬁcient 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 efﬁciency, despite having the highest absolute performance.
C ONCLUSION
Explicit regularization using divergence measures between policy representations has been a common theme in recent work on policy optimization for RL.
While prior works have previously focused on the KL divergence, Pacchiano et al.
(2019) showed that a Wasserstein regularizer over behavioral distributions provides a powerful alternative framework.
Both approaches implicitly deﬁne a form of natural gradient, depending on which divergence measure is chosen.
Through the introduction of WNPG and WNES, we 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 on the WNG to problems where the reparameterization trick is not available, as well as to black-box methods like ES.
Moreover, using the WNG in conjunction with a WD penalty allows the WNG to take advantage of the local geometry induced by the regularization, further improving performance.
We also provide a novel comparison between the WNG and FNG, showing that the former has signiﬁcant advantages on certain problems.
We believe this framework opens up a number of avenues for future work.
Developing a principled way to identify useful behavioral embeddings for a given RL task would allow to get the highest beneﬁt form WNPG and WNES.
From a theoretical perspective, it would be useful to characterize the convergence boost granted by the combination of explicit regularization and the corresponding natural gradient approach.
Acknowledgments The authors would like to thank Jack Parker-Holder for sharing his code for
BGPG and BGES, as well as colleagues at Gatsby for useful discussions.
A B ACKGROUND A. P OLICY O PTIMIZATION
An agent interacting with an environment form a system that can be described by a state variable s belonging to a state space S .
In the Markov Decision Process (MDP) setting, the agent can interact with the environment by taking an action a from a set of possible actions A given the current state s of the system.
As a consequence, the system moves to a new state s (cid:48) according to a probability transition function P ( s (cid:48) | a, s ) which describes the probability of moving to state s (cid:48) given the previous state s and action a .
The agent also receives a partial reward r which can be expressed as a possibly randomized function of the new state s (cid:48) , r = r
( s (cid:48) ) .
The agent has access to a set of possible policies π θ ( a | s ) parametrized by θ ∈ R p and that generates an action a given a current state s .
Thus, each policy can be seen as a probability distribution conditioned a state s .
Using the same policy induces a whole trajectory of state-action-rewards τ = ( s t , a t , r t ) t ≥ 0 which can be viewed as a sample from a trajectory distribution P θ deﬁned over the space of possible trajectories τ .
Hence, for a given random trajectory τ induced by a policy π θ , the agent receives a total discounted reward
R ( τ ) :
= (cid:80) ∞
t =1 γ t − 1 r ( s t ) with discount factor 0 < γ < 1 .
This allows to deﬁne the value function as the expected total reward conditioned on a particular initial state s :
[Equation]
When the gradient of the score function ∇ log π θ ( a | s ) is available, the policy gradient theorem allows us to express the gradient of R ( θ ) :
[Equation]
where the expectation is taken over trajectories τ under P θ and A θ ( s, a ) represents the advantage function which can be expressed in terms of the value function V θ ( s ) in terms of
[Equation]
The agent seeks an optimal policy π θ (cid:63) that maximizes the expected total reward under the trajectory distribution:
[Equation]
B W ASSERSTEIN N ATURAL GRADIENT
Connection to the Fisher natural gradient and proximal methods.
Both WNG and FNG are obtained from a proximity measure between probability distributions:
Proposition 2 Let D ( θ, θ (cid:48) ) be either the KL-divergence KL ( π θ , π θ (cid:48) ) or the Wasserstein-2 distance between the behavioral distributions W 2 ( q θ , q θ (cid:48) ) and let g D be either the FNG g F or WNG g W ,
[Equation]
τ then
[Equation]
Equation (18) simply states that the both WNG and FNG arise as limit cases of penalized objectives provided the strength of the penalty β diverges to inﬁnity and the step-size is shrank proportionally to β − 1 .
An additional global rescaling by β of the total objective prevents it from collapsing to 0 .
Intuitively, performing a Taylor expansion of Equation (18) recovers an equation similar to
Equation (8).
Equation (18) shows that using a 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-scaling.
This also informally shows why both natural gradients are said to be invariant to re-parametrization (Arbel et al., 2020, Proposition 1), since both KL and W 2 remains unchanged if q θ is parameterized in a different way.
C A LGORITHM FOR ESTIMATING WNG
1: Input mini-batch of samples { X n } Nn =1 distributed according to q θ , gradient direction ˆ g , basis functions h 1 , ..., h M , regularization parameter (cid:15) .
2: Output Wasserstein Natural gradient ˆ g W
: Compute a matrix C of shape M × Nd using C m, ( n,i ) ←
∂ i h m ( X n ) .
4: Compute similarity matrix L ← 1 N CC T . 5: Compute surrogate vector V using Equation (11).
6: for iteration = 1 , 2 , ...
M do 7: Use automatic differentiation on V m to compute Jacobian matrix J in Equation (
D A DDITIONAL E XPERIMENTAL D ETAILS D. P OLICY G RADIENT T ASKS
We conserve all baseline and shared hyperparameters used by Pacchiano et al.
More precisely, for each task we ran a hyperparameter sweep over learning rates in the set { 1e-5 , 5e-5 , 1e-4 , 3e-4 } , and used the concatenation-of-actions behavioral embedding Φ( τ ) =
[ a 0 , a 1 , . .
. , a T ] with the base network implementation the same as Dhariwal et al.
The WNG hyperparameters were also left the same as in Arbel et al.
Speciﬁcally, the number of basis points was set as M = 5 , the reduction factor was bounded in the range [0 . 25 , 0 . 75] , and (cid:15) ∈
[ 1e-10 , 1e5 ] .
D. E VOLUTION S TRATEGIES T ASKS
As with the policy gradient tasks, we conserved all baseline and shared hyperparameters used by
Pacchiano et al.
Speciﬁcally, for the point task, we set the learning rate to be η = 0 .
1 , the standard deviation of the noise to be σ = 0 .
01 , the rollout length H was 50 time steps, and the behavioral embedding function to be the last state Φ( τ ) =
s H .
For the quadruped task we set η = 0 .
02 , σ = 0 .
02 , H = 400 , and Φ( τ ) =
(cid:80)
Ht
=
r t (cid:16)(cid:80) ti =
e i (cid:17) (reward-to-go encoding; see
Pacchiano et al.
for more details).
Both tasks used 1000-dimensional random features and embeddings from the n = 2 previous policies to compute the WD.
For WNG, the same hyperparameters were used as in the policy gradient tasks.
D. E XPERIMENTAL SETTING OF F IGURE
The Objective We consider a function ψ ( x ) is the sum of sinc functions over all dimensions of x ∈
[Equation]
[Equation]
Such function is highly non-convex and admits multiple bad local minima with the global minimum of ψ ( x ) reached for x (cid:63) = 0 .
However, we do not make use of this information during optimization.
To alleviate the non-convexity of this loss, we consider a gaussian relaxation objective L ( θ ) obtained by taking the expectation of ψ ( x ) over the 100 dimensional vector x w.r.t.
to a gaussian q θ with parameter vector θ .
Thus the objective function to be optimized is a function of θ :
[Equation]
The parameter vector θ is of the form θ = ( µ, v ) , where µ is the mean of the gaussian q θ and v is a vector in R 100 parameterizing the covariance matrix Σ of the gaussian q θ .
We will later consider two parameterizations for the covariance matrix.
The minimal value of L ( θ ) is reached when the gaussian q θ is degenerate with Σ = 0 and mean µ = x (cid:63) = 0 .
Hence, the mean parameter of the global minimum of L ( θ ) recover the global optimum of ψ .
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-diagonal .
For the diagonal parameterization the Covariance Σ ii = v i and for the log-diagonal we set Σ ii = exp(2 v i ) .
Optimization methods We consider different optimization methods using the same objective
L ( θ ) .
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 as ∇ W L ( θ ) for (WNG) and ∇ F L ( θ ) for FNG and express them in terms of the euclidean/standard gradient ∇L ( θ ) :
• Diagonal parameterization: –
[Equation] [Equation]
• Log-diagonal parameterization: –
[Equation] [Equation]
Training details Training is up to 4000 gradient iterations, with λ = .
9 and β = .
1 unless they are varied.