Batch size-invariance for policy optimization
Authors: Jacob Hilton
Karl Cobbe
John Schulman
Venue: NeurIPS 2022 Accept (2022)
Keywords: reinforcement learning
policy gradient
learning rate

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 efficient use of stale data.

Batch size-invariance for policy optimization
Jacob Hilton OpenAl jhilton@openai.co com
Karl Cobbe OpenAl karl tarl@openai.com
John Schulman OpenAl
joschu@openai.cmm com
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 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. 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 efficient use of stale data.
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 al. 2015a 2017]. The usual justification for this is that cannot trust updates that take us too far from the policy used to collect experience, called the behavior policy. this work we identify a subtle flaw with this: the behavior policy is irrelevant to the justification. Instead, what matters is that 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 preximal policy. It does not matter whether the proximal policy is also the behavior policy; only matters how old the proximal policy is. We demonstrate this by running PPO with stale data collected using 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- rize-ivvrraant, 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-weigee
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 B. sometimes under the name "perfect scaling". It of practical benefit because the batch size has big influence training, but is often constrained by computational resources such GPU memory. an algorithm batch size-invariint, 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.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

In Section
we explain the concept of batch size-invariance, and how applies to SGD
and Adam Kingma and Ba 2014
we introduce PPO-EWMA and PPG-EWMA, variants of PPO and PPG (Cobbe et al. 2020 that make of our decoupled objectives, and show how make them batch
In Section
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 beneficial, and that it allows us achieve batch size-invariant policy optimization.
Finally, in Section
6
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 Taaaa 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)]:
L. KLPEN (0) : E,
To (at 8c)
A. (KL [To (-1 st) ( siil
Toaad (at) 82)
where A, is estimator of the advantage at timestep t, and Eg [...] indicates the empirical average over finite batch of timesteps t. The first use of Taai in this expression 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 T The second use of Te recent target to pull the current policy towards,
call this the proximal policy To
Our key insight that the proximal policy need not equal the behavior policy. 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
Ty (a, 5t)
L
KLPEN decouupeed
(0) = E,
A, - KL [T (t) m
st)
Tuuuaay (az: 88))
where Thpaha is the policy used for sampling, and pprox is recent policy yet be specified.
is less obvious how to decouple the clipped PPO objective, because maaa only appears once that expression [Schulman et al. 2017 equation (7)):
I, CLIP (0) = B min (rx (0) A,, clip (rx 1 oc
where rr (0)
(a(ac lec)
However, can rewrite this objective
Jo (ar (8e)
I, CLIP (0)
min Ao, clip (TO (1 -c) Tosaa (1 6) mo A.)
(omitting the policy arguments (at 8x) for brevity). Now the first use of Taai is as part of an importance sampling ratio, for which must use the behavior policy, and the second and third uses are in applying the implicit KI. penalty, for which we can use the proximal policy.
We therefore define the decoupled clipped objective

---

To proox (at 88)
CLIP decoupped (0) = Eg
min (0) (0) A., clip (0), - ())
"T BBbbbax (at 86)
T (ax $t)
where Tt (0)
=
7TB pross (ac 8p))
As a sanity check, note that if we set the KL penalty coefficient 3 = 0 the clipping parameter e = oc. then the dependence the proximal policy disappears, and recover the vanilla (importance- sampled) policy gradient objective [Schulman al. 2017 equation (6)].
Batch size-invariance
We say algorithm is batch size-invariant to mean that when the batch size is changed, the original behavior can approximately recovered by adjusting other hyperparameters to compensate. Here we consider behavior as 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. (Shalue et al. (2018] and Zhang et al. 2019) refer to this as "perfect scaling".
We treat batch size-invriiance 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 which the batch size is changed.
descriptive property that can hold some degree, rather than as
3.1 Batch size-invariance for stochastic gradient descent
Stochastic gradient descent (SGD) is batch iize-invariant, up until the batch size approaches some critical batch size. This is the batch size at which the gradient has a signal-to-noiss ratio of around
1. At smaller batch sizes than this, changes to the batch size be compensated for by directly
proportional adjustment to the learning rate. This core observation has been made many times before [Mandt et al. 2017 Goyal al. 2017 (Smihh al. 2017 Hardin 2017 Ma al., [018 Shallue
et al. 2018 McCandlish et al. 2018). A discussion of this and other previous work can be found in Section
Sketch explanation. For the benefit of the reader's intuition, we sketch the explanation for SGD's batch size-invariance. For much more thorough explanation, we refer the reader to Mandt et al.
2017
Consider running SGD on a loss function L (0; x) of parameter vector 0 and data point r. Two steps with batch size n and learning rate a corresponds the update rule
a Ou42 = 0.
Σ VoL VoL B
Σ VoL (0++1:)) BEB
YI
where B, and B are the next two batches of size n. On the other hand, single step with batch size 2n and learning rate 2a corresponds to the update rule
O242 = =0, =
2a
Σ
VoL (0;z))
2n
ZEBLUBB+1
These update rules are very similar, the only difference being whether the gradient for B:41 is evaluated at Og or O.+1. If the batch size is small compared to the critical batch size, then the difference between 6g and d.+1 is mostly noise, and moreover this noise small compared to the total noise accumulated by 6 over previous updates. Hence the two update rules behave very similarly.
good mental model of SGD in this small-batch regime is of 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 difference to gradient evaluations.

---

In more formal terms, SGD is numerically integrating 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- iizerinvariant 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 or by Vc.
(Raise the exponential decay rates 31 and Bg 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
The second adjustment is much less important in practice, since Adam is fairly robust to the 81 and 32 hyperparameters (hence it has been parenthesized). Note also that 81 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 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 each alternation between sampling and optimization, which we call the iteration butch 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 definition is that this is the effect of changing the degree of data aata-parallelism.
If the optimization algorithm (such as SGD or Adam) used by PPO is batch size-invariant, then by definition this makes PPO optimization batch size-invariant. In the next section, we show how to make PPO iteration batch size-invare invariant, and therefore batch inzaaiinttinnt outright.
To measure batch size- iitarinvariance for policy optimization, there are many features of the algorithm's behavior could look at. As simple metric, use the final 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 final performance different batch sizes should be small.
4 PPO-EWMA and PPG-EWMA
We now introduce a simple modification that can be made to any PPO-based algorithm:
Maintain ggponentially-weightes moving average (EWMA) of the policy network, updat- ing it after every policy gradient step using some decay rate 3prox.
Use this 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 fixed number of steps ago the proximal policy, but this requires storing copy of the network from every intermediate step, which is prohibitive this number of steps is large. Using an EWMA allows 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. 2 and the age of the proximal policy can still be controlled by adjusting Sprox-
We refer to this modification using the -EWMA suffix. 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 and code may be found at httph:://gtiii com/ openai/ppo- ewma
To see how this modification helps us to achieve batch size-invariance, note that the main effect of changing the iteration batch size in PPO is 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 influences 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 modification enables.
More specifically, to achieve batch size- invaarinee for PPO- and PPG-EWMA, 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 or the Adam step size by vc. (We use Adam.)
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.
Modify Pprox such that T- Peone
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 e 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 sufficiently small. We also require that the number of policy epochs (denoted E in PPO or E, in PPG) is 1. This because when the iteration batch size is very small, using multiple policy epochs essentially amounts to training on the same data multiple times in row, which is redundant (modulo changing the learning rate). Our batch size-invariance experiments therefore use PPG-EWMA, where F.- 1 is the default.
Note that PPG has 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 Procgen Benchmark (Cobbe al. 2019). which we found to serve as useful testbed due to the difficulty 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
5.1 Artificial staleness
To investigate our decoupled policy objectives, we introduced artificial 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 fixed number of steps later. Despite being artificial, similar staleness is often encountered in asynchronous training setups, where it is known to cause problems for on-policy algorithms like PPO [OpenAl et al., 2019]. We measure staleness in iterations, with one iteration being single alternation between sampling and optimization.
With artificial staleness, the original PPO objectives are underspecified, since there are two natural choices for Thws the policy immediately preceding the current iteration, denoted T and the behavior policy Thucher However, the decoupled objectives allow us to take the proximal policy proo to be the recent policy, while continuing to use the behavior policy for importance sampling.
'The effective sample size, sometimes called the span, of an EWMA with decay rate 3 is equal 1-3 L.
5

---

PPO with RA.a Wiccos
PPO with decoupled ohjective
PPO with RaAM Bauhaer
0.6
Mean normalized neturn
0.1
0.2
Stuleness
Nane
iterations
iterations
ierrtions
itenaions
0.2
itenaioos
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.1
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Environment steps
x10*
Environment steps
x10#
Environment steps
10
(a) Using the recent policy for im- portance sampling introduces bias that makes training highly unstable for even small amounts of staleness.
(b) The decoupled objective allows the correct importance sampling ra- tio be used while maintaining the age of the proximal policy, pre- venting performance from degrad- ing 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.
Figure I: PPO with artificial staleness, averaged over all 16 Procgen environments. One iteration corresponds 65, 536 environment steps with our hyperparameters. Mean and standard deviation over 4 seeds shown.
This allows the KL penalty (or clipping) to have 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 shown in Figure With both choices for the original objective, even small amount of staleness hurts performance. However, with the decoupled objective, performance is robust to surprising amount of staleness, with minimal degradation until staleness of around iterations (over 500,000 environment steps). This demonstrates that the decoupling the proximal policy from the behavior policy can be beneficial.
5.2 Batch size-invariance
We tested method of achieving batch size-invariance for PPG-EWMA described in Section 4 Since the optimization batch size is required to be sufficiently small, we started from our default batch size, which uses 256 parallel copies of the environment, and reduced it by factors of until were running just single parallel copy of the environment.
Our results are shown in Figures and We were able to achieve high degree of batch size- invariance, with difference in final 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 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 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 reflects the fact that PPG is relatively robust to changes to the 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.

---

PPG-EWMA with all batch size- invariance adjustments
(a) No Adam step size adjustment
(b) No adv. norm. adjustment
0.8
0.8
0.6
0.4
0.4
0.0
Mean normalzzed
Batch size
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
0.4
Default
x10
X10
Default 4
(c) No EWMA adjustment
(d) No EWMA at all, just PPG
Default
0.8
0.2
Default /
Default / 256
0.6
0.6
PPG with KL penalty coeff
0.0
0.2
0.2
1/256
0.0
0.0
0.2
0.8
0.0
1.0
0.0
0.8
1.0
Environment steps
x10*
Environment steps
x10*
Environment steps
x10*
Figure 2: PPG-EWMA at different batch sizes, with hyperparameters adjusted to achieve batch size-invariance, averaged over all 16 Procgen environments. For reference, we also show PPG (at the default batch size) with the KL penalty coefficient (3 the KKLPNN policy objective) reduced to 1/256, which serves too weak. On the right we show ablations with all but one of the adjustments. Mean and standard deviation over 3 seeds shown.
approximate lower bound on PPG's performance with KL penalty that is
Diflerence in normalized retur
1.25
0.75-
(square rott scale)
0.50
0.25
0.00
PPG-EWMA with all batch size-inve invaiivaee adjustments
(a) No Adam step size adjustment
(b) No adv. adjustment
(c) No EWMA adjustment
(d) No EWMA. just PPG
Experimentul setup and environment
Figure 3: For the results shown in Figure we measure the degree of batch size-invariance for each individual environment by calculating the difference in normalized return between the largest and smallest batch sizes, averaged over the last 4 million timesteps (the length of single PPG phase). use a square root scale to make small differences more visible. Mcan and standard error over seeds shown.
Environments from left right: CoinRun, StarPilot, CaveFlyer, Dodgeball, FruitBot, Chaser, Miner, Jumper, Leaper, Maze, BigFish, Heist, Climber, Plunder, Ninja, BossFight.

---

0.20
0.15
Difference in mozzalized retum
Mean nornalized return
0.6
0.10
(square rook scale)
0.05
0.4
0.00
0.2
PPO
PPO-EWMA
PPG
PPG-EWMA
-0.05
0.0
0.2
0.4
0.6
0.8
1.0
PPO-EWMA vs PPO
PPG-EWMA vs PPG
Environment steps
x10
Algorithm and environment
Figure 4: Performance of all 4 algorithms on Procgen. Left:: learning curves, with the mean and standard deviation over 4 seeds shown. Right: difference in normalized return between the EWMA and the non-EWMA algorithms, split up by environment, and averaged over the last 4 million timesteps, with the mean and standard error over seeds shown.
To check the statistical significance 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 final 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 % level.
5.3 EWMA comparison
Finally, tested the outright benefit of the EWMA modification by doing sead-tt-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 Sproxs but that this was tuned only on the first of the 16 Procgen environments (and only on PPG), and so the algorithms are "complete" on the last 8 environments in the sense of Jordan al. 2020). Our results are shown in Figure Al
We found the benefit of the EWMA to be small but remarkably consistent across environments and algorithms, outperforming the baseline on all of the last 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 discussed in Appendix G
Note that this benefit comes at the cost additional memory to store the weights of the EWMA network, and 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 natural policy gradient method
Our experiments provide strong empirical support that decoupling the proximal policy from the behavior policy can be beneficial: can be used to make more efficient use stale data, to achieve
"These costs are calculated as follows. PPO environment step, to which PPO-EWMA adds 3 forward-only passes. PPG has I forward-only and 7 forward- backward passes of both networks per environment step. to which PPG-EWMA adds I forward-only pass of the policy network. forward-backwant pass has times the cost of forward-only pass.
forward-only and forward-backward passes per
'In practice, including the time taken to step the environment, the EWMA increased wall-clock time by 19% for PPO and by 3% for PPG, but our PPG- PPG-EWMA implementation included an additional unnecessary forward pass of EWMA of the value network.

---

batch size- invariance, and to slightly improve sample efficiency outright. This implies that the usual justification for PPO's surrogate objectives, that they approximate trust region methods, subtly flawed. Trust region methods keep the policy close to the behavior policy, but it does not matter how far from the behavior policy we move specifically, 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 natural policy gradient method [akadee 2001]. These methods select updates that efficiently improve performance relative to how much the policy is changed.
This conflicts with the results of [Schulman et al. [2015]], which found constraining updates relative to the behavior policy to be beneficial. With the benefit 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 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 benefit 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 first been able to tune hyperparameters at a larger batch size. We therefore attempt to distill our findings 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 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 g as long as the batch size remains small.
Consider setting the number of policy epochs (E in PPO or Eg in PPG) to 1, at least initially. This is the easiest way maintain stability even if not enough ratios are being clipped. Furthermore, multiple policy epochs are less likely to be beneficial 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 beneficial to increase the iteration batch size or to use Sprox. If it is much more than 10% with I policy epoch 20% with multiple policy epochs, then this is often a sign that the learning rate is too high.
with a high
If using advantage normalization, monitor the advantage standard deviation estimates. If estimates oscillate by factor of 10 or more, then it is probably beneficial to perform normalization using data from more iterations.
7
Related work
7.1 Policy optimization
Policy gradient methods have a long history, going at least as far back as Williams [192], and were introduced in their modern formulation by /Stttnn et al. [1999]. 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 simplified by the use of surrogate objectives in PPO, which was introduced by (Shhlman et al. [217]. Other key ingredients of PPO include actor-critic setup similar to that of Mnih et al. (016)) and generalized advantage estimation, which was introduced by (Schulman et al. (0015b].
PPO's surrogate objectives can also be viewed through the lens of mirror descent, explained by Neu et al. 2017 This perspective also makes it clear that decoupling the proximal policy from the
*The iteration batch size can be increased without changing the sampling or optimization batch size by simultaneously increasing the number of timesteps per rollout (2') and the number of minibatches per epoch. However, 7 also affects the amount of bootstrapping performed, and so the GAE bootstrapping parameter (A) may also need to be adjusted to compensate.

---

behavior policy is possibility. Indeed, a similar decoupled objective is proposed in concurrent work of Tomar et (2000 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 stochastic differential equation (SDE) is long-established. going at least as far back as Kushner and Yin [200]] 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 batch size-invariant at small batch sizes. Mcanwhile, Goyal et al. [2017
empirically validated this invariance on ImageNet. |Smith and Le [217]] and [Smith 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. 22018. Since then a number of works have studied how the critical batch size varies between problems, mostly with the motivation of improving training efficiency at large batch sizes (in contrast our work, which focuses on maintaining performance 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- invarinnre. Golmant et al. (2018) also studied the effect of dataset complexity and size on the critical batch size. McCandlish al. (2018) measured the critical batch size in 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 gradiens
noise scale. Finally, Zhang al. (019] studied the effect of curvature the critical batch size using a noisy quadratic model, and showed that preconditioning can be used to increase the critical batch
size.
8 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 used to collect experience. We introduced PPO-EWMA and PPG-EWMA, variants of PPO and PPG in which the proximal policy is an exponentialll-weiihhe of the current policy. These variants allow stale data to used more efficiently, and are slightly more sample efficient 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 the batch size is not too large). We discussed our findings, which have both theoretical and practical implications for policy optimization.
moving average
9 Acknowledgments
We thank David Farhi, Chris Hardin and Holly Mandel for helpful discussions and comments, and anonymous reviewers for helpful and detailed feedback.
References
K. Cobbe, C. Hesse, J. Hilton, and J. Schulman. Leveraging procedural generation to benchmark reinforcement learning. arXiv preprint laXiiv 2019.
K. Cobbe, J. Hilton, O. Klimov, and J. Schulman. Phasic policy gradient. arXiv preprint
(rXhv:2009.04416 2020.
L. Espeholt, H. Soyer, R. Munos, K. Simonyan, Mnih, T. Ward, Y. Doron, V. Firoiu, T. Harley, I. Dunning, et al. IMPALA: Scalable distributed deep-RL with importance weighted actor-learner architectures. In International Conference on Machine Learning, pages 1407-1416. PMLR, 2018.
N. Golmant, N. Vemuri, Z. Yao, V. Feinberg, A. Gholami, K. Rothauge, M. W. Mahoney, and Gonzalez. On the computational inefficiency of large batch sizes for stochastic gradient descent. arXiv preprint arxhv:1811.12949) 2018
10

---

P. Goyal, P. Dollár, R. Girshick, P. Noordhuis, L. Wesolowski, A. Kyrola, A. Tulloch, Y. Jia, and K. He. Accurate, large minibatch SGD: Training ImageNet in hour. arXiv preprint preprintarxiv: wrw.i::706.00000000 2017.
C. Hardin. Does batch size matter?, 2017. does- baae-bbaaa--iiz--tteerr/
URL httpa:////loo..aaestaoote com/
Izmailov, D. Podoprikhin, T. Garipov, D. Vetrov, and A. G. Wilson. Averaging weights leads to wider optima and better generalization. arXiv preprint
2018.
S. Jordan, Y. Chandak, D. Cohen, M. Zhang, and P. Thomas. Evaluating the performance of reinforcement learning algorithms. In International Conference on Machine Learning, pages 4962-4973. PMLR, 2020.
S. M. Kakade. A natural policy gradient. Advances in neural information processing systems, 14: 1531-1538, 2001.
D. P. Kingma and Ba. Adam: A method for stochastic optimization. arXiv preprint 1110: 2014.
H. Kushner and G. G. Yin. Stochastic approximation and recursive algorithuss and applications, volume 35, Springer Science & Business Media, 2003.
L. M. Vange and G. G. Koch. Rank score tests. Circulation, 1 (144233)2528-25533 2006
S. Ma. R. Bassily, and M. Belkin, The power of interpolation: Understanding the effectiveness of SGD in modern over per-parametrized learning. In International Conference on Machine Learning. pages
PMLR, 2018.
S. Mandt, M. D. Hoffman, and D. M. Blei. Stochastic gradient descent as approximate Bayesian inference. arXiv preprint arxiv: 704894289999. 2017.
McCandlish, Kaplan, D. Amodei, and the OpenAl Dota Team. An empirical model of large-batch training. arXiv preprint grXi:: 2018
V. Mnih, A. P. Badia, M. Mirza, A. Graves, T. Lillicrap, T. Harley, D. Silver, and K. Kavukcuoglu. Asynchronous methods for deep reinforcement learning. In International conference on machine learning, pages
PMLR, 2016.
G. Neu, A. Jonsson, and V. Gómez. A unified view of entropy-regullarize markov decision processes. arXiv preprint 17X5.07005.07993 2017.
OpenAl, C. Berner, G. Brockman, B. Chan, V. Cheung, P. Dbiak, C. Dennison, D. Farhi, Q. Fis- cher, S. Hashme, C. Hesse, R. Józefowicz, S. Gray, C. Olsson, J. Pachocki, M. Petrov, H. P. de Oliveira Pinto, J. Raiman, Salimans, J. Schlatter, Schneider, S. Sidor, I. Sutskever, J. Tang, F. Wolski, and S. Zhang. Dota 2 with large scale deep reinforcement learning. arXiv preprint
1922v:191..06680 2019.
J. Schulman, S. Levine, P. Abbeel, M. Jordan, and P. Moritz. Trust region policy optimization. In International conference on machine learning, pages 1889-1897, 2015a.
Schulman, P. Moritz, Levine, M. Jordan, and Abbeel. figh-ddimensinna continuous control using generalized advantage estimation. arXiv preprint 106.001556.00238 2015b.
J. Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: 1707.06777 2017.
C. Shallue, J. Lee, J. Antognini, J. Sohk-Dickstein, R. Frostig, and G. E. Dahl. Measuring the effects of data parallelism on neural network training. arXiv preprint arxh::1871.03600 2018.
S. L. Smith and Q Le. A Bayesian perspective on generalization and stochastic gradient descent. arXiv preprint prx:i:1110.06444 2017.
I.. Smith, P.-J. Kindermans, C. Ying, and Q. Le. Don't decay the learning rate, increase the batch size. arXiv preprint arXiv: 11t.x::1111000883 2017.

---

R. S. Sutton, D. McAllester, S. Singh, and Y. Mansour. Policy gradient methods for reinforcement learning with function approximation. Advances in neural information processing systems, 12. 1999.
M. Tomar, L. Shani, Y. Efroni, and M. Ghavamzadeh. Mirror descent policy optimization. arXiv preprint 0rxv1:2055.098/ 2020.
van Elteren. On the combination of independent two sample tests of Wilcoxon. Bull Inst Intern Staist, 37:351-361, 1960.
R. J. Williams. Simple statistical learning. Machine learning, 8(3):229-22 256, 1992.
algorithms for connectionist reinforcement
G. Zhang, L. Li, Z. Nado, J. Martens, Sachdeva, G. Dahl, Shallue, and R. B. Grosse. Which algorithmic choices matter at which batch sizes? Insights from noisy quadratic model. Advances in neural information processing systems, 92:8196-8207, 2019.
Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations your work? [Yes] There are careful ablations and a
discussion of the additional computational cost in Section [
(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 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 experi- mental results (either the supplemental material or as 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 the random seed after running experi- ments 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 [] from which the total computation cost of the experiments can be calculated. The type of resources used proprietary information.
4. If you are using existing assets (c.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'l 're
using/curating? [N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable 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, applicable? [N/A]
12

---

(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent participant compensation? [N/A]
13


Reviews:
[{"review_id": "rLZRJLfem7", "reviewer": "Reviewer_PXCQ", "summary": "# \n\nThis paper studies how to achieve batch-invariance in the policy optimization methods, e.g., PPO. Changing batch-size (the samples an on-policy algorithm collects per update) has a significant influence on policy learning for existing algorithms; batch-invariance will help to increase the batch size with limited computational resources.\n\nIn vanilla gradient descent, one can achieve batch invariance by adjusting learning rates, but the learning speed of the policy is affected by both the learning rate and the KL constraints. The paper first distinguishes the behavior policy and the proximal policy and observes that the two do not need to be equal. The reason is that though we need to use the exact behavior policy that collects the data to compute the objective for a good gradient estimation, its age does not matter too much. In contrast, we do not need to use the proximal policy in the KL constraint to sample data, but its age is directly related to how fast the policy can change.\n\nThe paper designs a decoupled clipped objective. Such decoupling allows the authors to use an exponentially-weighted moving average (EWMA) of the policy network as the proximal policy to maintain invariance when the algorithm has different policy update sizes. With additional batch-invariance techniques in the gradient descent, the method can maintain similar performance when the batch size changes.", "questions": "see above", "limitations": "yes", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Actually I like this paper. \n\n- Batch size has huge influence on policy learning performance but is not emphasized in many publications. I especially appreciate the advice in 6.2. I have seen many beginners ignore its importance and struggle to tune the parameters in PPO. The comprehensive study of the batch size and the practical advice would be excellent guidance for them to understand the optimization process. I would recommend my colleague to read the paper when it is published.\n- The decoupled objective is novel to me. It provides a new perspective on the constraints and objectives in PPO and could help us understand the dynamics of policy optimization. It reminds me of the distributional RL, where we can train multiple agents but have to use a centralized policy to constrain their differences. I wonder if there is any connection between the two.\n- I think this paper provides a comprehensive recipe to make PPO batch-invariant. The modifications is very practical and perform good in experiments. The technique contribution is solid. It is a pity that EWMA does not have too much effects.\n\nAs for the weakness:\n\n- I only found experiments on procgen. I wonder how it will behave on continuous control tasks, where the original PPO paper did its experiments. I am also curious about if PPG will differ from PPO in those domain.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "Y1JZirSE3dU", "reviewer": "Reviewer_vc65", "summary": "This paper investigates how to make on-policy surrogate-objective-based algorithms PPO and PPG batch-size invariant: adjustments to their batch size can be compensated for by automatic adjustments to other hyperparameters. To enable this, they discuss a different between two uses of the old policy in the calculation of the PPO/PPG objective, and propose a decoupling of these uses. This decoupling then allows them to propose variants of PPO and PPG that are batch-size invariant, and they demonstrate across the procgen task suite that their adjustments mostly produce batch-size invariant results. They show their variants also outperform the standard versions of PPO and PPG.", "questions": "L171: what exactly do you mean here? What hyperparameter is being set to the output of this calculation?", "limitations": "I think the limitations are well-discussed.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper is well-written and clearly describes batch size invariance, the decoupling of behaviour and proximal policy in the PPO objective, and the adjustments to PPO and PPG to make them batch-size invariance. The experimental results back up their claims that the adjustments do make PPO/PPG batch-size invariant, and the fact that their adjustment improves these algorithms is a nice added benefit. The contribution is novel and of some use to the community.\n\nI think the main weakness is that the desire to have these algorithms batch-size invariant isn't as well-motivated as it could be. There is some discussion in the introduction, but it would be useful to discuss practical cases where batch-size invariance would be beneficial. Further, the discussion that there is (L32-35) motivates that increasing the effective batch size (to reduce variance) through these methods would be useful, but there aren't any experiments showing that these adjustments can be used to increase the batch size to actually reduce the gradient variance, and so an improvement to the paper would be to formulate those kinds of experiments.\n\nAlso, given you're evaluating on Procgen already, it would be great to see generalisation experiments - training on the standard subset of levels and reporting the evaluation score on the unseen levels. I'd expect the results to look similar to the current ones (i.e. that batch size invariance is maintained even when testing for generalisation performance), but it would be great to empirically validate that.\n\nEDIT: I've read the rebuttal and feel it's addressed my concerns, but I'll keep my score at a 7 as I think that's the level of the paper.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "NRVkI4KiM6M", "reviewer": "Reviewer_SfJ4", "summary": "The authors propose an approach to make policy optimization algorithms such as PPO batch size invariant, where essentially performance achieved by the algorithm can be preserved even after changing batch sizes by changing relevant hyperparameters accordingly.\n\nThanks to the authors for putting in the effort in doing this work!", "questions": "Questions:\n- I am not clear why batch size invariance would be a property we want. Is it that larger batches are preferred due lower variance in gradient updates (as mentioned in the paper), but that is not possible due to GPU load, so ideally we want an algorithm where we can have a smaller batch size but achieve performance as though we were using a larger batch size?\n- On line 76, the paper mentions another work that proposes the similar objective, but the paper doesnt make any comparisons. I am curious how their modifications relate to yours and why theirs could not be used here?\n- In Figure 1(c), the paper says “holds back learning”. But I don’t see why this is true. Comparing to Figure 1(b), the performances in 1(c) don’t seem to be “held back”\n- Maybe I am misunderstanding, but shouldn’t Figure 4 (right) be 4 sets of bar graphs: one for PPO-EWMA, PPO, PPG-EWMA, and PPG? Here there seem to be only two.\n\nSuggestions:\n- The paper lacks a dedicated literature review, which would be useful for context.\n- As presented, the proposed improvement seems to be marginal (and at memory cost). I wonder if the benefit of this suggestion can be made clearer in a setting where the error due to coupled objective compounds and then introducing the decoupled objective can avoid that error compounding (this is a very abstract suggestion, but I am trying to see how the proposed suggestion can show much larger performance improvements). Are there actor critic algorithms that also do this type of coupling? I wonder if in those cases there is error in actor and critic estimations due to coupling, but by introducing this decoupling concept one could improve learning in both actor and critic, thus having a much better improvement over the vanilla version.\n- Minor: FYI, the appendix was attached with the main paper.\n", "limitations": "Yes. This is more fundamental work with no direct societal impact.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The paper is well written and flows well.\n- Labelling of the equations in Section 2 was very helpful.\n- Figure 3 was helpful.\n- I appreciate the acknowledgement in line 244 that the improvement is small with EWMA.\n\nWeaknesses:\n- The purpose of creating a batch size-invariant algorithm is not absolutely clear. It may be obvious to some folks but it wasn't to me. Line 35 attempts to give some explanation regarding data parallelism, but still was not absolutely clear (see questions below).\n- Given the marginal improvement (as they’ve also acknowledged) and the additional memory cost (with EWMA) it's unclear to me that this would be readily adopted. I think it's a very interesting find, but it's unclear if it makes enough improvement that this would be used instead of vanilla PPO (see suggestions below).\n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "IMmpLdqVKa", "reviewer": "Reviewer_GBX9", "summary": "In the context of reinforcement learning (RL), the paper introduces one property: batch size-invariance and achieving it by decoupling the proximal policy from the behavior policy. The authors use an exponentially-weighted moving average version of learning policy to achieve this property and show improvements on several experiments.", "questions": "What is the main difference between gradients accumulation and your methods?\n\nCould you please give some intuitive examples of the necessity of batch size-invariance?\n", "limitations": "In the absence of a more explicit application scenario for their method, it would be better to find a practical application like asynchronous and multi-machine.", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Originality: the proposed method is novel. It is somewhat iterative, decoupled policy objective is used in recent works but applies exponentially-weighted moving average for batch size-invariance that has not been noticed before.\n\nClarity: The paper is well written and clear. The method is soundness and easy to implement. The paper contains many various and multi-level experiments which prove the effectiveness of their method. \n\nSignificance: Although the authors mention in the introduction and section 6.2 that using a large batch-size will give a low variance estimation of policy gradient, and it is often prohibited by computational resources such as GPU memory, so we should use their method. But there is a more simple way to achieve this goal like accumulating gradients from multiple back propagation and updating until calculating as many samples as we need. Is there anything I mistake? And as experiments shown (fig 1), the best choice of behavior policy is the most recent policy, which is obvious for on-policy methods and most researchers use this situation as default. Overall, I think this paper needs more effort to make the motivation.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept