On BatchNorm Forward Modes in Value-Based Reinforcement Learning

Daniel Palenicek; Mikael Henaff; Scott Fujimoto; Koustuv Sinha

cs.LGSep 6, 2026

On BatchNorm Forward Modes in Value-Based Reinforcement Learning

Authors: Daniel Palenicek, Mikael Henaff, Scott Fujimoto, Koustuv Sinha

Organizations: FAIR at Meta · Technical University of Darmstadt

Abstract

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.

Figures & tables

Figure 1: Bootstrap forward modes and replay scaling in C51. All panels report performance as a percentage of unnormalized C51 at UTD

0.25

, with C51 fixed to 100%. Panels (a,b) report AUC and final IQM at UTD

0.25

. BN running uses evaluation-mode bootstrapping and reads stored target moments. BN batch uses training-mode bootstrapping and statistics from the current next-state minibatch. The three BN running (

\alpha

) conditions update stored target moments from that minibatch with PyTorch update weight

\alpha\in\{0.001,0.01,0.1\}

before constructing the target in a separate evaluation-mode forward. Panels (c,d) report AUC and final IQM across UTD for C51, LayerNorm (LN), BN, and their fixed-norm projection variants (WN). C51, LN, and BN use a constant learning rate, whereas LN WN and BN WN use linear decay. Every displayed bar or arm–UTD point contains 26 games and three seeds per game, each trained for 5M frames. Bars and curves show IQM; error bars and bands are marginal 95% stratified bootstrap confidence intervals across 26 Atari games and 3 seeds each.

Figure 2: Forward-mode sensitivity in target-free PQN. Panels (a,b) report AUC and mean return over the last 4M frames of a 40M-frame Pong ablation. Gray uses running statistics for acting and bootstrapping, orange changes only bootstrapping to batch statistics, purple changes only acting to batch statistics, and blue uses batch statistics for both. The horizontal green line is LayerNorm; the gray star marks the released running/running configuration at Flax momentum

0.99

. Lines and points show means, and shaded regions show standard errors across five seeds each. Panel (c) full IQM training curves over 400M frames on Atari-26 aggregated across ten seed. Intervals are 95% stratified bootstrap confidence intervals within-game seed-bootstrap intervals. The dashed gray curve is the released running-statistic BN configuration from Gallici et al. (2025) ; green and blue show scalar-regression LayerNorm and batch-statistic BN, and orange and purple show their respective HL-Gauss variants.

Figure 3: Representation diagnostics in hard-target C51 and target-free PQN. For each of the five C51 arms, panels (a,b) use all 78 matched game–seed identities at UTD 4. C51, LayerNorm, and BN use a constant learning rate and the two WN-projection arms (LN WN, BN WN) use a linear-decay learning rate, matching Figure 1 . Panels (c,d) average PQN’s ten per-game seeds across 26 games. Both algorithms compute effective rank from centered post-ReLU 512-dimensional penultimate features. Confidence intervals are pointwise 95% stratified bootstrap confidence intervals obtained by resampling seed identities within each fixed game and recomputing the IQM. Color identifies the normalizer; in the C51 panels line style denotes the fixed-norm projection (WN), and in the PQN panels line style compares scalar regression and HL-Gauss for both normalizers.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 4: Aggregate C51 learning curves across replay ratios. Human-normalized IQM over 26 games and three seeds at UTD

0.25

1

2

4

8

, and

12

. Each panel uses its own vertical scale. All BN arms use batch statistics at the bootstrap forward; “WN” denotes the fixed-norm projection. C51, LayerNorm, and BN use a constant learning rate, while the two fixed-norm-projection arms (LN WN, BN WN) use a linear-decay learning rate, matching Figure 1 . Bands are marginal 95% seed-stratified bootstrap intervals conditional on the fixed game set. These curves are the longitudinal counterparts of Figure 1 c,d.

Architecture
Hyperparameter	Value
Convolutional encoder	Conv $(32,8{\times}8,\mathrm{s}4)$ –Conv $(64,4{\times}4,\mathrm{s}2)$ –Conv $(64,3{\times}3,\mathrm{s}1)$
Dense (penultimate) layer	512 units
Output head	per-action, 51 atoms
Hidden activation	ReLU
Initialization	Kaiming-normal (ReLU gain); biases $0$

Appendix

Table 1: C51 shared training hyperparameters (Atari 26-game suite).

Architecture
Hyperparameter	Value
Convolutional encoder	Conv $(32,8,4)$ –Conv $(64,4,2)$ –Conv $(64,3,1)$
Dense (penultimate) layer	512 units
Normalization	every hidden layer (LayerNorm or BatchNorm); head unnormalized
Output head	scalar (regression) or 51-atom (HL-Gauss)
Optimizer

Appendix

Table 2: PQN hyperparameters (target-free, on-policy;

40

M Pong diagnostic and

400

M Atari-26 experiment).

Figure 5: Per-environment C51 learning curves at UTD

0.25

. Curves are raw evaluation returns, averaged over three seeds; bands are one standard error. All BN arms use batch statistics at the bootstrap forward. WN denotes the fixed-norm projection. C51, LayerNorm, and BN use a constant learning rate, while LN+WN and BN+WN use a linear-decay learning rate, matching Figure 1 .

Figure 6: Per-environment C51 learning curves at UTD

1

and

2

. The top grid reports UTD

1

and the bottom grid UTD

2

. Curves are raw evaluation returns, averaged over three seeds; bands are one standard error. Conventions follow Figure 5 .

Figure 7: Per-environment C51 learning curves at UTD

4

and

8

. The top grid reports UTD

4

and the bottom grid UTD

8

. Curves are raw evaluation returns, averaged over three seeds; bands are one standard error. Conventions follow Figure 5 .

Figure 8: Per-environment C51 learning curves at UTD

12

. Curves are raw evaluation returns, averaged over three seeds; bands are one standard error. Conventions follow Figure 5 .

Figure 9: Per-environment coupled greedy-test learning curves over 400M Atari frames. For each game, curves show the mean raw return and one standard error across all ten seeds. The gray curve (crosses) shows the released running-statistic BN configuration from Gallici et al. (2025) . Green and orange identify scalar-regression and HL-Gauss LayerNorm; blue and purple show the corresponding batch-statistic BN configurations. Circles and triangles distinguish the scalar and HL-Gauss objectives.

Figure 10: Actor–critic Atari aggregate. Human-normalized IQM over 26 games (five seeds) for XQC, MR.Q, and TD3 to 5M frames (

5\times 10^{6}

frames at action repeat 4). Bands are 95% seed-bootstrap intervals; the dotted line marks human-level performance.

Figure 11: Aggregate metrics (actor–critic Atari). Median, IQM, mean, and optimality gap of the human-normalized score at 5M frames over 26 games (five seeds), with 95% stratified bootstrap intervals ( Agarwal et al., 2021 ) . Higher is better except for the optimality gap.

Figure 12: Per-environment actor–critic Atari curves. Raw evaluation return for XQC, MR.Q, and TD3 on each of the 26 games, averaged over five seeds; bands are one standard error. Evaluation is every

2\times 10^{5}

frames to

5\times 10^{6}

frames.

Reinforcement Learning

Reinforcement learning studies how agents learn sequential decisions from rewards and interaction, balancing exploration, performance, safety, and adaptation.

See all 1,742 papers

Accelerating Q-learning through Efficient Value-Sharing across Actions
Jun 29, 2026Prabhat Nagarajan, Brett Daley, Martha White +1Deep Q-Learning Value Function Estimation
Department of Computing Science, University of Alberta, Edmonton, AB, Canada. · Alberta Machine Intelligence Institute · Meta (independent work, not on Meta’s behalf). +1
RN-D: Discretized Categorical Actors for On-Policy Reinforcement Learning
Jan 30, 2026Yuexin Bian, Jie Feng, Tao Wang +3Policy Optimization Continuous Control
University of California San Diego, La Jolla, USA
Decoupled Continuous-Time Reinforcement Learning via Hamiltonian Flow
Feb 16, 2026Minh NguyenReinforcement Learning Continuous-Time RL
Google · University of Texas at Austin

Jun 29, 2026J/K move · Enter open · S save

Accelerating Q-learning through Efficient Value-Sharing across Actions

Prabhat Nagarajan, Brett Daley, Martha White, Marlos C. Machado

Department of Computing Science, University of Alberta, Edmonton, AB, Canada. · Alberta Machine Intelligence Institute · Meta (independent work, not on Meta’s behalf).+1

Deep Q-Learning Value Function Estimation Q-Learning

Open paper AskPDF arXiv

Action values are foundational to many control algorithms such as Q-learning. Therefore, efficient action-value learning is central to reinforcement learning (RL). However, learning them can be slow, requiring many updates to move values from their initialization, typically near zero, to their true values, which may be far from zero. Moreover, action-value learning algorithms typically update each state-action pair independently, without learning a value that is common to all actions within a state. In this paper, we address these inefficiencies by introducing the mean-expansion layer, which accelerates action-value learning by sharing values across actions within a state and by changing the problem from directly learning potentially large action-values to learning a lower-norm representation of them. In deep RL, this layer can be applied as a parameter-free addition to Q-network architectures without altering the underlying algorithm. Applied to deep Q-networks and implicit quantile networks, it improves aggregate performance across 57 Atari 2600 games while increasing action gaps and dramatically reducing value overestimation.

Abstract

Reinforcement Learning

Explore similar work

Accelerating Q-learning through Efficient Value-Sharing across Actions

RN-D: Discretized Categorical Actors for On-Policy Reinforcement Learning

Decoupled Continuous-Time Reinforcement Learning via Hamiltonian Flow