Memory-Constrained Policy Optimization

1 INTRODUCTION . Reinforcement learning ( RL ) combined with neural networks is the current workhorse in machine learning . Using neural networks to approximate value and policy functions enables classical approaches such as Q-learning and policy gradient to achieve promising results on many challenging problems such as Go , Atari games and robotics ( Silver et al. , 2017 ; Mnih et al. , 2015 ; Lillicrap et al. , 2016 ; Mnih et al. , 2016 ) . Compared to Deep Q-learning , deep policy gradient ( PG ) methods are often more flexible and applicable to both discrete and continuous action problems . However , these methods tend to suffer from high sample complexity and training instability since the gradient may not accurately reflect the policy gain when the policy changes substantially ( Kakade & Langford , 2002 ) . This is exacerbated for deep policy networks where numerous parameters need to be optimized and small updates in parameter space can lead to huge changes in policy space . To address this issue , one solution is to regularize each policy update by restricting the Kullback– Leibler ( KL ) divergence between the new policy and the previous one , which can guarantee monotonic policy improvement ( Schulman et al. , 2015a ) . However , jointly optimizing the approximate advantage function and the KL term does not work in practice . Therefore , Schulman et al . ( 2015a ) proposed Trust Region Policy Optimization ( TRPO ) to constrain the new policy within a KL divergence radius , which requires second-order gradients . Alternatives such as Proximal Policy Optimization ( PPO ) ( Schulman et al. , 2017 ) use a simpler first-order optimization with adaptive KL or clipped surrogate objective while still maintaining the reliable performance of TRPO . Recent methods recast the problem through a new lens using Expectation-Maximization or Mirror Descent Optimization , and this also results in first-order optimization with KL divergence term in the loss function ( Abdolmaleki et al. , 2018 ; Song et al. , 2019 ; Yang et al. , 2019 ; Tomar et al. , 2020 ) . An issue with the above methods is that the previous policy used to restrict the new policy may be suboptimal and thus unreliable in practice . For example , due to stochasticity and imperfect approximations , consider that the new policy may fall into a local optimum even under trust-region optimizations . Then in the next update , this policy will become the “ previous ” policy , and will continue pulling the next policy to stay in the local optimum , thus slowing down the training progress . For on-policy methods using mini-batch updates like PPO , the situation is more complicated as the “ previous ” policy is defined as the old policy that was used to collect data , which can be either very far or close to the current policy . There is no guarantee that the old policy defines a reasonable trust region for regulating the new policy . In this paper , we propose a novel constrained policy iteration procedure , dubbed as MemoryConstrained Policy Optimization ( MCPO ) , wherein a virtual policy representing past policies takes part in regularizing each policy update . The virtual and the old policy together form two trust regions that attract the new policy by minimizing two KL divergences . The virtual policy is designed to complement the old policy , attracting the new policy more when the old policy performs badly . As such , we assign different coefficient weights to the two KL terms . The coefficient weights are computed dynamically based on the performance of the two policies ( the higher performer yields higher coefficient weights ) . To create the virtual policy , we maintain a memory buffer of past policies , from which we build a mixture of policies . In particular , we use a neural network–named the attention network , which takes the context surrounding the current , the old and the virtual policy , to generate the attention weights to each policy in the memory . Then , we use the attention weights to perform a convex combination of the parameters of each policy , forming the virtual policy . The attention network is optimized to maximize the approximate expected advantages of the virtual policy . To train our system , we jointly optimize the policy and attention networks , alternating between sampling data from the policy and updating the networks in a mini-batch manner . We verify our proposed MCPO through a diverse set of experiments and compare ours with the performance of recent constrained policy optimization baselines . In our experiment on classical control tasks , amongst tested models , MCPO shows minimal sensitivity to hyperparameter changes , consistently achieving good performance across tasks and hyperparameters . Our testbed on 6 Mujoco tasks shows that MCPO with a big policy memory consistently outperforms others where the attention network plays an important role . We also demonstrate MCPO ’ s capability of learning efficiently on sparse reward and high-dimensional problems such as navigation and Atari games . Finally , our ablation study highlights the necessity of each component in MCPO . 2 BACKGROUND : POLICY OPTIMIZATION WITH TRUST REGION . In this section , we briefly review some fundamental constrained policy optimization approaches . A general idea is to force the new policy πθ to be close to a recent policy πθold . In this paper , we refer to a policy as its parameters ( i.e . policy θ means policy πθ ) . Conservative Policy Iteration ( CPI ) The method starts with a basic objective of policy gradient algorithms , which is to maximize the expected advantage Ât . LCPI ( θ ) = Êt [ πθ ( at|st ) πθold ( at|st ) Ât ] where the advantage Ât is a function of returns collected from ( st , at ) by using πθold ( see Appendix A.2 ) and Êt [ · ] indicates the empirical average over a finite batch of data . To constrain policy updates , the new policy is a mixture of the old and the greedy policy : θ̃ = argmax LCPI ( θ ) . That is , θ = αθold + ( 1 − α ) θ̃ where α is the mixture hyperparameter ( Kakade & Langford , 2002 ) . As the data is sampled from previous iteration ’ s policy θold , the objective needs importance sampling estimation . Hereafter , we denote πθ ( at|st ) πθold ( at|st ) as ratt ( θ ) for short . KL-Regularized Policy Optimization ( with fixed or adaptive KL coefficient ) Another way to enforce the constraint is to jointly maximize the expected advantage and minimize KL divergence between the new and old policy , which ensures monotonic improvement ( Schulman et al. , 2015a ) . LKLPO ( θ ) = Êt [ ratt ( θ ) Ât − βKL [ πθold ( ·|st ) , πθ ( ·|st ) ] ] where β is a hyperparameter that controls the degree of update conservativeness , which can be fixed ( KL Fixed ) or changed ( KL Adaptive ) during training ( Schulman et al. , 2017 ) . Trust Region Policy Optimization ( TRPO ) The method optimizes the expected advantage with hard constraint ( Schulman et al. , 2015a ) . This is claimed as a practical implementation that is less conservative than the theoretically justified algorithm using KL regularizer mentioned above . LTRPO ( θ ) = Êt [ ratt ( θ ) Ât ] st δ ≥ KL [ πθold ( ·|st ) , πθ ( ·|st ) ] where δ is the KL constraint radius . Proximal Policy Optimization ( PPO ) PPO is a family of constrained policy optimization , which uses first-order optimization and mini-batch updates including KL Adaptive and clipped PPO . In this paper , we use PPO to refer to the method that limits the change in policy by clipping the loss function ( clipped PPO ) ( Schulman et al. , 2017 ) . LPPO ( θ ) = Êt [ min ( ratt ( θ ) Ât , clip ( ratt ( θ ) , 1− , 1 + ) Ât ) ] where is the clip hyperparameter . In the above equations , θ is the currently optimized policy , which is also referred to as the current policy . θold represents a past policy , which can be one step before the current policy or the last policy used to interact with the environment . In either case , the rule to decide θold is fixed throughout training . If for some reason , θold is not optimal , it is unavoidable that the following updates will be negatively impacted . We will address this issue in the next section . 3 MEMORY-CONSTRAINED POLICY OPTIMIZATION ( MCPO ) . 3.1 TWO KL DIVERGENCE CONSTRAINTS . In trust-region methods with mini-batch updates such as PPO , the old policy θold is often chosen as the last policy that is used to collect observations from the environment . Before the next environment interaction , this old policy is fixed across policy updates , and can be one or many steps before the current policy depending on the mini-batch size and the number of update loops . This can be detrimental to the optimization if this old policy is poor in quality , forcing the following updates to be close to a poor solution . To tackle this issue , we propose to constrain the new policy not only to the policy θold , but also to a changeable policy that is representative of many past policies . Let ψ denote the virtual policy that represents the history of policies . ψ is dynamically computed based on the past policies using attention mechanism ( see Sec . 3.2 ) . In contrast to the fixed θold , depending on the attention weights , ψ can represent a further or closer checkpoint to the current policy than θold . We use both ψ and θold to construct the objective function as follows , L1 ( θ ) = Êt [ ratt ( θ ) Ât ] − βÊt [ ( 1− αt ( ·|st ) ) KL [ πθold ( ·|st ) , πθ ( ·|st ) ] ] ( 1 ) − βÊt [ αt ( ·|st ) KL [ πψ ( ·|st ) , πθ ( ·|st ) ] ] where αt ( ·|st ) is the coefficient weight resembling a forget gate , and β is the scaling coefficient of the KL constraint terms . In this paper , the expectation is estimated by taking average over t in a mini-batch of sampled data . The forget gate determines how much the new policy should forget the virtual policy from the memory and focus on the θold . Intuitively , if the virtual policy is better than the old policy , the new policy should be kept close to the virtual policy and vice versa . Hence , αt = e Rt ( ψ ) eRt ( ψ ) +eRt ( θold ) where Rt ( ψ ) measures the performance of the policy ψ , which can be estimated by weighted importance sampling . That is Rt ( ψ ) = ratt ( ψ ) Ât . Besides deciding which trust region the new policy should rely on , we dynamically weigh the whole KL terms via adjusting β . Using a fixed threshold dtarg to change β ( e.g in KL Adaptive , if KL [ πθold ( ·|st ) , πθ ( ·|st ) ] > dtarg , increase β ( Schulman et al. , 2017 ) ) showed limited performance since dtarg should vary depending on the current learning . We instead make use of ψ as a reference for selecting β . Let d ( a , b ) = Êt [ KL [ πa ( ·|st ) , πb ( ·|st ) ] ] denote the “ distance ” between 2 policies πa and πb , we propose a switching-β rule as follows , { β = βmax if d ( θold , θ ) > d ( θold , ψ ) β = βmin otherwise ( 2 ) The intuition is that we encourage the enforcement of the constraint when πθ is too far from πθold using the distance between πθold and πψ as a reference . Algorithm 1 Memory-Constrained Policy Optimization ( for 1 actor ) . Require : A policy bufferM , an initial policy πθold . T , K , B are the learning horizon , number of update epochs , and batch size , respectively . 1 : Initialize ψold ← θold , θ ← θold 2 : for iteration = 1 , 2 , ... do 3 : Run policy πθold in environment for T timesteps . Compute advantage estimates Â1 , ... , ÂT 4 : for epoch = 1 , 2 , ... K do 5 : for batch = 1 , 2 , ... T/B do 6 : Compute ψ ( Eq . 4 ) using ψold , θ , θold , then optimize LMCPO wrt θ ( Eq . 5 ) 7 : if d ( θ , ψ ) > d ( θold , ψ ) then add θ toM 8 : if |M| > N then remove the last item fromM 9 : ψold ← ψ 10 : end for 11 : end for 12 : θold ← θ 13 : end for

Summary: This paper proposes a new approach to regularize on-policy methods utilizing two soft trust region terms. The first term encourages the new policy to stay close to old policy like PPO and the second term enforces the new policy to stay close to the combination of past policies called virtual policy. The virtual policy is built by a convex combination of past policies in the memory through attention weights. The motivation to use attention is to ensure the virtual policy has good performance on current data. To evaluate their method, this paper uses 3 control tasks (i.e. Pendulum, LunarLander and Bipedal- Walker), MiniGrid library for sparse reward environments evaluation, 6 Mujoco tasks, and 6 Atari games.