Batch-wise Reward Normalization

• Batch-wise reward normalization is a technique that adaptively rescales rewards per training batch to mitigate gradient instability in reinforcement learning.
• Methods like Pop-Art and BNPO dynamically adjust normalization parameters, replacing static reward clipping with principled, data-driven scaling.
• Empirical results demonstrate that these approaches maintain consistent gradient norms and enhance sample efficiency across diverse tasks.

Batch-wise reward normalization refers to techniques that, on each training batch, adaptively rescale rewards or learning targets during reinforcement learning (RL) or policy optimization. This normalization mitigates instability caused by dynamic or disparate reward scales, stabilizes the variance of gradient updates, and improves the learning signal delivered to functional approximators. Key implementations include Pop-Art normalization for value-based deep RL and the Beta Normalization Policy Optimization (BNPO) method for policy-gradient methods with binary-valued rewards. Both enable training across tasks or environments of varying reward magnitudes without resorting to arbitrary clipping or static normalization.

1. Motivation and Background

Reinforcement learning algorithms, including classic value-function estimation and modern policy-gradient methods, are generally sensitive to the scale and distribution of rewards. Without normalization, environments with high-magnitude or highly variable rewards can cause gradient estimates to become unstable or excessively high-variance. Historically, this led to domain-specific heuristics such as reward clipping, particularly for value-based deep RL agents trained on Atari games, where direct regression to true rewards produced brittle learning and poor cross-game generalization (Hasselt et al., 2016).

Policy-gradient methods using binary-valued, rule-based rewards (e.g., for reasoning-focused LLMs) are also affected by scale sensitivity. Static normalization can become suboptimal as the policy’s induced reward distribution evolves throughout training (Xiao et al., 3 Jun 2025). The need for adaptive, batch-wise normalization is especially acute in large-scale and heterogeneously-scored RL settings.

2. Pop-Art: Adaptive Target Normalization for Value-Based RL

The Pop-Art (Preserving Outputs Precisely while Adaptively Rescaling Targets) methodology maintains online estimates of the mean $\mu_t$ and variance $\sigma_t^2$ of target values. It normalizes each batch’s targets via an affine transformation:

$\hat{Y}_t = \frac{Y_t - \mu_t}{\sigma_t}$

where $Y_t$ are the raw targets (e.g., TD targets in value-based RL). The mean and variance are updated either per sample or per batch using exponential moving averages, which track the evolving statistics under nonstationary policies:

• $\mu_t = (1-\beta_t) \mu_{t-1} + \beta_t \bar{Y}_t$
• $$\nu_t = (1-\beta_t) \nu_{t-1} + \beta_t \overline{Y^2_t}$$
• $\sigma_t = \sqrt{\nu_t - \mu_t^2}$

with $\beta_t$ either $1/t$ or a small constant to retain adaptivity in nonstationary settings. Batch-wise normalization aggregates statistics over the minibatch at each update step (Hasselt et al., 2016).

A unique property of Pop-Art is the preservation (“POP step”) of the unnormalized function output: when updating $\mu_t$ and $\sigma_t$, the last linear network layer parameters are rescaled so the true (unnormalized) network predictions remain unchanged. This ensures continuity of the learned value function despite rescaling.

Pop-Art replaces target and reward clipping, letting the agent optimize for true reward structure and variance across tasks.

3. BNPO: Beta-Based Batch-Wise Normalization in Policy Optimization

BNPO extends the batch-wise normalization principle to policy-gradient methods designed for binary (or multi-binary) reward settings (Xiao et al., 3 Jun 2025). It fits, on each batch, a Beta distribution to success probabilities $p(q)$ of the current policy for each context $q$:

• For a batch of $n$ contexts, $p(q_i) \approx \text{Mean}_{j=1}^m R(q_i, o_{ij})$, where each output $o_{ij}$ is sampled from $\pi_\theta(\cdot | q_i)$.

The batch mean and variance of $p(q)$ are estimated, and method-of-moments yields parameters $(a, b)$ for the data Beta:

• $$a = \left( \frac{E[p](1-E[p])}{\text{Var}[p]} - 1 \right) E[p]$$
• $$b = \left( \frac{E[p](1-E[p])}{\text{Var}[p]} - 1 \right) (1-E[p])$$

The optimal normalization Beta, shown theoretically to minimize the variance of the policy-gradient estimate, sets scaling parameters as $\alpha = 1 + a/3$, $\beta = 1 + b/3$.

Each reward is then normalized by dividing the mean-subtracted reward by the probability density of its $p(q)$ under the Beta fit:

$$A_{\alpha, \beta}(q,o) = \frac{R(q,o) - p(q)}{f_N(p(q); \alpha, \beta)}$$

where $f_N$ is the Beta density. This yields batch-adaptive scaling of reward deviations, dynamically tuned to policy-induced reward statistics at each update. BNPO further generalizes to multiple binary reward channels via per-channel normalization and summation.

4. Statistical Properties and Theoretical Analysis

Both Pop-Art and BNPO provide explicit variance reduction for gradient updates by adapting normalization scales to batch or running statistics. In Pop-Art, Theorem 2 bounds normalized targets within an explicit interval depending only on $\beta$ ($$\left[ -\sqrt{(1-\beta)/\beta}, +\sqrt{(1-\beta)/\beta} \right]$$), preventing outlier-induced divergence (Hasselt et al., 2016).

For BNPO, Theorem 1 proves that the batch-wise choice of $(\alpha, \beta)$ minimizes the variance of the normalized gradient estimator, contingent on mild regularity conditions for finiteness (i.e., $\alpha < (a+3)/2$, $\beta < (b+3)/2$). The variance formula is given in closed-form in terms of Beta functions, and the minimum is unique. The method is provably lower-variance than vanilla REINFORCE or static-normalized baselines.

Both frameworks address the problem of dynamic reward distributions during training, ensuring that learning rates and gradient magnitudes remain well-behaved across different tasks and as the policy evolves.

5. Practical Algorithms and Implementation Details

For Pop-Art, the standard workflow for each batch includes:

1. Compute raw targets (e.g., TD targets in Q-learning).
2. Update running or batch mean and variance estimates.
3. Apply the POP step: rescale the last layer’s parameters to preserve unnormalized output.
4. Forward pass: obtain normalized Q-value prediction.
5. Compute normalized error and backpropagate only this error through the network (Hasselt et al., 2016).

For BNPO, the per-batch procedure is:

1. Sample contexts and outputs; compute binary rewards.
2. Estimate per-context reward means; fit Beta distribution $(a,b)$ to $p(q)$.
3. Derive normalization parameters $(\alpha,\beta)$.
4. Compute Beta-normalized advantage for each sample.
5. Use these advantages in PPO-style policy-gradient updates for stable, variance-reduced optimization (Xiao et al., 3 Jun 2025).

For multiple binary reward channels, BNPO computes a separate normalization for each and aggregates the normalized advantages.

6. Impact on Training Stability and Empirical Results

Pop-Art eliminates the need for arbitrary reward or target clipping, enabling value-based deep RL agents to faithfully represent and optimize the true reward structure in diverse domains. Empirical results in Atari benchmarks show that Pop-Art maintains gradient norms within tight, game-invariant ranges, supporting stable training with a single learning rate across environments. Scores using Pop-Art normalization match or surpass those of clipped Double DQN agents; importantly, Pop-Art permits learning of nuanced policy behaviors sensitive to true reward magnitude (e.g., differentiating pursuit of pellets versus ghosts in Ms. Pac-Man) (Hasselt et al., 2016).

BNPO demonstrates consistent gradient-norm stability and superior sample efficiency in reasoning-focused LLMs. On mathematical reasoning tasks such as MATH500 and AMC23, BNPO achieves state-of-the-art pass@1 accuracy, outperforming or matching all static-normalization and baseline methods across multiple model sizes. Figure 1 of (Xiao et al., 3 Jun 2025) shows that gradient norms under BNPO remain most stable, and Figure 2 confirms that normalization parameters track empirical statistics throughout learning.

7. Advancements, Limitations, and Broader Implications

Batch-wise reward normalization provides a principled alternative to reward clipping and static normalization. It enables RL systems to operate robustly over highly non-stationary and task-heterogeneous reward landscapes, which is essential for generalizable agents and large-scale RL for LLMs. Adaptive normalization, as exemplified by Pop-Art and BNPO, increases the stability and reproducibility of policy optimization while preserving true reward semantics (Hasselt et al., 2016, Xiao et al., 3 Jun 2025).

A plausible implication is that as RL and LLM fine-tuning regimes become increasingly large-scale and diverse, batch-wise adaptive normalization will remain a core component. Current methods are most effective when batch statistics are well estimated; however, small-batch or highly non-iid settings may present new challenges. The extension to multi-variate and continuous reward scenarios is active research territory. Empirical evidence suggests that per-component normalization and decomposition yield additive benefits in multi-reward settings (Xiao et al., 3 Jun 2025).