📊 statistics

Scalable On-Policy Reinforcement Learning via Adaptive Batch Scaling

This paper challenges the conventional belief that large-batch training is incompatible with Reinforcement Learning by proposing Adaptive Batch Scaling (ABS), a method that dynamically adjusts batch sizes based on policy stability to successfully combine large networks and large batches for superior performance.

Original authors: Jongchan Park

Published 2026-05-22

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Jongchan Park

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Goldilocks" Dilemma of AI Training

Imagine you are teaching a dog a new trick.

Small Batch Size: You give the dog one treat at a time, correcting it immediately after every single move. This is great when the dog is confused and learning fast. You can change your teaching style instantly if the dog isn't getting it.
Large Batch Size: You wait until the dog has practiced for an hour, then you give it a massive pile of treats and corrections all at once. This is efficient and gives a very clear picture of what the dog is doing on average, but it's slow to react if the dog suddenly starts doing something weird.

In the world of Artificial Intelligence (specifically Reinforcement Learning, or RL), there is a long-held belief that you must stick to the "Small Batch" approach (teaching one step at a time).

Why? Because the AI is constantly changing its own mind. As it learns, the world it sees changes. If you wait too long to give it feedback (using a large batch), the AI might have already changed so much that the feedback is useless or even confusing. It's like giving a driver instructions for a road that no longer exists.

The New Idea: "Adaptive Batch Scaling" (ABS)

The authors of this paper say, "Wait a minute. The AI doesn't change its mind at the same speed the whole time."

They propose a method called Adaptive Batch Scaling (ABS). Think of it as a smart teacher who changes their teaching style based on how stable the student is acting.

The Early Days (Chaos Mode): When the AI first starts, it is wild, exploring, and making huge mistakes. Its behavior changes rapidly from one second to the next.
- The Strategy: Use Small Batches. Keep the feedback coming fast and frequent. This allows the AI to stay flexible ("plastic") and adapt quickly to its own mistakes.
The Later Days (Stable Mode): As the AI gets better, it stops making wild swings. It starts to settle into a good routine. Its behavior becomes predictable and stable.
- The Strategy: Switch to Large Batches. Now that the AI is stable, you can wait longer to gather more data. This gives a super-precise, high-quality correction that helps the AI perfect its skills and reach the top performance faster.

The Secret Tool: "Behavioral Divergence"

How does the AI know when to switch from "Chaos Mode" to "Stable Mode"? It uses a new measuring stick called Behavioral Divergence.

The Analogy: Imagine the AI is a dancer.
- High Divergence: The dancer is flailing their arms, spinning wildly, and changing moves every second. The teacher sees this and says, "Okay, let's stop and correct you immediately!" (Small Batch).
- Low Divergence: The dancer is now performing a smooth, consistent routine. The teacher sees this and says, "Great, you're stable. Let's watch a whole minute of your dance before I give you a detailed critique." (Large Batch).

The paper introduces a math formula to measure exactly how much the AI's "dance moves" (actions) are changing between updates. If the moves are changing a lot, the batch size stays small. If the moves are steady, the batch size grows.

The Results: Breaking the Rules

For a long time, experts thought you couldn't use "Large Batches" in Reinforcement Learning because the data was too messy. They also thought that making the AI's "brain" (the neural network) bigger would make it harder to train.

This paper proves them wrong.

By using their "Adaptive Batch Scaling" method:

They broke the limit: They successfully trained AI agents using much larger batches than ever before, which usually causes AI to fail.
They scaled up: They combined these large batches with massive AI brains (larger neural networks). In almost every other field of AI (like image recognition or language models), bigger brains + bigger batches = better results. In RL, this was thought to be impossible.
The Outcome: Their method worked better than the old "small batch only" methods on standard video game tests (Atari games). It learned faster at the start and finished stronger at the end.

Summary

The paper argues that the "one size fits all" approach to training AI is wrong. Instead of forcing a small batch size forever, we should let the batch size grow as the AI becomes more stable. By measuring how much the AI's behavior is changing, the system automatically switches between "fast, flexible learning" and "slow, precise polishing," unlocking the ability to train much smarter and larger AI agents.

Technical Summary: Scalable On-Policy Reinforcement Learning via Adaptive Batch Scaling

1. Problem Statement

Conventional wisdom in Reinforcement Learning (RL) posits that large-batch training is fundamentally incompatible with the field. Unlike supervised learning, where increasing batch sizes stabilizes gradient estimation and enables the training of massive models (Scaling Laws), RL practitioners typically rely on modest batch sizes (e.g., 256–1024). This limitation is attributed to the inherent non-stationarity of the data distribution in RL; as the policy ( $\pi_\theta$ ) updates, the data distribution shifts, creating a bias-variance tradeoff.

The prevailing view suggests that increasing batch sizes beyond a critical threshold yields diminishing returns or performance degradation. Large batches are believed to cause agents to overfit to transient, unstable trajectories during the early, volatile stages of training. Consequently, standard practice favors small batches to maintain model plasticity and adapt to shifting distributions, even when using large-scale architectures. This paper challenges the static nature of this view, arguing that the degree of non-stationarity is not constant but evolves throughout training.

2. Methodology: Adaptive Batch Scaling (ABS)

The authors propose Adaptive Batch Scaling (ABS), a framework that dynamically adjusts the effective batch size based on the stability of the learning policy. The core hypothesis is that RL training consists of two distinct regimes:

Early Stage (High Non-Stationarity): Characterized by rapid behavioral shifts. Small batches are required to preserve plasticity and prevent overfitting to unstable trajectories.
Late Stage (Near-Stationarity): As the policy converges, the data distribution stabilizes. Large batches become beneficial to suppress gradient noise and enable precise convergence.

Key Components

Behavioral Divergence ( $\delta_\pi$ ): A novel metric introduced to quantify policy non-stationarity. It measures the fraction of states where the preferred action changes between the current policy ( $\pi_\theta$ $π_{θ}$ ) and the previous policy ( $\pi_{\theta_{old}}$ $π_{θ_{o l d}}$ ) on a reference batch.
- Formally: $\delta_\pi(\theta', \theta) := \frac{1}{M} \sum_{i=1}^{M} \mathbb{I} [\pi_{\theta'}(s_i) \neq \pi_{\theta}(s_i)]$ .
- High $\delta_\pi$ indicates high volatility (early stage); low $\delta_\pi$ indicates stability (late stage).
Dynamic Scaling Schedule: The effective batch size ( $|B|$ $∣ B ∣$ ) is scaled inversely proportional to the behavioral divergence.
- The rollout length ( $L_{adapt}$ ) is adjusted logarithmically between a minimum ( $L_{min}$ ) and maximum ( $L_{max}$ ) based on the normalized divergence.
- When $\delta_\pi$ is high (rapid change), $L_{adapt}$ is set to $L_{min}$ (small batch).
- When $\delta_\pi$ is low (stable), $L_{adapt}$ scales to $L_{max}$ (large batch).
Integration: The method is integrated with the Parallelised Q-Network (PQN) algorithm. The calculation of Behavioral Divergence relies solely on forward passes, making it computationally efficient compared to methods requiring backward passes (e.g., Gradient Noise Scale).

3. Key Contributions

Quantification of RL Non-Stationarity: The introduction of Behavioral Divergence as a practical, forward-pass metric to measure action-level shifts between successive policy updates.
Adaptive Batch Scaling Framework: A simple yet effective algorithm that dynamically scales batch sizes (rollout lengths) in response to policy stability, bridging the gap between the need for plasticity in early training and stability in late training.
Scalability Verification: Empirical demonstration that ABS enables performance improvements with larger batch sizes on the ALE benchmark, particularly when combined with larger network architectures (e.g., ResNet-style multi-skip models), a scaling behavior previously thought unattainable in RL.

4. Experimental Results

The method was evaluated on the Atari-57 and Atari-10 benchmarks using the PQN algorithm, and extended to continuous control (MuJoCo) and off-policy settings (BTR).

Performance vs. Fixed Batches: ABS outperformed vanilla PQN and fixed small/large batch baselines across most environments. It achieved higher early-stage scores (due to small batches) and lower late-stage variance (due to large batches).
Comparison with GNS: ABS outperformed the Gradient Noise Scale (GNS) approach in both sample efficiency and final performance, while being significantly more computationally efficient as it avoids additional backward passes.
Scalability with Large Models: When applied to larger architectures (PQN-L and PQN-XL), ABS consistently enhanced performance. Crucially, while using a fixed large batch degraded performance on large models, ABS (scaling from small to large) unlocked the potential of these models, achieving the best performance. This confirms a positive correlation between model scale and the efficacy of adaptive batching.
Generalization:
- Continuous Control: Applied to PPO on MuJoCo using KL Divergence as the non-stationarity proxy, ABS improved sample efficiency and final performance.
- Off-Policy RL: Applied to BTR with minimal modification (replacing rollout length with batch size range), ABS improved performance in 4 out of 5 Atari-5 environments.

5. Significance and Claims

The paper claims to challenge the "conventional wisdom" that large-batch training is incompatible with RL. By demonstrating that non-stationarity is an evolving property rather than a fixed constraint, the authors show that larger networks and larger batch sizes can achieve the best performance in RL when managed adaptively.

The significance of this work lies in its potential to unlock the scaling potential of reinforcement learning, allowing high-capacity RL agents to utilize the same large-batch training advantages that have driven success in supervised learning (e.g., LLMs and Vision Transformers). The authors position ABS as a foundational step toward scalable RL, bridging the performance gap between RL and supervised learning at large scales.

Limitations & Future Directions:
The authors note that while the current work focuses on on-policy algorithms, future work includes extending ABS to offline RL, validating on Transformer-based architectures, and integrating the method into Large Language Model alignment techniques (e.g., GRPO) for scalable RLHF.