Scalable Reinforcement Learning via Adaptive Batch Scaling

Conventional wisdom holds that large-batch training is fundamentally incompatible with Reinforcement Learning (RL) — beyond a modest threshold, increasing batch sizes typically yields diminishing returns or performance degradation due to the inherent non-stationarity of the data distribution. We challenge this view by observing that non-stationarity is not a fixed property of RL, but evolves throughout training: early stages exhibit rapid behavioral shifts that demand small batches for plasticity, whereas late stages approach a quasi-stationary regime where large batches enable precise convergence. Motivated by this observation, we propose Adaptive Batch Scaling (ABS), that dynamically adjusts the effective batch size according to the stability of the learning policy. Central to ABS is Behavioral Divergence, a novel metric that quantifies policy non-stationarity by measuring action-level shifts between consecutive updates, which we use to scale batch size inversely to policy volatility. Integrated with the Parallelised Q-Network (PQN) algorithm and evaluated on the ALE benchmark, ABS seamlessly reconciles early-stage plasticity with late-stage stable convergence. Strikingly, contrary to conventional wisdom, our results reveal that the combination of larger networks and larger batch sizes achieves the best performance — a scaling behavior previously thought to be unattainable in RL, now unlocked through adaptive batch control.

The Challenge Conventional wisdom holds that large-batch training is fundamentally incompatible with Reinforcement Learning (RL). Increasing batch sizes beyond a modest threshold typically leads to diminishing returns or performance degradation due to the inherent non-stationarity of RL data. ### Our Insight We challenge this view by observing that non-stationarity is not a fixed property, but evolves throughout training: * Early stages exhibit rapid behavioral shifts that demand small batches to maintain learning plasticity. * Late stages approach a quasi-stationary regime where large batches enable precise and stable convergence. ### Our Solution: Adaptive Batch Scaling (ABS) We propose Adaptive Batch Scaling (ABS), a novel method that dynamically adjusts the effective batch size based on policy stability. Central to ABS is Behavioral Divergence, a metric that quantifies non-stationarity by measuring action-level shifts between consecutive updates. ABS scales the batch size inversely to this policy volatility. ### Key Results Integrated with the Parallelised Q-Network (PQN) algorithm and evaluated on the ALE (Atari) benchmark, ABS seamlessly reconciles early-stage plasticity with late-stage stable convergence. Strikingly, our results reveal that the combination of larger networks and larger batch sizes achieves the best performance—unlocking a scaling behavior previously thought to be unattainable in RL.