Batch size-invariance for policy optimization
Abstract
We say an algorithm is batch size-invariant if changes to the batch size can largely be compensated for by changes to other hyperparameters.
Stochastic gradient descent is well-known to have this property at small batch sizes, via the learning rate.
However, some policy optimization algorithms (such as PPO) do not have this property, because of how they control the size of policy updates.
In this work we show how to make these algorithms batch size-invariant.
Our key insight is to decouple the proximal policy (used for controlling policy updates) from the behavior policy (used for off-policy corrections).
Our experiments help explain why these algorithms work, and additionally show how they can make more efﬁcient use of stale data
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., 2015a, 2017].
The usual justiﬁcation for this is that we cannot 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 ﬂaw with this: the behavior policy is irrelevant to the justiﬁcation.
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 ﬁrst 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 sufﬁciently 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 beneﬁt because the batch size has a big inﬂuence on training, but is often constrained by computational resources such as GPU memory.
If an algorithm is batch size-invariant, then the batch size may be freely adjusted according to computational constraints, while other hyperparameters are adjusted formulaically to compensate.
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 and 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 beneﬁcial, 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.
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)]:
[Equation]
where ˆ A t is an estimator of the advantage at timestep t , and ˆ E t [ . .
. ] indicates the empirical average over a ﬁnite batch of timesteps t .
The ﬁrst 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 deﬁne the decoupled KL penalized objective
[Equation]
where π θ behav is the policy used for sampling, and π θ prox is a recent policy yet to be speciﬁed.
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)]:
[Equation]
where r t ( θ ) :
= π θ ( a t | s t )
π θ old ( a t | s t ) .
However, we can rewrite this objective as
[Equation]
(omitting the policy arguments ( a t | s t ) for brevity).
Now the ﬁrst 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 deﬁne the decoupled clipped objective
L CLIPdecoupled ( θ ) :=
[Equation] E [Equation]
(cid:20) [Equation] prox [Equation]
π [Equation] behav ( a [Equation]
s [Equation] ) min (cid:16) r t ( θ ) ˆ A t , clip ( r t ( θ ) , 1 − (cid:15), 1 + (cid:15) ) ˆ A t (cid:17)(cid:21) , where r t ( θ ) :
= π θ ( a t | s t ) π θ prox ( a t
| s t ) .
As a sanity check, note that if we set the KL penalty coefﬁcient β = 0 or the clipping parameter (cid:15)
= ∞ , then the dependence on the proximal policy disappears, and we recover the vanilla (importancesampled) policy gradient objective [Schulman et al., 2017, equation (6)].
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
Section 7.2.
Sketch explanation.
For the beneﬁt 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.
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
[Equation]
where B t and B t +1 are the next two batches of size n .
On the other hand, a single step with batch size 2 n and learning rate 2 α corresponds to the update rule
[Equation]
These update rules are very similar, the only difference being whether the gradient for B t +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 and 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 β
and β 2 to the power of 1 / c .)
The ﬁrst 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 C.
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 D.
3.3 Batch size-invariance for policy optimization
In policy optimization algorithms like PPO, there are two different batch sizes: the number of environment steps in each gradient step, which we call the optimization batch size , and the number of environment steps in each alternation between sampling and optimization, which we call the iteration batch size .
When we say that such an algorithm is batch size-invariant, we mean that changes to both batch sizes by the same factor simultaneously can be compensated for.
The motivation for this deﬁnition is that this is the effect of changing the degree of data-parallelism.
If the optimization algorithm (such as SGD or Adam) used by PPO is batch size-invariant, then by deﬁnition this makes PPO optimization batch size-invariant.
In the next section, we show how to make PPO iteration batch size-invariant, and therefore batch size-invariant outright.
To measure batch size-invariance for policy optimization, there are many features of the algorithm’s behavior we could look at.
As a simple metric, we use the ﬁnal performance of the algorithm, since this is the primary quantity of interest to most practitioners.
If an algorithm has a high degree of batch size-invariance, then the difference in ﬁnal performance at different batch sizes should be small.
PPO-EWMA and PPG-EWMA
We now introduce a simple modiﬁcation that can be made to any PPO-based algorithm:
• Maintain an exponentially-weighted moving average (EWMA) of the policy network, updating it after every policy gradient step using some decay rate β prox .
• Use this as the network for the proximal policy in one of the decoupled policy objectives.
The motivation for using an EWMA is as follows.
We would like to be able to use a policy from some ﬁxed number of steps ago as the proximal policy, but this requires storing a copy of the network from every intermediate step, which is prohibitive if this number of steps is large.
Using an EWMA allows us to approximate this policy using more reasonable memory requirements.
Although this approximation is not exact, averaging in parameter space may actually improve the proximal policy [Izmailov et al., 2018], and the age of the proximal policy can still be controlled by adjusting β prox .
We refer to this modiﬁcation using the -EWMA sufﬁx.
Thus from PPO we obtain PPO-EWMA, and from Phasic Policy Gradient (PPG)
[Cobbe et al., 2020] we obtain PPG-EWMA.
Pseudocode for PPO-EWMA may be found in Appendix A, and code may be found at https://github.com/ openai/ppo-ewma .
To see how this modiﬁcation helps us to achieve batch size-invariance, note that the main effect of changing the iteration batch size in PPO is to change the age of the behavior and proximal policies (which are coupled).
The age of the behavior policy affects how on-policy the data is, but this does not matter much, as long as it is not too large.
However, the age of the proximal policy affects the strength of the KL penalty (or the implicit KL penalty in the case of the clipped objective), which inﬂuences how fast the policy can change.
We therefore need to maintain the age of the proximal policy as the iteration batch size is changed, which is what our modiﬁcation enables.
More speciﬁcally, to achieve batch size-invariance for PPO- and PPG-EWMA, we make the following adjustments to compensate for the optimization and iteration batch sizes being divided by some constant c :
• Adjust the optimization hyperparameters as described in the previous section, i.e., divide the vanilla SGD learning rate by c or the Adam step size by √ c .
(We use Adam.)
• Modify β prox such that 1 1 − β prox − 1 is multiplied by c .
(This expression is the center of mass of the proximal policy EWMA, measured in gradient steps.)
This adjustment is what keeps the age of the proximal policy constant, measured in environment steps.
• If using advantage normalization, multiply the number of iterations used to estimate the advantage mean variance by c .
(In practice, we use EWMAs to estimate the mean and variance, and multiply their effective sample sizes, measured in iterations, by c . 1 )
This keeps the overall sample sizes of these estimates constant, preventing their standard errors becoming too large.
• For PPG, multiply the number of policy iterations per phase N π by c .
(We use PPG.)
For these adjustments to work, we require that the optimization batch size is sufﬁciently small.
We also require that the number of policy epochs (denoted E in PPO or E π in PPG) is 1.
This is because when the iteration batch size is very small, using multiple policy epochs essentially amounts to training on the same data multiple times in a row, which is redundant (modulo changing the learning rate).
Our batch size-invariance experiments therefore use PPG-EWMA, where E π = 1 is the default.
Note that PPG has a third batch size: the number of environment steps in each alternation between phases, which we call the phase batch size .
The effect of our adjustment to N π is to simply hold the phase batch size constant, thereby preserving the dynamics of the policy and auxiliary phases.
Experiments
To validate our analysis, we ran several experiments on Procgen Benchmark [Cobbe et al., 2019], which we found to serve as a useful testbed due to the difﬁculty and diversity of the environments.
Hyperparameters for all of our experiments can be found in Appendix B, and full results on each of the individual environments can be found in Appendix F.
5.1 Artiﬁcial staleness
To investigate our decoupled policy objectives, we introduced artiﬁcial staleness.
By this we mean that once data has been sampled through interacting with the environment, it is not immediately used for optimization, but is instead placed in a buffer to be used a ﬁxed number of steps later.
Despite being artiﬁcial, similar staleness is often encountered in asynchronous training setups, where it is known to cause problems for on-policy algorithms like PPO [OpenAI et al., 2019].
We measure staleness in iterations, with one iteration being a single alternation between sampling and optimization.
With artiﬁcial staleness, the original PPO objectives are underspeciﬁed, since there are two natural choices for π θ old : the policy immediately preceding the current iteration, denoted π θ recent , and the behavior policy π θ behav .
However, the decoupled objectives allow us to take the proximal policy π θ prox to be the recent policy, while continuing to use the behavior policy for importance sampling. to
Environment steps × 10− 0 .0 .0 .0 .0 .
M ea n no r m a li
ze
d r e t u r n
PPO with π θ old =
π θ recent 0 .0 .0 .
Environment steps ×
PPO with decoupled objective 0 .0 .0 .0 .0 .1 .
Environment steps ×
PPO with π θ old =
π θ behav
Staleness
None1iteration2iterations4iterations8iterations16iterations32iterations
(a) Using the recent policy for importance sampling introduces bias that makes training highly unstable for even small amounts of staleness.
(b) The decoupled objective allows the correct importance sampling ratio to be used while maintaining the age of the proximal policy, preventing performance from degrading much until the data is very stale.
(c) Using the behavior policy to control the size of policy updates holds back learning unnecessarily for small amounts of staleness, but the additional stability is helpful for very stale data
This allows the KL penalty (or clipping) to have a consistent effect in terms of controlling how fast the policy changes, while avoiding harmful bias from incorrect importance sampling.
In our experiments, we compare the decoupled objective to both choices for the original objective.
Our results are shown in Figure 1.
With both choices for the original objective, even a small amount of staleness hurts performance.
However, with the decoupled objective, performance is robust to a surprising amount of staleness, with minimal degradation until a staleness of around 8 iterations (over 500,000 environment steps).
This demonstrates that the decoupling the proximal policy from the behavior policy can be beneﬁcial.
5.2 Batch size-invariance
We tested our method of achieving batch size-invariance for PPG-EWMA described in Section 4.
Since the optimization batch size is required to be sufﬁciently small, we started from our default batch size, which uses 256 parallel copies of the environment, and reduced it by factors of 4 until we were running just a single parallel copy of the environment.
Our results are shown in Figures 2 and 3.
We were able to achieve a high degree of batch sizeinvariance, with a difference in ﬁnal mean normalized return between the largest and smallest batch sizes of 0.052.
Moreover, there was a single outlier environment, Heist, without which this difference is reduced to 0.019.
We conducted further experiments to try to explain this outlier, which we discuss in Appendix D, but we were not successful.
We conducted ablations in which all but one of the adjustments was removed, the results of which are also shown in Figures 2 and Figure 3.
The Adam step size adjustment is the most important at every batch size, and training becomes highly unstable at the smallest batch sizes without this.
The advantage normalization adjustment does not matter at the largest batch sizes, but matters a lot at the smallest batch sizes in some environments, which are the ones with particularly noisy advantage standard deviation estimates (see Figure 13 in Appendix F.2).
The adjustment to the EWMA matters a little at every batch size, which reﬂects the fact that PPG is relatively robust to changes to the KL penalty.
We did not run an ablation for the adjustment to the PPG hyperparameter N π , but can infer from Cobbe et al.
[2020, Figure 5] that it comes immediately after the Adam step size adjustment in importance.
Environment steps × 100 .0 .0 .0 .0 .
M ea n no r m a li
ze
d r e t u r n
PPG-EWMA with size invariance adjustments
Batch size
Default
Default / Default /
Default / Default /
PPG with KL penalty coeff β
= 1 /0 .
8 1 .× 100 .0 .0 .0 .0 .(a)
No Adam step size adjustment 0 . 0 0 .
8 1 .× 100 .0 .0 .0 .
.(b)
No adv. norm.
adjustment 0 .
Environment steps × 100 .0 .0 .0 .0 .(c)
No EWMA adjustment 0 . 0 0 .
Environment steps × 100 .0 .0 .0 .0 .(d)
No EWMA at all, just PPG
[Figure]
adjustments
[Figure]
D i ff e r e n ce
i n no r m a li
ze d r e t u r n ( s qu a r e r oo t s ca l e )
Environment steps × 100 .0 .0 .0 .0 .
M ea n no r m a li
ze
d r e t u r n
PPO
PPO-EWMA
PPG
PPG-EWMA
PPO-EWMA vs PPO
PPG-EWMA vs PPG
Algorithm and environment − 0 .0 .0 .0 .0 .0 .
D i ff e r e n ce
i n no r m a li
ze d r e t u r n ( s qu a r e r oo t s ca l e )
To check the statistical signiﬁcance of the effects produced by our ablations, we conducted hypothesis tests, which we describe in Appendix
E. Our null hypothesis was that the ablation had no effect on the difference in ﬁnal normalized return at different batch sizes in any of the environments.
For the comparison between the largest and smallest batch sizes, we rejected the null hypothesis for all of the ablations at the 0.1% level.
5.3 EWMA comparison
Finally, we tested the outright beneﬁt of the EWMA modiﬁcation by doing a head-to-head comparison of PPO against PPO-EWMA and of PPG against PPG-EWMA.
It is important to note that the EWMA introduces an additional hyperparameter β prox , but that this was tuned only on the ﬁrst 8 of the 16 Procgen environments (and only on PPG), and so the algorithms are “complete” on the lastenvironments in the sense of Jordan et al.
Our results are shown in Figure 4.
We found the beneﬁt of the EWMA to be small but remarkably consistent across environments and algorithms, outperforming the baseline on all of the last 8 environments for both PPO and PPG.
We believe that this is the result of the EWMA reducing the variance of the proximal policy.
Further evidence that the variance of the proximal policy matters is discussed in Appendix G.
Note that this beneﬁt comes at the cost of additional memory to store the weights of the EWMA network, and an additional forward pass of the EWMA network for each policy gradient step.
With our hyperparameters, this increases the computational cost of PPO by 30% and of PPG by 2.3%, not including the cost of stepping the environment.
Discussion 6.1 PPO as a natural policy gradient method
Our experiments provide strong empirical support that decoupling the proximal policy from the behavior policy can be beneﬁcial: it can be used to make more efﬁcient use of stale data, to achieve batch size-invariance, and to slightly improve sample efﬁciency outright.
This implies that the usual justiﬁcation for PPO’s surrogate objectives, that they approximate trust region methods, is subtly ﬂawed.
Trust region methods keep the policy close to the behavior policy, but it does not matter how far from the behavior policy we move speciﬁcally, only that we stay close to some recent policy, or in other words, that we do not move too fast.
Instead, we speculate that PPO is better viewed as a natural policy gradient method [Kakade, 2001].
These methods select updates that efﬁciently improve performance relative to how much the policy is changed.
This conﬂicts with the results of Schulman et al.
[2015a], which found constraining updates relative to the behavior policy to be beneﬁcial.
With the beneﬁt of hindsight, we believe that at that time, constraint methods had hyperparameters that were easier to tune, but that with the advent of PPO’s clipped objective and various normalization schemes, this tends to no longer be the case.
6.2 Practical advice for policy optimization at small batch sizes
Batch size is an important hyperparameter in reinforcement learning, since it controls gradient variance, which can be high in challenging environments.
But it is often constrained by computational resources such as GPU memory.
The beneﬁt of batch size-invariance is that it allows hyperparameters to be tuned at one batch size, and then adjusted to work at a different batch size, which can be selected based on computational constraints.
However, when working in a new domain, we may be constrained to use a small batch size, without having ﬁrst been able to tune hyperparameters at a larger batch size.
We therefore attempt to distill our ﬁndings into practical advice for getting policy optimization to work well in a new domain at small batch sizes.
Our advice, much of which is already folklore, is as follows:
• By far the most important hyperparameter to tune is the learning rate (or Adam step size).
Once it has been tuned for a certain batch size, it can be adjusted formulaically for use at other batch sizes using the rules given in Section 3, as long as the batch size remains small. • Consider setting the number of policy epochs ( E in PPO or E π in PPG) to 1, at least initially.
This is the easiest way to maintain stability even if not enough ratios are being clipped.
Furthermore, multiple policy epochs are less likely to be beneﬁcial when the iteration batch size is small.
• If using clipping, monitor the fraction of ratios clipped.
If it is much less than 1%, then it is probably beneﬁcial to increase the iteration batch size 4 , or to use PPO-EWMA with a high β prox .
If it is much more than 10% with 1 policy epoch or 20% with multiple policy epochs, then this is often a sign that the learning rate is too high.
• If using advantage normalization, monitor the advantage standard deviation estimates.
If estimates oscillate by a factor of 10 or more, then it is probably beneﬁcial to perform normalization using data from more iterations.
Related work 7.1 Policy optimization
Policy gradient methods have a long history, going at least as far back as Williams [1992], and were introduced in their modern formulation by Sutton et al.
Trust regions were introduced into policy optimization by Schulman et al.
[2015a], who already noted the similarity with natural policy gradient methods [Kakade, 2001].
The trust region approach was simpliﬁed by the use of surrogate objectives in PPO, which was introduced by Schulman et al.
Other key ingredients of PPO include an actor-critic setup similar to that of Mnih et al.
[2016], and generalized advantage estimation, which was introduced by Schulman et al.
[2015b].
PPO’s surrogate objectives can also be viewed through the lens of mirror descent, as explained by
Neu et al.
This perspective also makes it clear that decoupling the proximal policy from the behavior policy is a possibility.
Indeed, a similar decoupled objective is proposed in concurrent work of Tomar et al.
[2020, Section 5.1] that is motivated by this perspective.
7.2 Batch size-invariance
There has been much previous work on batch size-invariance.
The underlying idea of modeling
SGD as numerically integrating a stochastic differential equation (SDE) is long-established, going at least as far back as Kushner and Yin [2003].
More recently, Mandt et al.
[2017] and Hardin [2017] observed that changing the learning rate in proportion the batch size leaves the SDE unchanged, and therefore that SGD is batch size-invariant at small batch sizes.
Meanwhile, Goyal et al.
[2017] empirically validated this invariance on ImageNet.
Smith and Le [2017] and Smith et al.
[2017] provided further empirical validation for this, as well as for rules describing how the optimal learning rate changes with momentum and training set size.
The term critical batch size for the batch size beyond which SGD is no longer batch size-invariant was introduced by Ma et al.
Since then a number of works have studied how the critical batch size varies between problems, mostly with the motivation of improving training efﬁciency at large batch sizes (in contrast to our work, which focuses on maintaining performance at small batch sizes).
Shallue et al.
[2018] studied the effect of architectures, datasets and different forms of momentum on the critical batch size, and introduced the term perfect scaling for the regime of batch size-invariance.
Golmant et al.
[2018] also studied the effect of dataset complexity and size on the critical batch size.
McCandlish et al.
[2018] measured the critical batch size in a range of domains, and showed that it can be predicted using a measure of the noise-to-signal ratio of the gradient known as the gradient noise scale .
Finally, Zhang et al.
[2019] studied the effect of curvature on the critical batch size using a noisy quadratic model, and showed that preconditioning can be used to increase the critical batch size.
Conclusion
Policy optimization algorithms such as PPO typically control the size of policy updates using a recent policy we call the proximal policy.
We have shown that this policy can be safely decoupled from the behavior policy, which is used to collect experience.
We introduced PPO-EWMA and PPG-EWMA, variants of PPO and PPG in which the proximal policy is an exponentially-weighted moving average of the current policy.
These variants allow stale data to be used more efﬁciently, and are slightly more sample efﬁcient outright.
Finally, we showed how to make these algorithms 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 not too large).
We discussed our ﬁndings, which have both theoretical and practical implications for policy optimization.
Acknowledgments
We thank David Farhi, Chris Hardin and Holly Mandel for helpful discussions and comments and anonymous reviewers for helpful and detailed feedback.
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes] There are careful ablations and a discussion of the additional computational cost in Section 5.
(c) Did you discuss any potential negative societal impacts of your work? [No] The main impact of the work is via an improved understanding of policy optimization in general, and the experiments are all on toy environments. (d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[N/A] (b) Did you include complete proofs of all theoretical results?
[N/A] 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[No] The additional computational cost compared to existing methods is clearly described in Section 5, from which the total computation cost of the experiments can be calculated.
The type of resources used is proprietary information.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] (b) Did you mention the license of the assets?
[Yes] (c) Did you include any new assets either in the supplemental material or as a URL?
[Yes] (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?
[N/A] (e) Did you discuss whether the data you are using/curating contains personally identiﬁable information or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
A Pseudocode for PPO-EWMA
We provide pseudocode PPO as well as PPO-EWMA to make it clear what changes need to be made:
Algorithm 1 PPO for iteration = 1 , 2 , . . .
do for actor = 1 , 2 , . .
. , N do
Run policy π θ old in environment for T timesteps
Compute advantage estimates ˆ A 1 , ˆ A 2 , . . .
, ˆ A T end for for epoch
= 1 , 2 , . .
, E do for each minibatch do
Optimize objective L with respect to θ on minibatch end for end for θ old ← θ end for
Algorithm 2 PPO-EWMA for iteration = 1 , 2 , . . .
do for actor = 1 , 2 , . .
. , N do
Run policy π θ behav in environment for T timesteps
Compute advantage estimates ˆ A 1 , ˆ A 2 , . . .
, ˆ A T end for for epoch
= 1 , 2 , . .
, E do for each minibatch do
Compute π θ prox for minibatch Optimize objective L decoupled with respect to θ on minibatch θ prox ← EWMA β prox ( θ ) end for end for θ behav ← θ end for
The expression EWMA β prox ( θ ) is shorthand for θ t + β prox θ t − 1 + β 2prox θ t − 2 + · · · + β t prox θ1 + β prox + β 2prox + · · · + β t prox , where θ 0 , θ 1 , . . .
θ t are the values of θ after each gradient step.
In practice, we compute this incrementally by initializing θ prox ← θ and w ← 1 , and treating the update θ prox ← EWMA β prox ( θ ) as shorthand for w new ← 1 + β prox w θ prox ← w new θ + β prox w w new θ prox w ← w new .
For PPG-EWMA, we make the same changes to the policy phase, while leaving the auxiliary phase unchanged.
However, the EWMA should be reinitialized at the start of each policy phase , since θ changes a lot during the auxiliary phase.
Code for both PPO-EWMA and PPG-EWMA may be found at https://github.com/openai/ ppo-ewma .
B Hyperparameters
All experiments were on Procgen’s hard difﬁculty, without frame stack, using the convolutional neural network from IMPALA [Espeholt et al., 2018].
Unless stated otherwise, experiments lasted formillion environment steps.
[Table] [Table]
[Table]
For the purpose of the artiﬁcial staleness experiments
[Table]
we clipped π θ [Table] to keep the ratio
[Table]
below 100, for numerical stability.
For PPG-EWMA [Table] we chose the default EWMA decay rate β prox such that the center of mass of the
EWMA, 1 1 − β prox − 1 , equaled the number of minibatches per policy phase iteration (8), so that the maximum age of the proximal policy is the same in PPG and PPG-EWMA.
We tuned this on the ﬁrst 8 of the 16 Procgen environments by also trying
[Equation]
= [Equation] and [Equation]
, but did not ﬁnd these to perform better.
We did not re-tune β prox on the last 8 Procgen environments or on PPO-EWMA
For the batch size-invariance experiments, we made the following changes to the above defaults:
• We reduced the number of parallel environments, ﬁrst by reducing the number of workers from 4 to 1, and then by reducing the number of parallel environments per worker fromto 16 to 4 to 1.
[Table]
• For the policy phase, we adjusted the Adam step size ( α ), the proximal policy EWMA decay rate ( β prox ), advantage normalization, and the number of policy iterations per phase ( N π ) in the way described in Section 4.
• For the auxiliary phase, we initially tried adjusting the Adam step size in the same way as for the policy phase.
This worked well in terms of batch size-invariance, but resulted in prohibitively large wall-clock times at small batch sizes, due to the large number of auxiliary epochs.
We therefore simply kept the auxiliary phase minibatch size per worker constant, and only adjusted the Adam step size when reducing the number of workers, not when reducing the number of parallel environments per worker.
C Adam square root step size adjustment
In Section 3, we stated that SGD and Adam have different learning rate adjustment rules.
To compensate for the batch size being divided by some constant c , one must divide the SGD learning rate by c , but divide the Adam step size by √ c
[Hardin, 2017].
The reason for the difference is that Adam divides the gradient by a running estimate of the root mean square gradient.
If the gradient vector at the current step is g t , then this denominator is approximately
[Equation]
where all operations including the variance operator are applied componentwise, n is the batch size, and B is a componentwise version of the gradient noise scale deﬁned by McCandlish et al.
[2018], a measure of the noise-to-signal ratio of the gradient that approximates the critical batch size.
Hence if the batch size is small compared to the critical batch size, then B (cid:29) n for most components, and so the Adam denominator is approximately proportional to 1 √ n .
It follows that if the batch size is divided by some constant c , then the Adam denominator is multiplied by approximately √ c (providing the batch size is small compared to the critical batch size).
Hence
Adam is effectively dividing the learning rate by √ c automatically, and so the step size α only needs to be adjusted by an additional √ c to effectively divide the learning rate by c overall.
This all ignores Adam’s (cid:15) hyperparameter, which is usually negligible, but is sometimes used to interpolate between Adam and momentum SGD.
To verify the square root rule for Adam, we conducted an ablation of our batch size-invariance experiments, in which we made the exact same adjustments, except that we divided the Adam step size by c instead of by √ c .
Our results are shown in Figure 5.
When compared with Figure 2, this clearly shows that the square root rule is superior in our setting.
Full results on each of the individual environments can be found in Appendix F.2. [Figure] [Figure] [Figure]
× [Figure] [Figure] [Figure]
[Figure]
[Figure]
li ze d r e t u r n
PPG-EWMA with step size adjustment
Batch size
Default
Default / Default /
Default / Default /
Our results are in tension with those of Smith et al.
[2017], who veriﬁed batch size-invariance for
Adam using the linear rather than the square root rule.
However, they achieved a lower degree of batch size-invariance with Adam than with SGD (see Figure 4 in that work), and moreover, our results show that the batch size needs to be reduced signiﬁcantly before the difference between the two rules is noticeable.
We believe that this accounts for their experimental results, and that the square root rule is superior in general (with the exception of when Adam’s (cid:15) hyperparameter is high enough for it to behave like momentum SGD).
D Adam β 1 and β 2 adjustments
As discussed in Section 3, there is an additional adjustment one should make when using Adam, other than to the step size α .
To compensate for the batch size being divided by some constant c , one should also raise the exponential decay rates β 1 and β 2 to the power of 1 / c
[Hardin, 2017].
We omitted this adjustment in most of our experiments, and were still able to achieve a high degree of batch size invariance.
For all except one environment, the difference in normalized return between the largest and smallest batch sizes at the end of training was at most 0.11 (see Figure 3).
For these environments, it would probably have required many additional experiments to detect any further improvement that adjusting β 1 and β 2 might provide.
However, for the Heist environment, this difference was 0.55.
We hypothesized that this might be explained by the fact that we did not adjust β 1 and β 2 .
We therefore conducted a version of our batch size-invariance experiments in which we either adjusted only β 2 using the above rule, or adjusted both β 1 and β 2 .
Our results are shown in Figure 6.
In both cases there was still a large difference in performance at the largest and smallest batch sizes.
Environment steps ×
R e t u r n
No further adjustments
Batch size
Default
Default / Default /
Default /
Default /0 .0 .0 .0 .0 .1 .
Environment steps ×
Adjustment to β 2 only 0 .0 .0 .0 .0 .1 .
Environment steps ×
Adjustments to β 1 and β
E Hypothesis tests for ablations
As discussed in Section 3, we conducted hypothesis tests to check the statistical signiﬁcance of the effects produced by the ablations to our batch size-invariance experiments.
Our primary metric for measuring batch size-invariance was the difference in ﬁnal performance of the algorithm at different batch sizes.
To get a complete picture of how important our ablations were at different batch sizes, we compared our default (largest) batch size with each of the other batch sizes, and tested the hypothesis that the difference was larger for the ablation.
This resulted in 16 hypotheses, corresponding to the 4 ablations and the 4 non-default batch sizes.
To test each hypothesis, we used a van Elteren test [van Elteren, 1960], a stratiﬁed version of the Mann–Whitney U -test, treating the different environments as strata.
This gives a non-parametric Z -test of the null hypothesis that for each of the environments, the probability of the ablation outperforming the original experiment is the same as the probability of the original experiment outperforming the ablation [LaVange and Koch, 2006].
To reduce noise (and thereby increase statistical power), we measured average performance over the last 4 million timesteps (the length of a single PPG phase).
We used a signiﬁcance level of 0.1% and applied a Bonferroni correction to account for multiple comparisons.
Our results are shown in Table 5.
For ablation (b), the difference is only signiﬁcant at the smallest batch size.
For all other ablations, the difference is signiﬁcant at every batch size, execpt for the largest batch size for ablation (d).
above 3.84).
Batch sizes:
Default vs ... size adjustment adjustment adjustment , just PPG
Default /0 .
046 ( 5 . 02 ) − 0 .
014 ( − 1 . 53 ) 0 .
Default /0 .
347 ( 7 . 31 ) − 0 .
033 ( − 1 . 53 ) 0 .
Default /0 .
( − 0 . 11 ) 0 .
Default /0 .
F Results on individual environments F.1 Artiﬁcial staleness
[Figure] [Figure]
R e t u r n
CoinRun
[Figure]
[Figure]
StarPilot
CaveFlyer −
Dodgeball
[Figure]
R e t u r n
[Figure]
[Figure] [Figure] [Figure]
Jumper
[Figure]
e t u r n
[Figure]
R e t u r n
CoinRun
StarPilot
CaveFlyer
Dodgeball
[Figure]
R e t u r n
[Figure]
Chaser
Miner
Jumper
[Figure]
R e t u r n
[Figure]
[Figure] [Figure]
[Figure]
None1iteration2iterations4iterations8iterations16iterations32iterations
R e t u r n
CoinRun
StarPilot
CaveFlyer
Dodgeball
[Figure]
R e t u r n
[Figure]
Chaser
Miner
Jumper
R e t u r n
[Figure]
[Figure] [Figure]
BigFish
Heist
[Figure]
None1iteration2iterations4iterations8iterations16iterations32iterations
F.2 Batch size-invariance
[Figure]
R e t u r n
CoinRun
[Figure]
StarPilot
[Figure]
CaveFlyer
[Figure]
Dodgeball
[Figure]
R e t u r n
[Figure]
R e t u r n
[Figure]
t u
[Figure]
n
[Figure]
R e t u r n
CoinRun
StarPilot
CaveFlyer 2 .5 .
Dodgeball
R e t u r n
FruitBot 0 .2 .
Chaser 2 .5 .
Miner
Jumper
[Figure] [Figure]
[Figure]
R e t u r n
[Figure]
u r n
[Figure]
R e t u r n
CoinRun
StarPilot
CaveFlyer 2 .5 .
Dodgeball
[Figure]
n
[Figure]
Miner
[Figure]
R e t u r n
[Figure]
[Figure]
dv a n t a g e s t d .
d e v .
e s ti m a t
[Figure]
n t a g
[Figure]
s t d .
d
[Figure]
s ti
[Figure]
t
[Figure]
no advantage normalization adjustment only once the estimates oscillate by factor of around 10 or more.
R e t u r n
CoinRun
StarPilot
CaveFlyer 2 .5 .
Dodgeball
[Figure]
R e t u r n
FruitBot
Chaser 2 .5 .7 .
[Figure]
[Figure] [Figure]
Miner
Jumper
[Figure]
R e t u r n
[Figure]
n
[Figure]
R e t u r n
CoinRun
StarPilot
CaveFlyer 2 .5 .
Dodgeball
[Figure]
R e t u r n
FruitBot
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
Miner
Jumper
[Figure]
e t u r n
[Figure]
t
[Figure]
r n
[Figure]
R e t u r n
CoinRun
StarPilot
CaveFlyer 2 .5 .
Dodgeball
[Figure]
R e t u r n
FruitBot
Chaser 2 .5 .
Miner
Jumper
[Figure] [Figure]
R e t u r n
[Figure]
.3 EWMA
R e t u r n
CoinRun
StarPilot
CaveFlyer 2 .5 .
Dodgeball
[Figure]
R e t u r n
FruitBot
Chaser 2 .5 [-.
-].
[Figure]
[Figure]
[Figure]
[Figure]
Miner
[Figure]
R e t u r n
[Figure]
n
[Figure]
G Role of the proximal policy EWMA decay rate
In this section we discuss the role of the hyperparameter β prox in PPO-EWMA and PPG-EWMA in more depth.
Recall that in PPO-EWMA and PPG-EWMA, the proximal policy network parameter vector θ prox is an exponentially-weighted moving average (EWMA) of the policy network parameter vector θ , meaning that
[Equation]
where θ 0 , θ 1 , . . .
θ t are the values of θ after each gradient step.
Rather than working with the decay rate β prox of this EWMA directly, it is conceptually clearer to work with the center of mass of this EWMA,
[Equation]
This is average age of a term in the EWMA in the limit as t → ∞ , and so if θ were to follow a straight line path for example, then θ − θ prox would be approximately proportional to COM prox .
Suppose then that we halve COM prox .
What is the effect of this?
Consider the gradient of the KL divergence from the current policy to the proximal policy, as a function of the proximal policy parameter vector,
[Equation]
Since KL divergence is always greater than or equal to 0 , with equality if and only if the input distributions are equal, Grad-KL ( θ ) = 0 , and hence, to ﬁrst-order, Grad-KL ( θ prox ) is a linear function of θ − θ prox .
Therefore halving COM prox should have a similar effect to halving the KL penalty coefﬁcient β .
In other words, we should be able to compensate for halving COM prox by doubling the KL penalty coefﬁcient.
Intuitively, the KL penalty acts like a rubber band pulling the policy towards the proximal policy.
Halving COM prox is analogous to attaching the rubber band to a point half as far away, while doubling the KL penalty coefﬁcient is analogous to doubling the thickness of the rubber band.
Doing both simultaneously results in the same overall force.
We tested this hypothesis using a hyperparameter grid search over the EWMA center of mass and the
KL penalty coefﬁcient, for PPG-EWMA on StarPilot.
We used our smallest batch size, along with our corresponding batch size-invariance adjustments, to allow the greatest scope for reducing C prox without making the EWMA degenerate into averaging over a single data point.
Our results are shown in Figure 18.
The diagonal banding clearly demonstrates the expected effect.
However, the effect only holds locally: as COM prox is continually halved and the KL penalty coefﬁcient is continually doubled, performance gradually degrades.
We believe that this is because reducing COM prox has a second-order effect, which is to increase the variance of θ − θ prox .
This is both because the EWMA is averaging over a smaller effective sample size, and because θ t − θ t − k has a lower signal-to-noise ratio as k decreases.
Therefore if COM prox has been halved too many times, we should expect to no longer be able to fully compensate for this by continuing to double the KL penalty coefﬁcient.
[Figure]
K L p e n a lt y c o e f ﬁ c i e n t ( β )
Return
[Figure] [Figure]