Batch size-invariance for policy optimization

1 INTRODUCTION . Policy gradient-based methods for reinforcement learning have enjoyed great success in recent years . The stability and reliability of these methods is typically improved by controlling the size of policy updates , using either a “ trust region ” ( TRPO ) or a surrogate objective ( PPO ) ( Schulman et al. , 2015 ; 2017 ) . The usual justification for this is that we can not trust updates that take us too far from the policy used to collect experience , called the behavior policy . In this work we identify a subtle flaw with this : the behavior policy is irrelevant to the justification . Instead , what matters is that we control how fast the policy is updated , or put another way , that we approximate the natural policy gradient ( Kakade , 2001 ) . Our key insight is that the “ old ” policy in these methods serves two independent purposes . The first purpose is for off-policy corrections , via importance sampling , for which the old policy must be the behavior policy . The second purpose is to control the size of policy updates , for which the old policy can be any recent policy , which we call the proximal policy . It does not matter whether the proximal policy is also the behavior policy ; it only matters how old the proximal policy is . We demonstrate this by running PPO with stale data collected using a policy from multiple iterations ago , which causes performance to quickly degrade unless the proximal policy is decoupled from the behavior policy . Our insight allows us to make PPO batch size-invariant , meaning that when the batch size is changed , we can preserve behavior , as a function of the number of examples processed , by changing other hyperparameters ( as long as the batch size is sufficiently small ) . We achieve this by using an exponentially-weighted moving average ( EWMA ) of the policy network ’ s weights as the network for the proximal policy . Batch size-invariance has been studied many times before ( see Section 3.1 ) , sometimes under the name “ perfect scaling ” . It is of practical benefit when we wish to increase the batch size to reduce gradient variance , but computational resources such as GPU memory do not allow this . In such a situation , we can instead adjust other hyperparameters formulaically , thereby spreading out the increased computational load over time . The remainder of the paper is structured as follows . • In Section 2 , we explain the difference between the proximal and behavior policies , and show how to decouple them in PPO ’ s objectives . • In Section 3 , we explain the concept of batch size-invariance , and how it applies to SGD and Adam ( Kingma & Ba , 2014 ) . • In Section 4 , we introduce PPO-EWMA and PPG-EWMA , variants of PPO and PPG ( Cobbe et al. , 2020 ) that make use of our decoupled objectives , and show how to make them batch size-invariant at small batch sizes . • In Section 5 , we provide experimental evidence for our central claims : that decoupling the proximal policy from the behavior policy can be beneficial , and that it allows us to achieve batch size-invariant policy optimization . • Finally , in Section 6 , we discuss the theoretical and practical implications of our results . 2 DECOUPLED POLICY OBJECTIVES . In this section we explain the difference between the proximal and behavior policies , and introduce new versions of PPO ’ s objectives in which they have been decoupled . PPO alternates between sampling data through interaction with the environment , and optimizing a surrogate objective . The policy used for sampling is denoted πθold , and is used by the objective in two different ways . This is easiest to see with the KL penalized objective ( Schulman et al. , 2017 , equation ( 8 ) ) : LKLPEN ( θ ) : = Êt [ πθ ( at | st ) πθold ( at | st ) Ât − βKL [ πθold ( · | st ) , πθ ( · | st ) ] ] , where Ât is an estimator of the advantage at timestep t , and Êt [ . . . ] indicates the empirical average over a finite batch of timesteps t. The first use of πθold in this expression is as part of an importance sampling ratio . In order for the policy gradient estimate to be unbiased , this policy needs to be the one that was used for sampling , so we call this the behavior policy πθbehav . The second use of πθold is as a recent target to pull the current policy towards , so we call this the proximal policy πθprox . Our key insight is that the proximal policy need not equal the behavior policy . As we will show experimentally , it matters how old the proximal policy is , but it does not matter whether or not the proximal policy was used for sampling . We therefore define the decoupled KL penalized objective LKLPENdecoupled ( θ ) : = Êt [ πθ ( at | st ) πθbehav ( at | st ) Ât − βKL [ πθprox ( · | st ) , πθ ( · | st ) ] ] , where πθbehav is the policy used for sampling , and πθprox is a recent policy yet to be specified . It is less obvious how to decouple the clipped PPO objective , because πθold only appears once in that expression ( Schulman et al. , 2017 , equation ( 7 ) ) : LCLIP ( θ ) : = Êt [ min ( rt ( θ ) Ât , clip ( rt ( θ ) , 1− , 1 + ) Ât ) ] , where rt ( θ ) : = πθ ( at|st ) πθold ( at|st ) . However , we can rewrite this objective as LCLIP ( θ ) = Êt [ 1 πθold min ( πθÂt , clip ( πθ , ( 1− ) πθold , ( 1 + ) πθold ) Ât ) ] ( omitting the policy arguments ( at | st ) for brevity ) . Now the first use of πθold is as part of an importance sampling ratio , for which we must use the behavior policy , and the second and third uses are in applying the implicit KL penalty , for which we can use the proximal policy . We therefore define the decoupled clipped objective LCLIPdecoupled ( θ ) : = Êt [ πθprox ( at | st ) πθbehav ( at | st ) min ( rt ( θ ) Ât , clip ( rt ( θ ) , 1− , 1 + ) Ât ) ] , where rt ( θ ) : = πθ ( at | st ) πθprox ( at | st ) . As a sanity check , note that if we set the KL penalty coefficient β = 0 or the clipping parameter =∞ , then the dependence on the proximal policy disappears , and we recover the vanilla ( importancesampled ) policy gradient objective ( Schulman et al. , 2017 , equation ( 6 ) ) . A similar decoupled policy objective is also propsed in Mirror Descent Policy Optimization ( Tomar et al. , 2020 , Section 5.1 ) . 3 BATCH SIZE-INVARIANCE . We say an algorithm is batch size-invariant to mean that when the batch size is changed , the original behavior can be approximately recovered by adjusting other hyperparameters to compensate . Here we consider behavior as a function of the total number of examples processed , so another way to put this is that doubling the batch size halves the number of steps needed . Shallue et al . ( 2018 ) and Zhang et al . ( 2019 ) refer to this as “ perfect scaling ” . We treat batch size-invariance as a descriptive property that can hold to some degree , rather than as a binary property . In practice , the original behavior can never be recovered perfectly , and the extent to which it can be recovered depends on both how much and the direction in which the batch size is changed . 3.1 BATCH SIZE-INVARIANCE FOR STOCHASTIC GRADIENT DESCENT . Stochastic gradient descent ( SGD ) is batch size-invariant , up until the batch size approaches some critical batch size . This is the batch size at which the gradient has a signal-to-noise ratio of around 1 . At smaller batch sizes than this , changes to the batch size can be compensated for by a directly proportional adjustment to the learning rate . This core observation has been made many times before ( Mandt et al. , 2017 ; Goyal et al. , 2017 ; Smith et al. , 2017 ; Hardin , 2017 ; Ma et al. , 2018 ; Shallue et al. , 2018 ; McCandlish et al. , 2018 ) . A discussion of this and other previous work can be found in Appendix C. Sketch explanation . For the benefit of the reader ’ s intuition , we sketch the explanation for SGD ’ s batch size-invariance . For a much more thorough explanation , we refer the reader to Mandt et al . ( 2017 ) . Consider running SGD on a loss function L ( θ ; x ) of a parameter vector θ and a data point x . Two steps with batch size n and learning rate α corresponds to the update rule θt+2 = θt − α n ∑ x∈Bt ∇θL ( θt ; x ) − α n ∑ x∈Bt+1 ∇θL ( θt+1 ; x ) , where Bt and Bt+1 are the next two batches of size n. On the other hand , a single step with batch size 2n and learning rate 2α corresponds to the update rule θt+2 = θt − 2α 2n ∑ x∈Bt∪Bt+1 ∇θL ( θt ; x ) . These update rules are very similar , the only difference being whether the gradient for Bt+1 is evaluated at θt or θt+1 . If the batch size is small compared to the critical batch size , then the difference between θt and θt+1 is mostly noise , and moreover this noise is small compared to the total noise accumulated by θt over previous updates . Hence the two update rules behave very similarly . A good mental model of SGD in this small-batch regime is of the parameter vector making small , mostly random steps around the loss landscape . Over many steps , the noise is canceled out and the parameter vector gradually moves in the direction of steepest descent . But a single additional step makes almost no difference to gradient evaluations . In more formal terms , SGD is numerically integrating a stochastic differential equation ( SDE ) . Changing the learning rate in proportion the batch size leaves the SDE unchanged , and only affects the step size of the numerical integration . Once the step size is small enough ( the condition that gives rise to the critical batch size ) , the discretization error is dominated by the noise , and so the step size stops mattering . 3.2 BATCH SIZE-INVARIANCE FOR ADAM . Adam ( Kingma & Ba , 2014 ) is a popular variant of SGD , and is also batch size-invariant until the batch size approaches a critical batch size ( which may be different to the critical batch size for SGD ) ( Zhang et al. , 2019 ) . To compensate for the batch size being divided by some constant c , one must make the following adjustments ( Hardin , 2017 ) : • Divide the step size α by √ c. • ( Raise the exponential decay rates β1 and β2 to the power of 1/c . ) The first adjustment should be contrasted with the linear learning rate adjustment for vanilla SGD . We discuss the reason for this difference and provide empirical support for the square root rule in Appendix D. The second adjustment is much less important in practice , since Adam is fairly robust to the β1 and β2 hyperparameters ( hence it has been parenthesized ) . Note also that β1 also affects the relationship between the current policy and the proximal policy in policy optimization . For simplicity , we omitted this adjustment in most of our experiments , but included it in some additional experiments that are detailed in Appendix E .

Summary: This paper studies the method to achieve the batch size-invariant for policy gradient algorithms (PPO, PPG). The paper achieves this by decoupling the proximal policy from the behavior policy. The experiments demonstrate the effectiveness of the method.