1 Introduction↩︎

The prevailing paradigm in modern Deep Learning is defined by the ‘Scaling Laws’, where performance monotonically improves with increased model size, dataset size, and computational budget [1]. In domains such as Computer Vision (CV) and Natural Language Processing (NLP), this has crystallized into a strategy of training massive models with correspondingly massive batch sizes. Large batch training not only stabilizes the gradient estimation by reducing variance but also maximizes parallelization efficiency on distributed hardware accelerators. Pioneering works in large-scale optimization, such as LARS [2] and LAMB [3], have enabled training ResNet-50 on ImageNet in minutes using batch sizes of up to 32k. Similarly, the success of Large Language Models (LLMs) like GPT-3, DeekSeek-R1 [4], [5] and Vision Transformers (ViT) [6], [7] relies heavily on global batch sizes comprising millions of tokens to effectively saturate compute capacity and maintain training stability.

However, Reinforcement Learning (RL) presents a stark contrast to this trend. Despite the adoption of scalable architectures like Transformers, the effective batch sizes in RL remain surprisingly small compared to their supervised counterparts. For instance, the Distributed Impala architecture [8] and even high-capacity Robotics Transformers (RT-1, RT-2) [9], [10] operate on surprisingly modest batch sizes (e.g., 256 to 1024) that do not scale in proportion to their massive parameter counts. This phenomenon extends to Large Language Model (LLM) training, where reinforcement learning stages utilize notably smaller batch sizes compared to the massive scales employed during supervised learning phases [5], [11]. Empirical studies consistently report that increasing the batch size beyond a critical threshold in RL often leads to diminishing returns or even performance degradation [12], [13].

We attribute this phenomenon to the inherent non-stationarity of both the evolving policy and the data distribution generated through environment interactions, and it creates a fundamental bias-variance tradeoff. Unlike supervised learning, which operates on a fixed data distribution, RL agents must learn from a stream of experience that co-evolves with the policy updates (\(\pi_{\theta}\)). Recent studies have highlighted the fundamental difficulty of optimizing deep models under such shifting distributions, even in supervised contexts [14]. This challenge is particularly acute during the initial stages of RL training, where rapid policy shifts exacerbate distributional instability. To mitigate these effects, employing modest batch sizes is often preferred, as it allows the agent to adapt more fluidly to the non-stationary data stream [12].

This paper challenges the static view of batch size in RL based on a simple yet overlooked insight: the degree of non-stationarity is not constant throughout training. Figure 1 (Left, Blue curve) illustrates the evolution of the policy’s behavior change rate between successive updates (i.e., defined ‘Behavioral Divergence’ in this work)—over aggregated Atari-10 [15] environments. In the initial phase, intensive exploration and rapid parameter adjustments lead to high behavioral volatility, resulting in severe distribution shifts. In this highly non-stationary regime, maintaining small batch sizes (i.e., short rollouts for RL) is essential to preserve model plasticity (i.e, rapidly evolving policies) for minimizing bias and prevent the agent from overfitting to transient, unstable trajectories. However, as training progresses and the policy converges toward optimal behavior, the rate of behavioral change diminishes significantly, and the trajectory distribution approaches a near-stationary state. In this ‘near-stationary’ regime, large batches (i.e., long rollouts for RL) can help suppress gradient noise and ensure convergence by minimizing variance. We hypothesize that adhering to small batches in this later stage is suboptimal; instead, scaling the batch size—as per standard Deep Learning wisdom—becomes necessary to refine the policy with high-precision gradients and accelerate convergence.

Based on this hypothesis, we propose a novel simple yet effective training framework that dynamically adjusts the batch size (i.e., the rollout length of an RL) in response to the stability of the learning process. Specifically, we introduce a mechanism that scales the batch size inversely proportional to the policy behavioral change rate (Figure 1 (Left)). We demonstrate the efficacy of this approach by applying it to PQN[16] on the ALE benchmark [17]. Empirical results demonstrate that our Adaptive Batch Scaling (ABS) yields substantial performance improvements by synergizing early-stage stability with high-throughput optimization in the final phases. Notably, this approach effectively scales to larger backbone architectures [14], surmounting the performance bottlenecks inherent in conventional RL training. Given its robust performance and ease of integration, we expect our methodology to be a foundational key to unlocking the scaling potential of reinforcement learning, bridging the performance gap with supervised learning at large scale architectures.

Our contributions are summarized as follows:

• Quantification of RL Non-stationarity: We introduce Behavioral Divergence—a novel metric that quantifies the non-stationarity of RL training by measuring action-level shifts between successive policy updates.

• Adaptive Batch Scaling Framework: We propose a novel simple yet effective algorithm that dynamically scales the batch size based on policy stability metrics (i.e., Behavioral Divergence).

• Scalability Verification: We empirically demonstrate that our method enables performance improvements with larger batch sizes on the ALE benchmark, particularly when combined with larger network architectures.

2 Related Work↩︎

2.0.0.1 Batch Size Scalability in Reinforcement Learning.

Recent endeavors in Reinforcement Learning (RL) have sought to replicate the ‘Scaling Laws’ of supervised learning by scaling up model parameters and architectures [8]–[10], [14], [18]. However, unlike model size, the batch size—a critical hyperparameter for training speed and stability—has remained disproportionately small. This discrepancy is particularly evident in the RL fine-tuning stages of Large Language Models (LLMs); while Pretraining and Supervised Fine-Tuning (SFT) utilize massive batch sizes, the subsequent RL phases (e.g., RLHF) revert to significantly smaller batches [5], [11], [19]. Empirical studies suggest that smaller batch sizes generally yield better performance in standard RL benchmarks due to the inherent non-stationarity of the learning process [12], [13]. To overcome this, frequent updates with small batches are typically employed to track the shifting data distribution.

In this work, we view RL training through the lens of two distinct regimes: the early and late phases. The early phase is characterized by high policy volatility and significant non-stationarity, necessitating small batches for rapid adaptation. Conversely, as training progresses, the policy stabilizes, and the data distribution becomes ‘near-stationary’, resembling a supervised learning task. We propose a methodology that bridges these regimes by starting with small batches to mitigate early-stage non-stationarity and progressively scaling them up to enhance stability and convergence speed in the later stages.

2.0.0.2 Dynamic Batch and Sampling Mechanisms.

Prior research has explored dynamically adjusting batch size to improve efficiency [13], [19], [20]. In Model-Based RL, [20] proposed dynamically adjusting the rollout length to balance the trade-off between sample efficiency and model accuracy. In the context of data selection, [19] maintained a fixed batch size but adjusted the number of candidate samples to ensure each batch contained high-information data (DAPO).

Most relevant to our work is [13], which introduced the Gradient Noise Scale (GNS) to determine the optimal batch size at any given step. However, leveraging GNS directly for online RL training presents practical challenges: it requires additional backward passes to compute per-sample gradients, incurring significant computational overhead. Furthermore, while their analysis showed that the optimal batch size (inferred from increasing GNS over steps) increases over time, they did not experimentally validate a dynamic batch sizing schedule in RL, instead concluding that fixed small batches were empirically superior for the tested RL tasks—Figure 1 is the result that expands the GNS based batch scheduling to RL by us. Our work aligns with the theoretical implication of [13]—that batch size should grow—but diverges in execution. We propose a metric based on the policy action change rate (Behavioral Divergence), which relies solely on the forward pass, making it computationally efficient. We demonstrate that this lightweight, dynamic batch size scheduling strategy effectively handles the stationarity transition and improves performance across various ALE environments.

2.0.0.3 Policy Churn and Behavioral Divergence.

The instability of greedy policies during training has been studied under the notion of policy churn [21], [22]. [21] first documented that a single gradient update can alter the greedy actions across a substantial fraction of states, hypothesizing that this rapid fluctuation serves as a mechanism for implicit exploration. [22] further characterized a chain effect in which value churn and policy churn compound cyclically, leading to training instability, and proposed a regularization method (CHAIN) to suppress these oscillations and stabilize learning.

Our notion of Behavioral Divergence is conceptually related to policy churn in that both quantify the discrepancy between successive policies. However, the objectives diverge fundamentally: prior work treats churn as a phenomenon to be explained (as implicit exploration) or mitigated (via regularization), whereas we leverage Behavioral Divergence as an metric for non-stationarity in the learning process.

3 Preliminaries↩︎

We consider a standard Markov Decision Process (MDP) defined by the tuple \((\mathcal{S}, \mathcal{A}, P, R, \gamma)\), where \(\mathcal{S}\) and \(\mathcal{A}\) denote the state and action spaces, \(P\) the transition dynamics, \(R\) the reward function, and \(\gamma \in [0, 1)\) the discount factor. The objective of the agent is to learn a parameterized policy \(\pi_\theta(a|s)\) that maximizes the expected cumulative reward \(J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta} [\sum_{t=0}^\infty \gamma^t R(s_t, a_t)]\).

3.1 Bias-Variance in RL↩︎

In modern on-policy style algorithms such as PPO [23], the policy gradient is typically estimated using a batch of trajectories \(\mathcal{B}\) collected by the behavior policy \(\pi_{\theta_\text{old}}\). The objective gradient is approximated as: \[\hat{\nabla}_\theta J(\theta) \approx \frac{1}{|\mathcal{B}|} \sum_{(s,a) \in \mathcal{B}} \frac{\pi_\theta(a|s)}{\pi_{\theta_\text{old}}(a|s)} \hat{A}(s,a) \cdot \nabla_\theta \log \pi_\theta(a|s)\] where \(\hat{A}\) represents the advantage estimator (e.g., GAE). The quality of this gradient estimate is governed by the bias-variance tradeoff, which is directly influenced by the batch size \(B = |\mathcal{B}|\).

3.1.0.1 Variance (Estimation Error).

The variance of the gradient estimator arises from the stochasticity of the environment (transitions and rewards) and the policy itself. The variance of the gradient estimate scales inversely with the batch size: \[\text{Var}(\hat{\nabla} J) \propto \frac{1}{B}\]

Small batches (\(B \downarrow\)) yield high variance results in ‘noisy’ gradients. While traditionally seen as detrimental, this stochastic noise acts as a source of implicit exploration, preventing the policy from prematurely converging to sharp local minima. In contrast, larger batches (\(B \uparrow\)) average out the noise, providing a high-fidelity estimate of the true gradient direction [24].

3.1.0.2 Bias (Approximation & Shift Error).

In RL, bias primarily stems from two sources: (1) the bias in advantage estimation (e.g., due to bootstrapping in GAE), and (2) the distributional shift during multi-epoch updates. When optimizing \(\theta\) over multiple epochs using a fixed batch \(\mathcal{B}\), the data distribution \(d^{\pi_{\theta}}\) diverges from the collection distribution \(d^{\pi_{\theta_\text{old}}}\).

While large batches provide a stable signal for the collected distribution, they encourage the optimizer to commit strongly to a specific policy update direction. If the underlying data distribution is non-stationary (i.e., the optimal policy changes drastically), ‘over-optimizing’ on a large batch introduces a bias where the agent overfits to transient dynamics that may no longer be valid after the update.

3.2 The Dynamics of Non-stationarity: Early vs. Late Training↩︎

The central premise of our work is that the optimal trade-off between bias and variance is not static but evolves as training progresses.

3.2.0.1 Early Stage: High Non-stationarity & Plasticity.

At the beginning of training, the policy \(\pi_\theta\) is initialized randomly, leading to high behavioral volatility. Consequently, the distribution shift \(d^{\pi_k} \to d^{\pi_{k+1}}\) is drastic between updates. In this regime, the learning landscape is highly non-stationary.

Calculating a precise gradient (Large Batch) for a transient policy is computationally wasteful and potentially harmful. The ‘accurate’ gradient may point towards a local optimum of a distribution that is about to vanish. Paradoxically, the high variance from small batches acts as an implicit regularizer, preserving model plasticity and preventing the agent from overfitting to unstable trajectories [25].

3.2.0.2 Late Stage: Near-Stationarity & Convergence.

As the agent begins to master the task, the policy \(\pi_\theta\) converges towards a high-performing region, and the rate of behavioral change diminishes. The trajectory distribution approaches a quasi-stationary state.

In this stable regime, the gradient noise from small batches transitions from a feature (exploration) to a bug (instability), preventing the model from settling into the sharp global optimum. Here, scaling up the batch size reduces variance, allowing for high-precision updates required to refine the policy and accelerate asymptotic convergence [24].

3.3 Parallelised Q-Network (PQN)↩︎

In our experiments, we utilize the Parallelised Q-Network (PQN) proposed by [16] as the primary learning framework. PQN simplifies deep temporal difference (TD) learning by integrating the stability benefits of proximal optimization into value-based methods. Specifically, instead of relying on slowly-updating target networks, PQN employs a proximal objective that clips the Q-value updates relative to the previous estimates, preventing catastrophic divergence during frequent gradient steps. By applying our Adaptive Batch Scaling (ABS) to PQN, we aim to demonstrate that even highly efficient, modern RL architectures can further benefit from dynamic batching, particularly when scaling to larger neural backbones.

4 Method↩︎

In this section, we introduce Adaptive Batch Scaling (ABS), a training framework designed to align the batch size with the evolving stationarity of the reinforcement learning process. Our core insight is to utilize small batches during the highly non-stationary early training phase to maximize gradient update frequency and plasticity, while transitioning to larger batches in the stationary later phase to ensure stable convergence. We quantify stationarity via the rate of policy behavioral change (i.e., behavioral divergence) and employ this metric to dynamically adjust the data collection horizon (rollout length). We demonstrate the implementation of ABS within the Parallelised Q-Network (PQN) algorithm.

4.1 Quantifying Non-stationarity via Behavioral Divergence↩︎

In RL, the learning agent alternates between data collection using the current policy \(\pi_\theta\) and optimization updates. Unlike the fixed data distribution in supervised learning, the distribution of collected experience in RL is non-stationary, shifting continuously as the policy evolves. While the environment dynamics \(P(s'|s,a)\) generally remain static, the policy distribution \(\pi_\theta(a|s)\) induces a covariate shift in the visited state-action space.

This distributional shift is most severe during the early stages of training when policy updates result in significant behavioral changes. As training progresses and the policy converges toward an optimal behavior, the updates become incremental, leading to a ‘near-stationary’ joint distribution of the policy and environment.

To operationalize this observation, we propose Behavioral Divergence (\(\delta_\pi\)) as a proxy metric for measuring non-stationarity. Specifically, we measure the rate at which the updated policy \(\pi_{\theta'}\) deviates from a rollout policy \(\pi_{\theta}\). Formally, \(\delta_\pi\) is defined as the fraction of states where the preferred actions differ:

\[\delta_\pi(\theta', \theta) := \frac{1}{M} \sum_{i=1}^M \mathbb{I} \left[ \pi_{\theta'}(s_i) \neq \pi_{\theta}(s_i) \right], \;s_i \sim \mathcal{B}_{ref} \label{eq:divergence}\tag{1}\]

where \(\mathcal{B}_{ref}\) is a reference batch of \(M\) states sampled from recent trajectories (i.e., collected by \(\pi_\theta\), \(M=2048\)), and \(\mathbb{I}[\cdot]\) is the indicator function. Since we build upon PQN, the policy is deterministically derived from the Q-values: \(\pi_\theta(s) = \arg\max_a Q_\theta(s,a)\). As illustrated in Figure 1, \(\delta_\pi\) typically starts high due to random exploration and rapid initial learning, then monotonically decreases as the agent stabilizes.

4.2 Adaptive Batch Scaling (ABS)↩︎

Standard RL algorithms which don’t use replay buffer for learning define the global batch size \(|\mathcal{B}|\) as the product of the number of parallel environments (\(E\)) and the rollout length (\(L\)): \(|\mathcal{B}| = E \times L\). While \(E\) is typically constrained by hardware resources, \(L\) is a flexible hyperparameter that directly dictates the batch size and the freshness of the data.

ABS dynamically scales \(|\mathcal{B}|\) by adjusting the rollout length \(L\) in response to the measured Behavioral Divergence \(\delta_\pi\). We maintain a fixed number of mini-batches \(N_\text{mb}\) and optimization epochs \(N_\text{epochs}\).

We define a scaling schedule where the rollout length \(L_\text{adapt}\) increases as the divergence \(\delta_\pi\) decreases. To bridge the varying scales of rollout lengths smoothly, we employ a logarithmic interpolation between a minimum length \(L_{\min}\) and a maximum length \(L_{\max}\):

\[L_{\text{adapt}} = L_{\max} - \alpha\cdot(L_{\max} - L_{\min}) \label{eq:scaling}\tag{2}\]

where the interpolation factor \(\alpha\) is determined by the normalized divergence: \[\alpha = \frac{\log(\delta_{\pi,\text{cliped}} / \delta_{\min})}{\log(\delta_{\max} / \delta_{\min})}\]

Here, \(\delta_{\pi,\text{cliped}}=\text{clip}(\delta_\pi, \delta_{\min}, \delta_{\max})\) and \([\delta_{\min}, \delta_{\max}]\) represents the sensitivity range of the policy change rate. If the policy is changing rapidly (\(\delta_\pi \geq \delta_{\max}\)), we use \(L_{\min}\) to prioritize frequent updates. Conversely, if the policy is stable (\(\delta_\pi \leq \delta_{\min}\)), we scale to \(L_{\max}\) to minimize gradient noise.

Finally, the effective batch size for the next update is determined as: \[|\mathcal{B}| = E \cdot L_\text{adapt}\] This allows the effective mini-batch size used for gradient descent to be \(|\mathcal{B}| / N_\text{mb}\). To balance computational overhead and enhance the learning stability, we re-calculate \(L_\text{adapt}\) every \(K\) updates.

4.3 Implementation on PQN↩︎

We integrate our Adaptive Batch Scaling (ABS) framework into the Parallelised Q-Network (PQN) [16], a prominent recent off-policy RL algorithm. The complete training procedure is summarized in Algorithm 2.

Our implementation is built upon the CleanRL library¹. We uploaded our implementation of ABS on our github page² for reproducibility. Further implementation details and specific hyperparameter settings are provided in Appendix 7 and Appendix 9, respectively.

5 Experiments↩︎

In this section, we empirically validate the effectiveness of our Adaptive Batch Scaling (ABS). Our experiments are designed to answer the following research questions:

• Q1 (Performance). Does ABS improve performance compared to standard fixed-batch baselines and a prior dynamic batch approach?

• Q2 (Scalability). Can ABS stabilize the training and increase the performance of larger network architectures with scaled batch?

• Q3 (Sensitivity). How sensitive is ABS to its hyperparameters, and which 54components are most critical to its performance?

• Q4 (Generalization). Is ABS applicable to continuous control domains (e.g., PPO) and even to a recent off-policy algorithm (e.g., BTR) which involves a replay buffer ?

Subsequence is the experimental setup to answer the above questions:

5.0.0.1 Environments.

We evaluate our method on the full Arcade Learning Environment (ALE) [17], utilizing the Atari-57 benchmark. For ablation studies, we use the Atari-10 suite [15], a curated subset of 10 games identified to capture the maximum variance in algorithm performance and highly correlated (\(>0.9\)) with the full benchmark. Additionally, we extend our evaluation to continuous control tasks using MuJoCo environments via Gymnasium.

5.0.0.2 Baselines & Implementation.

Our primary baseline is PQN [16], a streamlined and high-throughput but effective RL algorithm. We adapt our ABS on PQN and compare to: (1) vanilla PQN, (2) PQN with smaller fixed batch size, (3) PQN with larger fixed batch size, and (4) PQN with Gradient Noise Scale (GNS) [13] based dynamic batch scheduling. For scalability experiments, we adopt the ResNet-style multi-skip architectures proposed in [14] (i.e., but using same optimizer with PQN, RAdam [26]).

5.0.0.3 Evaluation Protocol.

We train all agents for 20 million steps (5M for MuJoCo) across 3 independent seeds following standard protocols [27], [28]. For evaluation, we report the mean score over 100 evaluation episodes per seed (total 300 episodes) after training concludes.

5.1 A1: Enhancing Performance on Full Atari-57↩︎

We first evaluate ABS on the full Atari-57 benchmark. Figure 3 illustrates the relative performance gains of ABS compared to the vanilla PQN baseline, showing that our method outperforms the baseline across the most of environments. To further analyze the training dynamics, Figure 4 presents the learning curves on Atari-10, averaged over three seeds.

Table 2 shows early training performance (first 10% of steps) and late training stability (last 10% of steps) separately. ABS achieves higher early-stage scores in 7/10 environments, supporting the hypothesis that smaller batches accelerate initial learning. For late-stage variance, ABS achieves lower variance in 6/10 environments. These results validate our hypothesis: leveraging smaller batch sizes in the early stages accelerates initial learning, while transitioning to larger batches in the later stages stabilizes and enhances final performance.

5.2 A1: Outperforming Fixed Batches↩︎

To isolate the impact of batch size, we compare ABS against fixed Small Batch and Large Batch configurations on the Atari-10 suite. The Small Batch setting utilizes a rollout length of \(16\) (\(2{,}048\) samples), while the Large Batch setting uses a rollout length of \(64\) (\(8{,}192\) samples), compared to the default PQN’s \(32\) (\(4{,}096\) samples). Table [tab:atari-batches] summarizes these results with IQM HNS ± 95% confidence intervals. While Small Batch often outperforms the vanilla baseline by prioritizing plasticity and rapid updates, ABS achieves the superior score in \(8\) out of \(10\) environments. This success stems from its ability to resolve the fundamental trade-off between these regimes: ABS leverages small batches to navigate early-stage instabilities and maintain plasticity, then scales to large batches to minimize gradient variance for precise late-stage refinement.

5.3 A1: Outperforming Gradient Noise Scale (GNS)↩︎

max width=0.90

We compare ABS with the Gradient Noise Scale (GNS) metric proposed by [13], which suggests batch sizes based on gradient variance. While GNS is theoretically grounded, it requires computationally expensive backward passes to estimate noise. In contrast, ABS relies solely on forward pass inference (behavioral divergence), making it significantly more efficient.

As shown in the learning curves (Figure 4) and evaluation results (Table [tab:atari-batches]), ABS outperforms GNS-based scheduling in terms of sample efficiency and final performance. This suggests that our behavioral divergence is a more direct and practical way for measuring non-stationarity and batch scheduling in RL than gradient noise alone. Details and hyperparameters of the GNS is in Appendix 7.2 and 9.

5.4 A2: Unlocking the Potential of Scaling RL with ABS↩︎

In general deep learning, larger models typically necessitate larger batch sizes to ensure convergence and maximize final performance [29]. However, in standard RL, simply increasing batch size often leads to performance degradation on larger models. A central hypothesis of this work is that ABS can significantly enhance the performance of large-scale RL models by unlocking the potential of their vast parameter capacity through larger batch in late learning stage.

To verify this, we adopt ResNet-style Multi-Skip architectures from [14], specifically the PQN-L (CNN + 10-layer MLP) and PQN-XL (CNN + 20-layer MLP), for scaled large RL policy. Table 1 presents the IQM HNS with 95% confidence intervals on the Atari-10 benchmark. We find that ABS consistently and significantly enhances the performance of scaled networks, while a large fixed batch size conversely degrades performance across all model capacities.

Figure 5 presents the mean improvement rate of ABS over each baseline across the Atari-10 environments, organized into three bar plots corresponding to three model scales: PQN, PQN-L, and PQN-XL. Within each plot, the x-axis denotes the rollout range \([L_{\min} \to L_{\max}]\) (\(16{\to}32\); \(2{,}048{\to}4{,}096\) samples, \(16{\to}64\); \(2{,}048{\to}8{,}192\) samples, \(16{\to}128\); \(2{,}048{\to}16{,}384\) samples), and the y-axis denotes the mean improvement rate (%) over the respective baseline.

As model capacity and maximum rollout length (batch size) increase, the improvement of ABS becomes more pronounced. For PQN-XL, the largest model, ABS (\(16{\to}128\); \(2{,}048{\to}16{,}384\) samples) achieves the most substantial gains. For detailed results, please refer to Appendix 8.

Through these experiments, we demonstrate that ABS is particularly effective for large-scale models. Specifically, ABS enables high-capacity RL agents to stably exploit large batch sizes—an advantage that has long been established in supervised learning but has remained difficult to harness in RL due to the inherent non-stationarity of the learning process. By adaptively scheduling the rollout length in response to Behavioral Divergence, ABS bridges this gap, allowing scaled RL models to benefit from the same large-batch training advantages that have driven the success of large-scale supervised learning, including more stable gradient estimates, improved sample diversity, and stable convergence. Consequently, by unlocking stable access to large batch sizes, ABS fully realizes the potential of scaled RL: allowing high-capacity models to exploit their full representational capacity, and ultimately translating model scale into tangible performance gains.

5.5 A3: Hyperparameter Sensitivity↩︎

Figure 6 presents the hyperparameter sensitivity of ABS with respect to the rollout range \(L\) (Top), policy change thresholds \(\delta\) (Center), and adapt frequency \(K\) (Bottom) on Atari-10.

5.5.0.1 Rollout Range \([L_{\min}, L_{\max}]\).

As shown in Figure 6 (Top), setting \(L_{\min}\) too small (\(8\); \(1{,}024\) samples) or large (\(32\); \(4{,}096\) samples) leads to high performance variance across environments. With an appropriate \(L_{\min}\) of \(16\) (\(2{,}048\) samples), setting \(L_{\max} = 64\) (\(8{,}192\) samples) yields the best overall performance (see also Figure 5).

5.5.0.2 Policy Change Thresholds \([\delta_{\min}, \delta_{\max}]\).

Figure 6 (Center) shows that \((15\%\)–\(85\%)\) achieves the highest average performance, with \((5\%\)–\(95\%)\) as a close second. However, \((15\%\)–\(85\%)\) exhibits performance degradation on BattleZone-v5, whereas \((5\%\)–\(95\%)\) yields consistent improvements across all environments. We therefore adopt \((5\%\)–\(95\%)\) as the default setting for policy change thresholds.

5.5.0.3 Adapt Frequency \(K\).

Figure 6 (Bottom) shows that both excessively small and large values of \(K\) lead to performance degradation in certain environments. We therefore select \(K = 50\) as the default adapt frequency, as it achieves the most stable performance across all environments.

5.6 A4: Application to Continuous Control (PPO)↩︎

To demonstrate the generality of our approach beyond discrete action spaces and value-based methods, we apply ABS to PPO on continuous control tasks (MuJoCo). In continuous spaces, measuring exact action mismatch is infeasible. Therefore, we revise our metric to use the KL Divergence between the old and new policies as a proxy for non-stationarity:

\[\delta_{\pi} \approx \mathbb{E}_{s \sim \mathcal{B}} [D_\text{KL}(\pi_{\theta}(\cdot|s) || \pi_{\theta_\text{old}}(\cdot|s))]\]

Figure 7 shows the learning curves for MuJoCo locomotion environments. PPO+ABS demonstrates stable learning and improved sample efficiency compared to standard PPO. This confirms that ABS is algorithm-agnostic and applicable to continuous domains. Further implementation details and specific hyperparameter settings are provided in Appendix 7 and Appendix 9, respectively.

5.7 A4: Application to Replay Buffer-Style RL (BTR)↩︎

To demonstrate the generality of ABS beyond RL which does not use a replay buffer for policy update, we apply it to BTR [30] with minimal modification. For this replay buffer-style RL, the only architectural change is replacing the rollout range with a batch size range (i.e., \([L_{\min}, L_{\max}] \to [B_{\min}, B_{\max}]\)), leaving all other components unchanged. For detailed hyperparameter settings, please refer to Table 7.

Table [tab:btr] reports the IQM HNS with 95% confidence intervals on Atari-5 [15]. Each configuration is trained over three independent seeds and evaluated for \(100\) episodes per seed. Despite the minimal adaptation required to apply ABS to an off-policy setting, ABS improves performance in 4 out of 5 environments, with the remaining environment showing comparable scores to the baseline. These results suggest that ABS effectively generalizes across various RL algorithms—both with and without replay buffers.

max width=0.65

6 Conclusion↩︎

In this paper, we proposed a novel training framework designed to training scaled Reinforcement Learning (RL) models. We identified the inherent non-stationarity of early training as the primary bottleneck preventing the adoption of large batch sizes in RL. By analyzing this phenomenon, we demonstrated that while small batches are essential for preserving model plasticity and minimizing bias during the volatile early stages, large batches become crucial for precise convergence as the policy stabilizes.

Our proposed method, Adaptive Batch Scaling (ABS), effectively navigates this trade-off by dynamically scaling the batch size based on the Behavioral Divergence—a proxy for non-stationarity. Empirical evaluations confirm that this approach not only accelerates convergence and improves final performance but also enables the stable training of high-capacity backbone models. We hope this work serves as a stepping stone toward scalable Reinforcement Learning, and opening new avenues for training high-capacity RL agents more effectively.

6.0.0.1 Limitations & Future Work

While this work focuses on some online RL algorithms, several promising directions remain for future research. First, we plan to extend our ABS, especially replay buffer-style, to offline RL settings. Second, although our experiments were conducted on MLP and CNN-based backbones, we aim to validate our approach on Transformer-based architectures. Ultimately, we envision applying this methodology to Large Language Models, integrating it with emerging alignment techniques such as GRPO to facilitate scalable and efficient RLHF training.

Impact Statement↩︎

This work presents Adaptive Batch Scaling (ABS), which successfully unlocks large-batch training and model scaling in Reinforcement Learning (RL). This scalability paves the way for training larger, more capable RL agents for complex real-world applications such as robotics, autonomous driving, and industrial optimization. While accelerating the deployment of powerful RL systems highlights the need for continued vigilance regarding safety and alignment protocols in automated decision-making, the enhanced efficiency of ABS fundamentally contributes to more accessible and computationally sustainable AI research.

7 Implementation Details↩︎

7.1 PQN (+ABS)↩︎

To ensure the stability of the training process and prevent performance degradation due to sudden shifts in the data distribution, we incorporate three stabilizing mechanisms into the Adaptive Batch Scaling (ABS) framework.

7.1.0.1 Smoothing the Policy Divergence Metric

While the instantaneous behavioral divergence \(\delta_\pi\) provides a signal for non-stationarity, it can be susceptible to high-frequency noise inherent in stochastic environments. To mitigate this, we utilize a sliding window average of the last \(W=10\) measured divergences: \[\bar{\delta}_\pi = \frac{1}{W} \sum_{i=1}^W \delta_{\pi, i}\] During the initial ‘burn-in’ phase of training (where fewer than \(W\) samples of \(\delta_\pi\) are available), we fix the rollout length to \(L_{\min}\). This ensures that the agent has established a baseline behavioral pattern before the adaptive scaling mechanism takes effect.

7.1.0.2 Exponential Moving Average for Rollout Transitions

Even with a smoothed divergence metric, a direct application of \(L_{adpt}\) can lead to sudden spikes in the batch size, which may destabilize the gradient updates. We apply an Exponential Moving Average (EMA) to the rollout length updates: \[L_\text{curr} \leftarrow (1 - \beta) L_\text{prev} + \beta L_\text{adapt}\] where \(\beta \in (0, 1]\) is the smoothing coefficient. This ensures a gradual transition between different batch scales, preserving the numerical stability of the optimizer.

7.1.0.3 Compensation for Sample Efficiency

In on-policy RL, the number of gradient updates per sample is a critical factor for sample efficiency. Since we keep the number of mini-batches \(N_\text{mb}\) constant, increasing the rollout length \(L\) naturally leads to fewer updates per transition if the number of epochs remains fixed. To maintain a consistent ‘gradient intensity’ (updates per transition), we scale the number of update epochs \(N_\text{epochs}\) proportionally to the rollout length. Specifically, when \(L_\text{adapt}\) increases from its baseline, \(N_\text{epochs}\) is adjusted as follows: \[N_\text{epochs}^\text{scaled} = N_\text{epochs}^\text{base} \times \frac{L_\text{adapt}}{L_{\min}}\] For example, a transition from \(L=32\) with \(N_\text{epochs}=2\) to \(L=64\) would result in \(N_\text{epochs}=4\). This compensation ensures that the model continues to extract sufficient information from the collected experience even as the batch size grows.

7.2 PQN (+GNS)↩︎

We employ a dynamic batch size adjustment strategy based on the Gradient Noise Scale (GNS), denoted as \(\mathcal{B}_{\text{simple}}\), following the empirical model proposed by [13]. The GNS serves as a critical metric for determining the optimal batch size where the noise in the stochastic gradient estimates begins to dominate the useful signal. It is defined as the ratio of the trace of the gradient covariance matrix \(\Sigma\) to the squared norm of the true gradient \(G\):

\[\mathcal{B}_{\text{simple}} = \frac{\text{tr}(\Sigma)}{\|G\|^2}\]

where \(G = \mathbb{E}[\nabla \mathcal{L}(\theta)]\) represents the true gradient over the data distribution, and \(\Sigma = \text{Cov}(\nabla \mathcal{L}(\theta))\) is the covariance matrix of the per-sample gradients.

7.2.0.1 Efficient Estimation Strategy

Direct computation of \(\Sigma\) and \(G\) is computationally prohibitive. To estimate \(\mathcal{B}_{\text{simple}}\) efficiently during training, we utilize a variance-based estimator using micro-batches. At every update step (or at a fixed frequency), the current batch \(B\) is split into \(m\) micro-batches, each of size \(b = B/m\). Let \(g_i\) denote the gradient computed from the \(i\)-th micro-batch.

We first compute the mean gradient of the micro-batches, \(\bar{g} = \frac{1}{m} \sum_{i=1}^{m} g_i\). The noise term, \(\text{tr}(\Sigma)\), is estimated using the sample variance of the micro-batch gradients, scaled by the micro-batch size:

\[\text{tr}(\Sigma) \approx b \cdot \frac{1}{m-1} \sum_{i=1}^{m} \|g_i - \bar{g}\|^2\]

The signal term, \(\|G\|^2\), is estimated by correcting the bias in the squared norm of the mean gradient:

\[\|G\|^2 \approx \|\bar{g}\|^2 - \frac{1}{B}\text{tr}(\Sigma)\]

Finally, the estimated GNS is given by: \[\hat{\mathcal{B}}_{\text{simple}} = \frac{b \sum_{i=1}^{m} \|g_i - \bar{g}\|^2}{(m-1)\|\bar{g}\|^2 - \frac{b}{B}\sum_{i=1}^{m} \|g_i - \bar{g}\|^2}\]

The complete procedure is summarized in Algorithm 8.

7.3 PPO (+ABS)↩︎

While both PQN and PPO benefit from dynamic horizon adjustments, the distinct nature of their action spaces—discrete for PQN and continuous for PPO—necessitates different metrics for measuring Behavioral Divergence (\(\delta\)). Here, we contrast the Adaptive Batch Scaling (ABS) used in PQN with the Adaptive Rollout Scaling (ARS) proposed for PPO.

7.3.0.1 Metric Selection: Behavioral vs. KL Divergence

For PQN, which operates in a discrete action space with a deterministic Q-policy, we employed Behavioral Divergence as the stability metric. This measures the fraction of states in a reference batch where the greedy action changes between updates: \(\delta_{\pi} = \mathbb{E}[\mathbb{I}(\pi_{\theta}(s) \neq \pi_{\theta_{\text{old}}}(s))]\). This binary metric is computationally efficient and directly captures the instability of discrete policies.

In contrast, PPO operates in a continuous action space where the policy \(\pi_\theta(\cdot|s)\) is typically modeled as a Gaussian distribution \(\mathcal{N}(\mu_\theta(s), \sigma_\theta(s))\). A binary comparison of actions is unsuitable here. Instead, we utilize the Kullback-Leibler (KL) Divergence to measure the distributional shift between the current and the previous policy snapshot. The metric \(\delta_{\text{PPO}}\) is computed analytically as:

\[\delta_{\pi,\text{KL}} = \mathbb{E}_{s \sim \mathcal{D}} \left[ D_{KL}(\pi_{\theta}(\cdot|s) \,||\, \pi_{\theta_{\text{old}}}(\cdot|s)) \right]\]

8 Additional Results↩︎

9 Hyperparameters↩︎

Tables 3 through 6 provide the detailed hyperparameter configurations used for our experiments on the Atari benchmark. Tables 3 and 4 list the settings for the standard PQN architecture and the high-capacity Multi-Skip Residual MLP backbone, respectively. Tables 5 specifies the hyperparameters for the PQN, incorporating GNS based dynamic batch scheduling method. Additionally, Table 6 specifies the hyperparameters for the PPO baseline, incorporating our proposed adaptive rollout mechanism to evaluate its scalability.

References↩︎

---

1. https://github.com/vwxyzjn/cleanrl↩︎

2. https://github.com/daisophila/ABS↩︎