Batch size-invariance for policy optimization

1 Introduction

• PPO’s usual behavior-policy justification is flawed because update stability depends on controlling policy-update speed, not matching the data-collection policy.The paper connects this control to approximating the natural policy gradient.Source: §1 · p. 1
• The old policy serves two roles: behavior-policy importance sampling and proximal-policy control of update size.These roles need not use the same policy.Source: §1 · p. 1
• PPO-EWMA and PPG-EWMA decouple those roles by using an exponentially weighted moving average for the proximal policy.The variants are introduced to support batch size-invariant policy optimization at small batch sizes.Source: §1 · pp. 1–2
• The paper presents experiments testing both the benefits of decoupling proximal and behavior policies and the resulting batch size-invariance.The introduction identifies these as the central experimental claims.Source: §1 · p. 2

2 Decoupled policy objectives

• PPO uses its old policy both for importance sampling and as a recent target that pulls the current policy toward it.The first role requires the sampling policy, while the second defines the proximal policy.Source: §2 · p. 2
• The proximal policy need not equal the behavior policy; its age matters, whereas whether it collected the data does not.This is the paper’s central conceptual distinction.Source: §2 · p. 2
• The clipped objective can likewise be rewritten with a ratio against the proximal policy, while retaining behavior-policy correction.When β = 0 or ϵ = ∞, the proximal-policy dependence disappears and vanilla importance-sampled policy gradient is recovered [Schulman et al., 2017, equation (6)].Source: §2 · pp. 2–3
• The decoupled KL objective uses the behavior policy for sampling correction and a separately specified recent proximal policy for the KL penalty.The roles are explicitly assigned in the reformulated objective.Source: §2 · p. 2

3 Batch size-invariance

• Batch size-invariance is approximate recovery of original behavior after changing batch size and compensating with other hyperparameters.The paper treats it as a degree rather than a binary property, and also calls it “perfect scaling.”Source: §3 · p. 3

3.1 Batch size-invariance for stochastic gradient descent

• SGD remains batch size-invariant below a critical batch size, where changing batch size can be offset by a proportional learning-rate adjustment.The critical batch size corresponds to a gradient signal-to-noise ratio of around 1.Source: §3.1 · p. 3

3.2 Batch size-invariance for Adam

• Adam is also batch size-invariant below a critical batch size, using a square-root step-size rule when batch size changes.For dividing batch size by c, the step size is divided by √c; β1 and β2 may also be raised to 1/c.Source: §3.2 · p. 4

3.3 Batch size-invariance for policy optimization

• Policy optimization has both optimization and iteration batch sizes, and invariance requires changing them together by the same factor.The paper measures invariance using differences in final performance across batch sizes.Source: §3.3 · p. 4
• Because PPO’s optimizer can inherit SGD or Adam invariance, the remaining challenge is making iteration batch size invariant.The next section addresses this by modifying policy optimization itself.Source: §3.3 · p. 4

4 PPO-EWMA and PPG-EWMA

• The method maintains an EWMA of policy weights and uses it as the proximal policy in a decoupled objective.This avoids storing network copies from every intermediate step.Source: §4 · p. 4
• Changing iteration batch size changes policy age, so the method adjusts the EWMA decay to keep proximal-policy age constant in environment steps.The adjustment multiplies (1/(1−βprox) − 1) by c when batch sizes are divided by c.Source: §4 · p. 5
• PPO-EWMA and PPG-EWMA adjust optimizer hyperparameters using SGD or Adam batch-size rules and scale advantage-estimation samples with batch size.For Adam, the step size is divided by √c; PPG also multiplies policy iterations per phase by c.Source: §4 · p. 5
• The invariance adjustments require a sufficiently small optimization batch size and one policy epoch.The experiments therefore use PPG-EWMA with Eπ = 1 by default.Source: §4 · p. 5
• For PPG, scaling policy iterations per phase holds the phase batch size constant and preserves policy- and auxiliary-phase dynamics.This addresses PPG’s additional phase batch size.Source: §4 · p. 5

5.1 Artificial staleness

• With decoupled objectives, performance degraded minimally until roughly 8 iterations of staleness, exceeding 500,000 environment steps.Both original-objective choices were harmed by even small staleness, whereas decoupling preserved performance much longer.Source: §5.1 · p. 6, Figure 1 · p. 6

5 Experiments

• The experiments used Procgen’s 16 environments and reduced parallel environment copies from 256 by factors of four to one.The batch-size experiments require the optimization batch size to remain sufficiently small.Source: §5 · p. 5, §5.2 · pp. 6–7

5.2 Batch size-invariance

• A difference of 0.052 in final mean normalized return separated the largest and smallest batch sizes, falling to 0.019 without Heist.PPG-EWMA was evaluated across batch sizes after adjusting hyperparameters for invariance; Heist remained an unexplained outlier.Source: §5.2 · p. 6, Figure 2 · p. 7
• The Adam step-size adjustment was most important, while advantage normalization mattered mainly at small batch sizes in noisy environments.Training became highly unstable at the smallest batch sizes without the Adam adjustment; the EWMA adjustment had a smaller effect.Source: §5.2 · p. 6, Figure 2 · p. 7

5.3 EWMA comparison

• EWMA produced small but consistent gains, outperforming non-EWMA PPO and PPG on all last 8 Procgen environments.The authors attribute this benefit to reduced proximal-policy variance.Source: §5.3 · p. 8, Figure 4 · p. 8
• EWMA gains increased computational cost by 30% for PPO and 2.3% for PPG, excluding environment-stepping cost.The modification also requires extra memory and an additional EWMA-network forward pass per policy-gradient step.Source: §5.3 · p. 8

6.1 PPO as a natural policy gradient method

• Decoupling the proximal and behavior policies supports stale-data reuse, batch size-invariance, and slight sample-efficiency gains, motivating a natural-policy-gradient interpretation of PPO.The authors argue that PPO need only avoid moving too fast relative to a recent policy, not remain close to the behavior policy.Source: §6.1 · pp. 8–9
• The claimed interpretation conflicts with earlier evidence favoring behavior-policy constraints, which the authors attribute to historically easier hyperparameter tuning for constraint methods.Source: §6.1 · p. 9

6.2 Practical advice for policy optimization at small batch sizes

• Batch size controls gradient variance and may be constrained by hardware, so batch size-invariance lets practitioners retune hyperparameters for computationally feasible batch sizes.Source: §6.2 · p. 9
• At small batch sizes, tune the learning rate first, initially use one policy epoch, and monitor clipping and advantage-normalization diagnostics to detect unsuitable settings.Source: §6.2 · p. 9

7.1 Policy optimization

• PPO evolved from trust-region policy optimization and incorporates actor-critic and generalized-advantage-estimation components from prior work [Schulman et al., 2015a, 2017; Mnih et al., 2016; Schulman et al., 2015b].Source: §7.1 · p. 9
• Mirror descent provides another perspective on PPO’s surrogate objective and supports decoupling the proximal policy from the behavior policy, paralleling concurrent work by Tomar et al. [2020].Source: §7.1 · pp. 9–10
• Iteration batch size and optimization batch size are distinct in policy optimization, and invariance changes both together to model altered data-parallelism.Source: §7.1 · p. 9

7.2 Batch size-invariance

• Prior SGD studies established small-batch invariance through learning-rate adjustments, with empirical validation and rules involving momentum and training-set size [Mandt et al., 2017; Goyal et al., 2017; Smith et al., 2017].Source: §7.2 · p. 10
• Subsequent work defined critical batch size and perfect scaling, and related scaling limits to architecture, data complexity, gradient noise, and curvature [Ma et al., 2018; Shallue et al., 2018; McCandlish et al., 2018; Zhang et al., 2019].Source: §7.2 · p. 10

8 Conclusion

• PPO-EWMA and PPG-EWMA use an exponentially weighted moving average as the proximal policy, enabling more efficient stale-data use, slight sample-efficiency gains, and batch size-invariance when batches are not too large.Source: §8 · p. 10

A Pseudocode for PPO-EWMA

• The PPO-EWMA pseudocode updates the proximal-policy weights through an EWMA while retaining the current policy as the behavior policy.Source: A Pseudocode for PPO-EWMA · p. 13
• For PPG-EWMA, the EWMA modification applies to the policy phase, is reinitialized before each policy phase, and leaves the auxiliary phase unchanged.Source: A Pseudocode for PPO-EWMA · p. 13

B Hyperparameters

• Experiments use Procgen hard difficulty, no frame stack, an IMPALA convolutional network, and generally 100 million environment steps.Source: B Hyperparameters · p. 14
• PPG-EWMA sets the EWMA center of mass to eight minibatches, matching the maximum proximal-policy age between PPG and PPG-EWMA; tuning used only the first eight of sixteen Procgen environments.Source: B Hyperparameters · p. 14
• Batch size-invariance experiments reduce parallel environments and adjust policy-phase step sizes, EWMA decay, normalization, and policy iterations, while limiting auxiliary-phase adjustments to avoid excessive wall-clock time.Source: B Hyperparameters · pp. 14–15

• When batch size is divided by c, SGD divides its learning rate by c, whereas Adam’s step size should be adjusted using a square-root rule.Source: C Adam square root step size adjustment · p. 16
• For batches below the critical batch size, Adam’s gradient-noise denominator scales approximately as √c after a c-fold batch reduction, motivating the additional √c step-size adjustment.Source: C Adam square root step size adjustment · p. 16
• The analysis usually neglects Adam’s ϵ, which can matter when it is large enough to interpolate toward momentum SGD.Source: C Adam square root step size adjustment · pp. 16–17
• The square-root adjustment clearly outperformed dividing Adam’s step size linearly by c in the ablation, compared with the main batch size-invariance results.Source: C Adam square root step size adjustment · p. 16, Figure 5 · p. 16

• For all except one environment, the largest-versus-smallest batch-size difference in normalized return was at most 0.11.Source: D Adam β1 and β2 adjustments · p. 18
• Heist showed a substantially larger difference of 0.55, motivating tests of β1 and β2 adjustments.Source: D Adam β1 and β2 adjustments · p. 18
• Adjusting only β2 or both β1 and β2 still left a large performance difference between the largest and smallest batch sizes.Source: D Adam β1 and β2 adjustments · p. 18, Figure 6 · p. 18

• The primary metric was the difference in final algorithm performance across batch sizes.Source: E Hypothesis tests for ablations · p. 19
• The study tested 16 ablation-by-batch-size hypotheses with stratified van Elteren tests across environments.Source: E Hypothesis tests for ablations · p. 19
• Tests averaged performance over the final 4 million timesteps, used a 0.1% significance level, and applied Bonferroni correction.Source: E Hypothesis tests for ablations · p. 19, Table 5 · p. 19
• Ablation (b) differed significantly only at the smallest batch size, whereas other ablations were significant at every batch size except the largest for ablation (d).Source: E Hypothesis tests for ablations · p. 19

F.1 Artificial staleness

• Figures 7–9 split the artificial-staleness results across individual environments for three PPO policy-coupling variants.Source: Figure 7 · p. 20, Figure 8 · p. 21, Figure 9 · p. 22

F.2 Batch size-invariance

• Without advantage normalization, performance degrades only when advantage-standard-deviation estimates oscillate by roughly a factor of 10 or more.Source: Figure 13 · p. 26

F.3 EWMA comparison

• The EWMA comparison and decay-rate experiments report environment-level or StarPilot performance using multiple seeds, with Figure 18 testing batch-size-adjusted PPG-EWMA settings.Source: Figure 17 · p. 30, G Role of the proximal policy EWMA… · p. 31, Figure 18 · p. 32

G Role of the proximal policy EWMA decay rate

• Locally, halving COMprox while doubling the KL penalty produces the expected diagonal performance pattern, but repeated adjustments gradually degrade performance.Source: G Role of the proximal policy EWMA… · p. 31
• The proposed compensation is limited because reducing COMprox increases the variance of θ −θprox through a smaller effective averaging sample and lower signal-to-noise ratio.Source: G Role of the proximal policy EWMA… · p. 31