Efficient Wasserstein Natural Gradients for Reinforcement Learning

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 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 . 1 INTRODUCTION . Defining efficient 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 & Popović , 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 define 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 efficient algorithms . Pacchiano et al . ( 2019 ) provide a first 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 defined 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 , 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 , 2018a ; b ; 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 a 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-parameterization trick is not applicable , but only the score function of the model is available . ∗Denotes equal contribution . Correspondence : ted @ gatsby.ucl.ac.uk . 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 find that such WNG-based methods are especially useful on tasks in which initial progress is difficult . 3- Finally , we demonstrate , to 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 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 . 2 BACKGROUND . 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 flavors , depending on whether training is on-policy ( left ) or off-policy ( right ) : L ( θ ) = Ê [ log πθ ( at|st ) Ât ] , or L ( θ ) = Ê [ πθ ( at|st ) πθk ( at|st ) Ât ] . ( 1 ) The expectation Ê is an empirical average over N trajectories τi = ( si1 , ai1 , ri1 , ... , siT , aiT , riT ) 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 Ât = rt + γV ( st+1 ) − V ( st ) ( 2 ) 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 first 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 : maximizeθ L ( θ ) − βÊ [ KL ( πθk ( ·|st ) , πθ ( ·|st ) ) ] . ( 3 ) 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 define 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 final state of a trajectory ( Φ ( τ ) = sT ) or its concatenated actions ( Φ ( τ ) = [ a0 , . . . , aT ] ) . Proximity between two policies πθ and πθ′ is then measured using the Wasserstein distance between their behavior distributions qθ and qθ′ . Although , the KL could also be used in some cases , the Wasserstein distance has the advantage of being well-defined 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 : maximizeθ L ( θ ) − β 2 W2 ( qθk , qθ ) , ( 4 ) 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 . ( 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 β 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 . ( 2017 ) . For a finite 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 efficient 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 efficient updates , thus converging faster than simply optimizing Equation ( 4 ) using standard gradients .

Summary: This 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.