Batch-Wise Reward Normalization

• Batch-wise reward normalization is an adaptive method that scales mini-batch rewards to stabilize gradient updates and preserve the true learning objective in reinforcement learning.
• It dynamically adjusts reward statistics for each batch, unlike fixed clipping, ensuring that diverse reward magnitudes from heterogeneous environments are accurately maintained.
• Empirical results using techniques like Pop-Art and BNPO show reduced gradient variance and improved performance, with benchmarks demonstrating tighter gradient ranges and more stable updates.

Batch-wise reward normalization refers to adaptive normalization schemes that operate over each mini-batch of rewards (or value targets) encountered during reinforcement learning (RL), such that surrogate learning targets have stabilized statistics and better-conditioned gradients throughout training. This approach addresses the instability and lack of invariance to scale that commonly afflict RL algorithms utilizing function approximation. Unlike fixed clipping—where targets are forcibly truncated to a predetermined interval—batch-wise normalization learns and maintains dynamic transformations to ensure numerically stable updates without distorting the true objective.

1. Motivation and Limitations of Fixed Clipping

Early value-based deep RL algorithms, such as DQN, often encountered severe instability due to large variations in the scale of temporal-difference (TD) targets: $$y_t = r_{t+1} + \gamma \max_{a} Q(s_{t+1},a;\theta^{-})$$ When training across heterogeneous environments (e.g., different Atari games), some tasks yield reward values of $\pm1$ and others reach magnitudes exceeding $\pm 1000$. Employing a global fixed learning rate leads to instability (huge gradients from large rewards), while tiny rewards result in slow convergence.

The standard workaround—reward clipping to $[-1,1]$—equalizes gradient scale. However, it fundamentally alters the learning objective from sum-of-true-rewards to sum-of-clipped-rewards. This can induce substantially different learned behaviors, causes loss of reward signal magnitude information, and is founded on a domain-specific heuristic that may not generalize to non-standard reward structures (Hasselt et al., 2016).

Batch-wise normalization circumvents these issues by learning batch or streaming statistics and transforming targets so that, per batch, they maintain controlled mean and variance. This normalizes gradient magnitudes while retaining the true underlying objective.

2. Batch-wise Normalization: Algorithms and Update Rules

A hallmark example of batch-wise normalization is Pop-Art, introduced by van Hasselt et al. (Hasselt et al., 2016). Consider a mini-batch of $B$ target values $y_{t,i}$ at update step $t$:

• Maintain running estimates:
  • $\mu_{t-1}$: previous batch-mean,
  • $\nu_{t-1}$: previous batch-second-moment,
  • $\beta$: smoothing constant (e.g., $10^{-4}$).

Compute new batch means: $$\bar y_t = \frac{1}{B} \sum_{i=1}^B y_{t,i}, \quad \overline{y^2}_t = \frac{1}{B} \sum_{i=1}^B y_{t,i}^2$$ Update moments: $$\mu_t = (1-\beta)\mu_{t-1} + \beta\,\bar y_t,\qquad \nu_t = (1-\beta)\nu_{t-1} + \beta\,\overline{y^2}_t$$ Estimated variance: $\sigma_t^2 = \nu_t - \mu_t^2$ For each target, normalization (to zero mean, unit variance) gives: $\hat y_{t,i} = \frac{y_{t,i} - \mu_t}{\sigma_t}$ The network is trained to predict these $\hat y_{t,i}$, but the original targets can always be recovered by inverting the transformation: $Q(s) = \sigma_t \hat q + \mu_t$

For normalization stability, Theorem 2 in (Hasselt et al., 2016) demonstrates that each normalized error is always contained within $$\left[-\sqrt{\frac{1-\beta}{\beta}}, \sqrt{\frac{1-\beta}{\beta}}\right]$$.

3. Architectural Adjustments to Preserve Unnormalized Outputs

Because normalization statistics ($\mu_t,\sigma_t$) change over time, naive batch-wise normalization would regularly invalidate the semantics of the model’s original outputs. To maintain the equivalence between the normalized and unnormalized outputs, Pop-Art introduces adaptive output-layer scaling ("POP" update):

Let $g(s)$ be the normalized network (last layer: $\hat q = Wh(s) + b$), and

$$f(s) = \sigma (W h(s) + b) + \mu = (\sigma W)h(s) + (\sigma b + \mu)$$

When updating to new statistics $(\mu',\sigma')$, find new parameters $(W', b')$ such that $f(s)$ is unchanged for all $h(s)$: $$W' = \frac{\sigma}{\sigma'} W,\qquad b' = \frac{\sigma b + \mu - \mu'}{\sigma'}$$ This reparametrization preserves the effect of all prior learning under the new normalization, ensuring that $Q(s)$ is not broken by updates to $\mu, \sigma$.

4. BNPO: Beta-based Batch-wise Reward Normalization

Recent methods such as Beta Normalization Policy Optimization (BNPO) generalize batch-wise normalization to policy-gradient RL with explicit variance minimization (Xiao et al., 3 Jun 2025). In this regime, for each batch of $N$ queries (e.g., natural language prompts) and $m$ samples per query, the observed binary rewards $R_{ij}$ are modeled via their empirical success probabilities $p_i$.

BNPO fits a $\mathrm{Beta}(a,b)$ distribution to batch-wise $p_i$ via method of moments, then computes optimal normalization parameters $(\alpha,\beta)$: $$\alpha = 1 + \frac{a}{3},\qquad \beta = 1 + \frac{b}{3}$$ Each raw reward is then normalized as: $$\tilde r_{ij} = \frac{R_{ij} - p_i}{f_N(p_i;\alpha,\beta)}$$ where $f_N(p_i;\alpha,\beta)$ is the normalization Beta PDF.

BNPO minimizes the variance of the policy gradient estimator under batch-dependent normalization, and empirically yields lower-variance, more stable updates than REINFORCE or static normalization schemes. Special cases (REINFORCE+baseline, GRPO) are recovered for different fixed $(\alpha,\beta)$ settings (Xiao et al., 3 Jun 2025).

5. Empirical Comparisons and Benefits

Pop-Art normalization for value-based RL and BNPO for policy-gradient RL produce consistent improvements in gradient stability and downstream performance.

In Atari-57, Pop-Art (with no clipping) produced median gradient norms across games spanning only two orders of magnitude, compared to six for unclipped and nearly four for clipped DQN. This enabled setting a global learning rate and, in score, Double DQN with Pop-Art outperformed or matched clipped Double DQN on 32 of 57 Atari games, with a mean improvement of +34% and median +0.4%, while eliminating reward-clipping-induced policy distortions (Hasselt et al., 2016).

For BNPO, average pass@1 on mathematical reasoning tasks was higher than with REINFORCE or GRPO at both small and large model scales, and training exhibited markedly reduced gradient-norm fluctuations (Xiao et al., 3 Jun 2025).

Method	Gradient Norm Range	Retains True Objective	Empirical Stability	Reference
Clipping	4–6 orders	No	Adequate	(Hasselt et al., 2016)
Pop-Art DQN	2 orders	Yes	Superior	(Hasselt et al., 2016)
BNPO (PG)	Reduced variance	Yes	Superior	(Xiao et al., 3 Jun 2025)

6. Batch-wise Normalization in Contemporary Policy Optimization

BNPO extends batch-wise normalization to multi-component binary rewards via per-component normalization and advantage decomposition, where for $K$ components, the final advantage is given by

$$A(q,o) = \frac{1}{K}\sum_{k=1}^{K} \frac{R^{(k)}(q,o) - p^{(k)}(q)}{f_N(p^{(k)}(q);\,\alpha^{(k)},\beta^{(k)})}$$

This decouples normalization across components and generalizes to settings beyond scalar rewards.

BNPO’s batch-wise normalization dynamically adjusts to changing policy distributions, aligning reward normalization with the on-policy statistics at every optimization step. The method directly reduces gradient variance, which is crucial for stable and efficient policy optimization when training with complex or sparse reward landscapes (Xiao et al., 3 Jun 2025).

7. Conclusion and Significance

Batch-wise reward normalization, as instantiated by Pop-Art for value-based RL and BNPO for policy-gradient settings, enables scale-invariant, distortion-free, and robust learning across diverse environments and reward regimes. Unlike static normalization or clipping, batch-wise normalization explicitly adapts to the empirical reward or value distribution of each batch, maintaining well-conditioned surrogate targets and stable gradient dynamics. Empirical results demonstrate substantial improvements in both performance and training stability, and the methodology generalizes to multi-component objectives and nonstationary reward structures (Hasselt et al., 2016, Xiao et al., 3 Jun 2025).