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: The paper proposes decoupling the behavior and proximal policies used in policy optimization algorithms such as PPO and PPG. Typically, the behavior policy itself is used as the proximal policy, i.e. the policy to which updates need to be close to (maintain trust region). In such a case, using stale or old data can lead to bad performance, leading to on-policy methods using data from the only most recent sampling iteration. The paper then talks about the batch invariance property, which refers to the performance remaining as is when changing the batch size. This is typically achieved by adjusting the optimization parameters such as learing rate (ex. if batch size is doubled, learning rate should be doubled as well). Finally, the authors perform experiments to test how the decoupling affects the usage of stale data, how batch invariance can be achieved with the decoupled objective, and how much better does the decoupling perform as compared to when the same policy is used as the behvior and proximal policies.