Differentiable Trust Region Layers for Deep Reinforcement Learning

1 INTRODUCTION . Deep reinforcement learning has shown considerable advances in recent years with prominent application areas such as games ( Mnih et al. , 2015 ; Silver et al. , 2017 ) , robotics ( Levine et al. , 2015 ) , and control ( Duan et al. , 2016 ) . In policy search , policy gradient ( PG ) methods have been highly successful and have gained , among others , great popularity ( Peters & Schaal , 2008 ) . However , often it is difficult to tune learning rates for vanilla PG methods , because they tend to reduce the entropy of the policy too quickly . This results in a lack of exploration and , as a consequence , in premature or slow convergence . A common practice to mitigate these limitations is to impose a constraint on the allowed change between two successive policies . Kakade & Langford ( 2002 ) provided a theoretical justification for this in the approximate policy iteration setting . Two of the arguably most favored policy search algorithms , Trust Region Policy Optimization ( TRPO ) ( Schulman et al. , 2015a ) and Proximal Policy Optimization ( PPO ) ( Schulman et al. , 2017 ) , follow this idea using the KullbackLeibler divergence ( KL ) between successive policies as a constraint . We propose closed-form projections for Gaussian policies , realized as differentiable neural network layers . These layers constrain the change in successive policies by projecting the updated policy onto trust regions . First , this approach is more stable with respect to what Engstrom et al . ( 2020 ) refer to as code-level optimizations than other approaches . Second , it comes with the benefit of imposing constraints for individual states , allowing for the possibility of state-dependent trust regions . This allows us to constrain the state-wise maximum change of successive policies . In this we differ from previous works , that constrain only the expected change and thus can not rely on exact guarantees of monotonic improvement . Furthermore , we propose three different similarity measures , the KL ∗Correspondence to fabian.otto @ bosch.com divergence , the Wasserstein L2 distance , and the Frobenius norm , to base our trust region approach on . The last layer of the projected policy is now the the trust region layer which relies on the old policy as input . This would result in a ever-growing stack of policies , rendering this approach clearly infeasible . To circumvent this issue we introduce a penalty term into the reinforcement learning objective to ensure the input and output of the projection stay close together . While this still results in an approximation of the trust region update , we show that the trust regions are properly enforced . We also extend our approach to allow for a controlled evolution of the entropy of the policy , which has been shown to increase the performance in difficult exploration problems ( Pajarinen et al. , 2019 ; Akrour et al. , 2019 ) . We compare and discuss the effect of the different similarity measures as well as the entropy control on the optimization process . Additionally , we benchmark our algorithm against existing methods and demonstrate that we achieve similar or better performance . 2 RELATED WORK . Approximate Trust Regions . Bounding the size of the policy update in policy search is a common approach . While Kakade & Langford ( 2002 ) originally focused on a method based on mixing policies , nowadays most approaches use KL trust regions to bound the updates . Peters et al . ( 2010 ) proposed a first approach to such trust regions by formulating the problem as a constraint optimization and provided a solution based on the dual of that optimization problem . Still , this approach is not straightforwardly extendable to highly non-linear policies , such as neural networks . In an attempt to transfer those ideas to deep learning , TRPO ( Schulman et al. , 2015a ) approximates the KL constraint using the Fisher information matrix and natural policy gradient updates ( Peters & Schaal , 2008 ; Kakade , 2001 ) , along with a backtracking line search to enforce a hard KL constraint . Yet , the resulting algorithm scales poorly . Thus , Schulman et al . ( 2017 ) introduced PPO , which does not directly enforce the KL trust region , but clips the probability ratio in the importance sampling objective . This allows using efficient first-order optimization methods while maintaining robust training . However , Engstrom et al . ( 2020 ) and Andrychowicz et al . ( 2020 ) recently showed that implementation choices are essential for achieving state-of-the-art results with PPO . Code-level optimizations , such as reward scaling as well as value function , observation , reward , and gradient clipping , can even compensate for removing core parts of the algorithm , e. g. the clipping of the probability ratio . Additionally , PPO heavily relies on its exploration behavior and might get stuck in local optima ( Wang et al. , 2019 ) . Tangkaratt et al . ( 2018 ) use a closed-form solution for the constraint optimization based on the method of Lagrangian multipliers . They , however , require a quadratic parametrization of the Q-Function , which can limit the performance . Pajarinen et al . ( 2019 ) introduced an approach based on compatible value function approximations to realize KL trust regions . Based on the reinforcement learning as inference paradigm ( Levine , 2018 ) , Abdolmaleki et al . ( 2018 ) introduced an actor-critic approach using an Expectation-Maximization based optimization with KL trust regions in both the E-step and M-step . Song et al . ( 2020 ) proposed an on-policy version of this approach using a similar optimization scheme and constraints . Projections for Trust Regions . Akrour et al . ( 2019 ) proposed Projected Approximate Policy Iteration ( PAPI ) , a projection-based solution to implement KL trust regions . Their method projects an intermediate policy , that already satisfies the trust region constraint , onto the constraint bounds . This maximizes the size of the update step . However , PAPI relies on other trust region methods to generate this intermediary policy and can not operate in a stand-alone setting . Additionally , the projection is not directly part of the policy optimization but applied afterwards , which can result in sub-optimal policies . In context of computational complexity , both , TRPO and PAPI , simplify the constraint by leveraging the expected KL divergence . Opposed to that , we implement the projections as fully differentiable network layers and directly include them in the optimization process . Additionally , our projections enforce the constraints per state . This allows for better control of the change between subsequent policies and for state-dependent trust regions . For the KL-based projection layer we need to resort to numerical optimization and implicit gradients for convex optimizations ( Amos & Kolter , 2017 ; Agrawal et al. , 2019 ) . Thus , we investigate two alternative projections based on the Wasserstein L2 and Frobenius norm , which allow for closed form solutions . Both , Wasserstein and Frobenius norm , have found only limited applications in reinforcement learning . Pacchiano et al . ( 2020 ) use the Wasserstein distance to score behaviors of agents . Richemond & Maginnis ( 2017 ) proposed an alternative algorithm for bandits with Wasserstein based trust regions . Song & Zhao ( 2020 ) focus on solving the trust region problem for distributional policies using both KL and Wasserstein based trust regions for discrete action spaces . Our projections are applicable independently of the underlying algorithm and only assume a Gaussian policy , a common assumption for continuous action spaces . Several authors ( Dalal et al. , 2018 ; Chow et al. , 2019 ; Yang et al. , 2020 ) used projections as network layers to enforce limitations in the action or state space given environmental restrictions , such as robotic joint limits . Entropy Control . Abdolmaleki et al . ( 2015 ) introduced the idea of explicitly controlling the decrease in entropy during the optimization process , which later was extended to deep reinforcement learning by Pajarinen et al . ( 2019 ) and Akrour et al . ( 2019 ) . They use either an exponential or linear decay of the entropy during policy optimization to control the exploration process and escape local optima . To leverage those benefits , we embed this entropy control mechanism in our differentiable trust region layers . 3 PRELIMINARIES AND PROBLEM STATEMENT . We consider the general problem of a policy search in a Markov Decision Process ( MDP ) defined by the tuple ( S , A , T , R , P0 , γ ) . We assume the state space S and action space A are continuous and the transition probabilities T : S × A × S → [ 0 , 1 ] describe the probability transitioning to state st+1 ∈ S given the current state st ∈ S and action at ∈ A . We denote the initial state distributions as P0 : S → [ 0 , 1 ] . The reward returned by the environment is given by a function R : S × A → R and γ ∈ [ 0 , 1 ) describes the discount factor . Our goal is to maximize the expected accumulated discounted reward Rγ = ET , P0 , π [ ∑∞ t=0 γ tR ( st , at ) ] . To find the optimal policy , traditional PG methods often make use of the likelihood ratio gradient and an importance sampling estimator . Moreover , instead of directly optimizing the returns , it has been shown to be more effective to optimize the advantage function as this results in an unbiased estimator of the gradient with less variance max θ Ĵ ( πθ , πθold ) = max θ E ( s , a ) ∼πθold [ πθ ( a|s ) πθold ( a|s ) Aπθold ( s , a ) ] , ( 1 ) where Aπ ( s , a ) = E [ Rγ |s0 = s , a0 = a ; π ] − E [ Rγ |s0 = s ; π ] describes the advantage function , and the expectation is w.r.t πθold , i.e . s ′ ∼ T ( ·|s , a ) , a ∼ πθold ( ·|s ) , s0 ∼ P0 ( s0 ) , s ∼ ρπθold where ρπθold is a stationary distribution of policy πθold . The advantage function is commonly estimated bygeneralized advantage estimation ( GAE ) ( Schulman et al. , 2015b ) . Trust region methods use additional constraints for the given objective . Using a constraint on the maximum KL over the states has been shown to guarantee monotonic improvement of the policy ( Schulman et al. , 2015a ) . However , since all current approaches do not use a maximum KL constraint but an expected KL constraint , the guarantee of monotonic improvement does not hold exactly either . We are not aware of such results for the W2 distance or the Frobenius norm . For our projections we assume Gaussian policies πθold ( at|st ) = N ( at|µold ( st ) , Σold ( st ) ) and πθ ( at|st ) = N ( at|µ ( st ) , Σ ( st ) ) represent the old as well as the current policy , respectively . We explore three trust regions on top of Equation 1 that employ different similarity measures between old and new distributions , more specifically the frequently used reverse KL divergence , the Wasserstein L2 distance , and the Frobenius norm . Reverse KL Divergence . The KL divergence between two Gaussian distributions with means µ1 and µ2 and covariances Σ1 and Σ2 can generally be written as KL ( { µ1 , Σ1 } ‖ { µ2 , Σ2 } ) = 1 2 [ ( µ2 − µ1 ) TΣ−12 ( µ2 − µ1 ) + log |Σ2| |Σ1| + tr { Σ−12 Σ1 } − d ] , where d is the dimensionality of µ1 , µ2 . The KL uses the Mahalanobis distance to measure the similarity between the two mean vectors . The difference of the covariances is measured by the difference in shape , i.e. , the difference in scale , given by the log ratio of the determinants , plus the difference in rotation , given by the trace term . Given the KL is non-symmetric , it is clearly not a distance , yet still a frequently used divergence between distributions . We will use the more common reverse KL for our trust region , where the first argument is the new policy and the second is the old policy . Wasserstein Distance . The Wasserstein distance is a distance measure based on an optimal transport formulation , for more details see Villani ( 2008 ) . The Wasserstein-2 distance for two Gaussian distributions can generally be written as W2 ( { µ1 , Σ1 } , { µ2 , Σ2 } ) = |µ1 − µ2|2 + tr ( Σ1 + Σ2 − 2 ( Σ 1/2 2 Σ1Σ 1/2 2 ) 1/2 ) . A key difference to the KL divergence is that the Wasserstein distance is a symmetric distance measure , i. e. , W2 ( q , p ) = W2 ( p , q ) . Our experiments also revealed that it is beneficial to measure the W2 distance in a metric space defined by the covariance of the old policy distribution , denoted here as Σ2 , as the distance measure is then more sensitive to the data-generating distribution . The W2 distance in this metric space reads W2 , Σ2 ( { µ1 , Σ1 } , { µ2 , Σ2 } ) = ( µ2 − µ1 ) TΣ−12 ( µ2 − µ1 ) + tr ( Σ−12 Σ1 + I− 2Σ−12 ( Σ 1/2 2 Σ1Σ 1/2 2 ) 1/2 ) . Frobenius Norm . The Frobenius norm is a matrix norm and can directly be applied to the difference of the covariance matrices of the Gaussian distributions . To measure the distance of the mean vectors , we will , similar to the KL divergence , employ the Mahalanobis distance as this empirically leads to an improved performance in comparison to just taking the squared distance . Hence , we will denote the following metric as Frobenius norm between two Gaussian distributions F ( { µ1 , Σ1 } , { µ2 , Σ2 } ) = ( µ2 − µ1 ) TΣ−12 ( µ2 − µ1 ) + tr ( ( Σ2 − Σ1 ) T ( Σ2 − Σ1 ) ) . The Frobenius norm also constitutes a symmetric distance measure .

Summary: The paper proposes a way to impose trust region restrictions via projections when doing policy optimisation in Reinforcement Learning. The projections have a closed form and enforce a trust region for each state individually. The authors propose three types of projections based on Frobenius, Wasserstein distances and KL divergence. They compare them to the existing methods (PPO, PAPI) and provide some insights about their behaviour.