- The paper presents ABS to dynamically adapt the rollout batch size based on Behavioral Divergence, effectively balancing bias and plasticity in RL.
- ABS scales with model capacity to enhance training stability and performance, achieving mean episodic return improvements up to +191.7% on Atari benchmarks.
- The method outperforms prior adaptive schemes, generalizing to PPO and off-policy algorithms, and setting a new standard for scalable reinforcement learning.
Overview of Scalable On-Policy RL via Adaptive Batch Scaling
The paper "Scalable On-Policy Reinforcement Learning via Adaptive Batch Scaling" (2605.21557) addresses the persistent problem of poor scalability in on-policy RL with regard to batch size. Unlike large-scale supervised learning, where the batch size can typically be increased alongside model size and compute, on-policy RL has historically suffered from performance degradation when batch sizes are scaled beyond a modest regime. The central contribution is the introduction of Adaptive Batch Scaling (ABS), a scheme which leverages an online, policy-driven stationarity metric—Behavioral Divergence—to dynamically control the rollout batch size during training.
Motivation and Background
The inability of RL to scale batch size in the same proportion as model capacity and compute presents a core bottleneck in leveraging large neural network architectures for sequential decision-making. Prior work demonstrates that large batch training stabilizes and accelerates optimization in supervised settings and is tightly linked to model performance scaling laws. Conversely, RL’s unique non-stationarity—arising from the evolving policy and its induced data distribution—results in a bias-variance tradeoff: small batches offer plasticity to track rapid policy shifts early, but induce variance in gradient estimates. The paper posits that the dynamic nature of non-stationarity in RL is not properly reflected by static batching, and proposes to explicitly measure and adapt to it.
Adaptive Batch Scaling: Methodology
ABS introduces a metric called Behavioral Divergence, quantifying action-level policy shifts between consecutive updates for a reference set of states. Formally, for deterministic policies, it measures the fraction of reference states for which the most probable action changes between two successive policies. The effective on-policy batch size is then adapted inversely to Behavioral Divergence: high divergence (policy volatility) mandates small batches for frequent updates, while low divergence (policy stationarity) supports larger batches to reduce gradient noise and expedite convergence. ABS is computationally efficient since it operates using only policy forward passes, without requiring gradient-level noise estimation.
Batch size scaling is implemented by dynamically adjusting the rollout length within predetermined bounds ([L_min, L_max]) according to the currently observed Behavioral Divergence, mapped via a log-linear schedule. Integration into Practical RL algorithms is demonstrated with minimal overhead.
Experimental Results and Quantitative Analysis
ABS is benchmarked primarily with the Parallelised Q-Network (PQN) on the Arcade Learning Environment (ALE), including full Atari-57 and the curated Atari-10 subset. The empirical claims can be summarized as follows:
- Consistent performance improvements are observed with ABS over vanilla PQN and both small and large fixed batch baselines. ABS achieves the best score in 8/10 Atari-10 environments, resolving the conventional trade-off between low-batch plasticity and high-batch stability.
- Scalability to large models: When combined with large, high-capacity backbones (multi-skip ResNet-style MLPs), ABS enables both increased final performance and training stability with large batch sizes—a regime where fixed large-batch training degrades performance.
- Efficiency relative to prior adaptive schemes: Compared with gradient noise scale (GNS) based dynamic batching, ABS significantly outperforms GNS both in sample efficiency and in practical compute, since it avoids per-sample backward passes.
- Generalization to other paradigms: ABS extends straightforwardly to PPO for continuous control (by replacing action-level divergence with KL divergence between successive policies) and to modern off-policy value-based algorithms such as BTR, yielding sample efficiency and stability benefits without algorithmic retuning.
Numerical highlights include mean improvements in episodic return ranging from +27.8% to +191.7% across various model and batch scaling settings on Atari, and faster, more stable convergence in MuJoCo continuous control tasks.
Theoretical and Practical Implications
The critical assertion—that optimal on-policy batch size ought to be dynamic, not fixed, and specifically scaled to policy stationarity—contradicts canonical RL wisdom. Moreover, the findings indicate that the “scaling laws” paradigm from supervised learning can be recovered in RL, provided that the right batching protocols are used.
The theoretical implication is that RL’s non-stationarity is a phase-dependent attribute. Early in training, distributional shifts and exploratory instability necessitate frequent updates with fresh data, whereas late in training the policy distribution approaches quasi-stationarity and benefits from high-throughput, low-variance optimization—aligning with the high-batch, large-model regime underpinning the recent success in supervised and self-supervised learning.
Practically, ABS enables new scaling strategies: high-capacity RL models, including those based on deep CNNs, ResNets, and, prospectively, Transformers, can leverage batch sizes commensurate with their expressive capacity and hardware scaling. This in turn can accelerate RL-driven applications in robotics, autonomous systems, and large-scale agent modeling, provided that distributional non-stationarity is duly accounted for.
Future Directions
The immediate extensions include:
- Systematic integration of ABS into Transformer-based RL architectures, as recent trends pursue agent scaling with attention-based backbones.
- Application to RLHF-style training of LLMs, where ABS could enable RL fine-tuning at batch scales previously considered unstable.
- Investigation of ABS in offline RL and high-dimensional, complex real-world domains, as well as co-design with automated exploration/exploitation balancing.
Safety and alignment concerns in scaling RL agents—especially as ABS facilitates more efficient large-model training—remain paramount and warrant parallel research.
Conclusion
"Scalable On-Policy Reinforcement Learning via Adaptive Batch Scaling" (2605.21557) provides a rigorous analysis and practical solution to the longstanding challenge of large-batch RL. Through Behavioral Divergence-driven dynamic batching, ABS reconciles the bias-plasticity trade-off and unlocks the stable scaling of batch size concomitant with model size. The framework demonstrates robust gains in both discrete and continuous control, across on-policy and off-policy regimes, establishing a new foundation for scalable RL. The research broadens the toolkit for scalable agent optimization, suggesting that RL can now more fully benefit from the scaling laws that have propelled progress in supervised learning.
