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: In this work, the authors propose the use a trust-region method in the vein of Proximal Policy Optimization (PPO). In contrast to previous work, the author propose to use two trust regions instead of a single one. Instead of constraining the policy to stay near to the previous policy, the author's proposed algorithm includes an extra virtual policy constructed from a memory buffer of past policies.