PaperlayerView exact revision on arXiv

Source-linked AI summary

On BatchNorm Forward Modes in Value-Based Reinforcement Learning

Daniel Palenicek, Mikael Henaff, Scott Fujimoto, Koustuv Sinha

arXiv:2609.06421v1cs.LGcs.AI

TL;DR

Discrete-action value learning has reported BN degradation despite lacking CrossQ’s action-input mismatch. This paper evaluates BN statistic choices at prediction, bootstrap, and action-selection forwards in C51 and PQN, finding that batch-statistic protocols reverse degradation and improve replay scaling and aggregate Atari performance.

  • Problem

    Reported BN degradation in discrete-action value learning is difficult to explain because discrete critics lack CrossQ’s action-input distribution mismatch.

  • Method

    The paper compares running- and batch-statistic forward protocols for target-based C51 and target-free PQN, including bootstrap, prediction, and action-selection passes.

  • Results

    Across C51 and PQN, batch-statistic protocols reverse BN degradation, support C51 scaling through UTD 12, and produce a higher final aggregate score than PQN with LayerNorm across 26 Atari games at 400M frames.

  • Takeaways & Limitations

    BN can substantially improve discrete-action value learning when its forward protocols are specified and configured carefully.

Abstract

from arXiv · show

Batch normalization (BN) substantially improves sample efficiency in continuous-control actor-critic methods such as CrossQ, yet recent studies report performance degradation in discrete-action value learning on Atari. These failures are surprising because discrete Q-networks lack the action-input distribution mismatch identified by CrossQ. We show for target-based C51 and target-free PQN that the simple choice between running and batch statistics at specific forward passes can reverse this degradation. In C51, switching the BN bootstrap forward to batch-statistic mode significantly improves performance over unnormalized and LayerNorm baselines and scales stably with update-to-data ratios up to 12. In PQN, using batch-statistics for both action selection and bootstrapping recovers performance from the failing running-statistic configuration. Across 26 Atari games at 400M frames, this configuration achieves a higher final aggregate score than PQN with LayerNorm. Our results show that carefully configured BN can substantially improve discrete-action value learning, and that its forward protocols are an essential part of the algorithm specification.

1 Introduction

BN failures in discrete-action value learning are surprising because discrete critics lack CrossQ’s action-input mismatch. The paper tests whether forward-pass statistic choices explain and reverse those failures across C51 and PQN.

  • Prior studies reported BN degradation in DQN, MinAtar, PQN, and Atari value-learning settings.
  • Discrete critics process states and emit action-value vectors, so CrossQ’s replay-versus-policy action mismatch cannot explain these failures.
  • Stored BN statistics can mismatch activations encountered by target, prediction, or action-selection forwards as state distributions evolve during training.
  • The study compares running and batch statistics across forward passes in target-based C51 and target-free PQN.
  • Batch-statistic bootstrapping reverses C51’s BN degradation, supports replay reuse through UTD 12, and recovers PQN performance when used for action selection and bootstrapping.
  • BN forward protocols must specify parameter sources, input batches, statistic sources, and running-state updates because they determine the function evaluated by the agent.

2 Value Learning and BatchNorm Semantics

The paper frames BN statistics as properties of a network–population pair rather than network parameters alone. Because RL populations and parameters co-evolve, each forward must explicitly define its data, statistics, and buffer-update behavior.

  • A one-step TD target combines reward, discounting, termination, and a maximum target-network action value.
  • C51 uses categorical Bellman projection with hard-copied target parameters, whereas PQN uses multi-step Q(λ) targets without replay or target parameters.
  • Batch-statistic forwards compute moments from the current input batch, while running-statistic forwards read exponential averages from earlier network–batch pairs.
  • RL changes policy visitation, replay contents, and forward populations, causing running statistics to chase a moving network–population target.
  • A complete BN forward specification includes network parameters, input population, statistic source, and whether stored statistics are updated.

3 Bootstrap Forward Modes and Replay Scaling in C51

C51’s BN outcome depends sharply on the bootstrap forward mode: running-statistic bootstrapping degrades performance, while batch-statistic bootstrapping improves it and supports scaling with replay reuse.

  • The evaluation uses 26 Atari games, three seeds per game, 5M training frames, IQM aggregation, and confidence intervals from stratified bootstrap resampling.
  • 3.1 Evaluation-mode versus training-mode bootstrapping: Evaluation-mode BN bootstrapping reduces C51 to 48% AUC and 72% final performance relative to unnormalized C51.
  • 3.1 Evaluation-mode versus training-mode bootstrapping: Training-mode BN bootstrapping reaches 154% AUC and 147% Final, outperforming the unnormalized C51 reference and a C51 LayerNorm baseline.
  • 3.1 Evaluation-mode versus training-mode bootstrapping: Increasing the responsiveness of stored moments monotonically improves performance toward the batch-statistic result, while direct training-mode bootstrapping is simpler.
  • 3.3 Scaling with replay reuse: Projected BN rises from about 1.3× the UTD-0.25 C51 reference to 11.5× at UTD 8 and remains at that level at UTD 12.
  • 3.3 Scaling with replay reuse: At UTD 8, projected BN exceeds matched unnormalized C51 on 25 of 26 games by AUC and all 26 games by final score.
  • C51’s results form a progression from severe evaluation-mode degradation, through more responsive running moments, to effective direct batch-statistic bootstrapping under increased replay reuse.

4 Forward Modes in Target-Free PQN

PQN performance depends strongly on whether acting and bootstrap forward passes use running or batch statistics. An all-batch protocol recovers BN performance and remains competitive with LayerNorm across 26 Atari games.

  • Forward-mode sensitivity: PQN’s released running/running BN configuration reproduces severe performance degradation, while forward-mode choices determine whether BN helps or harms.The diagnostic varies statistic sources for acting and bootstrapping while prediction always uses batch statistics.
  • Bootstrap and acting roles: At momentum 0.9, changing only bootstrapping to batch statistics raises Pong AUC from −13.70 to 7.63 and final return from −2.77 to 17.65.The partial recovery disappears at momenta 0.99 and 0.999.
  • Bootstrap and acting roles: Changing only acting to batch statistics approximately matches LayerNorm at momenta 0 and 0.9 and retains a positive final return of 18.40 at momentum 0.99.The reported AUCs are 15.17 and 13.94, with final returns 20.62 and 20.17 at momenta 0 and 0.9.
  • Selected protocol: Using batch statistics for both acting and bootstrapping reaches Pong AUC 15.62 and final return 20.45, with no forward pass depending on running statistics.This is the strongest Pong AUC point estimate and the protocol selected for the full-budget experiment.
  • Full-budget evaluation: Across Atari-26 at 400M frames, scalar-regression batch-statistic BN reaches final IQM 4.59 versus 4.14 for LayerNorm and significantly higher final return.The running-statistic BN configuration remains poor throughout training; matched HL-Gauss final IQMs are both approximately 4.67.
  • Interpretation: PQN’s role-specific forward semantics remain consequential even though prediction, bootstrapping, and acting reuse the same parameters.These calls still consume different input populations, motivating separate forward-mode choices.

5 Plasticity Diagnostics

Representation diagnostics show that normalization choices are associated with dormant-neuron fractions and effective rank across C51 and PQN. BN generally exhibits lower dormancy and, in several comparisons, higher effective rank than LayerNorm.

  • C51 diagnostics: Across 26 games at UTD 4, unnormalized C51 has the highest dormant fraction and lowest effective rank.The diagnostics use three seeds per game and matched game–seed identities.
  • C51 diagnostics: Both C51 BN variants have lower dormant-fraction IQMs than LayerNorm and higher effective-rank IQMs after the first diagnostic.The WN variants remain close to their unprojected counterparts relative to the difference between normalizers.
  • PQN diagnostics: HL-Gauss BN has a higher effective-rank IQM than its LayerNorm counterpart, whereas scalar BN and scalar LayerNorm have more similar effective ranks.The shared ordering is described as a qualitative regularity rather than a mechanism.

6 Conclusion

The paper finds that BN is not uniformly harmful in discrete-action value learning: its outcome depends on the forward mode and its application. Properly configured BN reverses reported degradation in C51 and PQN and can outperform LayerNorm on final aggregate performance.

  • Conclusion: Carefully configured BN substantially improves discrete-action value learning and is not uniformly harmful in value learning.The conclusion challenges the prevailing view that BN is ill-suited to discrete-action value learning.
  • Conclusion: In C51, replacing stored target moments with current-minibatch statistics reverses BN degradation and supports stable update-to-data ratio scaling up to 12.The conclusion reports improvement over the failing stored-statistics configuration.
  • Conclusion: In PQN, changing acting and bootstrap protocols recovers BN performance without replay or a separate target network.The selected batch-statistic configuration is competitive over training and achieves a higher final aggregate score than LayerNorm.
  • Conclusion: Forward protocols are an essential part of the algorithm specification because each prediction, bootstrap, acting, and evaluation forward can use different statistics and input batches.These choices determine the function evaluated by the agent and can significantly impact performance.

A Implementation Details

The implementation section defines the evaluation, preprocessing, architectures, optimizers, and training schedules used alongside role-specific BatchNorm semantics.

  • Implementation details: Evaluation, Atari preprocessing, network architectures, optimizers, and training schedules complement the role-specific BatchNorm semantics.The section specifies these implementation components without reporting an additional result.

A.1 Evaluation and uncertainty

The evaluation aggregates human-normalized returns with IQM and reports uncertainty using seed-stratified bootstrap intervals. The experiments cover a fixed 26-game Atari suite and specify shared C51 and PQN configurations.

  • Evaluation metrics: Human-normalized scores are aggregated with IQM, while C51 AUC averages the evaluation grid and the final statistic uses the 5M-frame evaluation.Intervals use 5,000 seed-stratified resamples, with matched-seed resampling for target-mode contrasts.
  • C51 evaluation design: The C51 replay-ratio sweep evaluates UTD 0.25, 1, 2, 4, 8, and 12 across 26 games and three seeds.Each panel uses its own vertical scale, and BN arms use batch statistics at the bootstrap forward.
  • PQN evaluation design: PQN evaluation uses 30 frame bins for the Pong diagnostic and 60 complete bins for each seed in the 400M-frame Atari-26 experiment.AUC averages human-normalized bin values, while the endpoint uses the final bin.
  • Experimental scope: Both studies use the same 26 Atari games and human and random reference scores from Mnih et al. (2015).The listed configurations report as-run values, including launcher-overridden defaults where applicable.

B Additional C51 Results

Additional C51 results provide per-environment learning curves across the replay-ratio sweep, using raw returns averaged over three seeds with one-standard-error bands. The figures distinguish UTD conditions and retain the BN, LayerNorm, and fixed-norm projection conventions.

  • UTD 0.25: At UTD 0.25, per-environment C51 curves show raw evaluation returns averaged over three seeds with one-standard-error bands.BN uses batch statistics at the bootstrap forward; WN denotes the fixed-norm projection.
  • UTD 1 and 2: At UTD 1 and 2, the top and bottom grids report per-environment C51 learning curves, respectively.The curves are raw evaluation returns averaged over three seeds with one-standard-error bands.
  • UTD 4 and 8: At UTD 4 and 8, the per-environment grids continue the same C51 raw-return evaluation protocol.The supplied figure text identifies the two replay ratios but does not report a performance outcome.
  • UTD 12: At UTD 12, the per-environment C51 curves use the same raw-return and one-standard-error conventions as the earlier grids.The figure follows the conventions of the UTD 0.25 presentation.

C Additional PQN Results

Additional PQN results present per-game learning curves over 400M Atari frames. The curves compare scalar-regression and HL-Gauss objectives with running-statistic and batch-statistic BN configurations alongside LayerNorm.

  • Per-game curves: PQN per-game curves span 400M Atari frames and report mean raw return with one-standard-error bands across ten seeds.The released running-statistic BN configuration is shown in gray, while batch-statistic BN configurations are shown in blue and purple.
  • Configuration encoding: Colors distinguish LayerNorm and BN configurations, while circles and triangles distinguish scalar-regression and HL-Gauss objectives.The figure compares the corresponding objective and normalization choices within each game.

D Discrete-Action Actor-Critic Instantiation on Atari

The paper also evaluates a discrete-action actor–critic instantiation of the BN recipe on Atari using MR.Q’s Gumbel-Softmax trick. Across 26 games, XQC exceeds TD3 at 5M frames but remains below MR.Q on the aggregate IQM.

  • Experimental setup: The actor–critic extension uses XQC, a CrossQ-style batch-statistic BN critic with fixed-norm weight projection, and compares it with MR.Q and TD3.The experiment applies MR.Q’s Gumbel-Softmax trick to discrete-action Atari.
  • Aggregate IQM: At 5M frames, final IQMs are 1.30 for MR.Q, 1.02 for XQC, and 0.55 for TD3 across 26 games.The reported 95% intervals are [1.24, 1.37], [0.93, 1.13], and [0.50, 0.60], respectively.
  • Aggregate comparison: XQC reaches roughly human-level aggregate performance and clearly exceeds TD3, while MR.Q remains ahead.Figure 11 additionally reports median, IQM, mean, and optimality-gap comparisons with 95% stratified bootstrap intervals.
  • Per-environment results: Per-environment actor–critic curves report raw evaluation returns every 2 × 10^5 frames through 5 × 10^6 frames, averaged over five seeds.The displayed methods are XQC, MR.Q, and TD3.
Loading 2609.06421v1…