Stable Gradients for Stable Learning at Scale in Deep Reinforcement Learning

TL;DR

The paper investigates why scaling deep reinforcement learning models often degrades performance, attributing the issue to the interplay of non-stationarity and gradient pathologies that worsen with larger architectures. It proposes two direct interventions—multi-skip residual connections and a Kron (Kronecker-factored) second-order optimizer—to stabilize gradient flow and enable robust learning at scale. Across a broad set of experiments, including Atari-10, the full ALE suite, PPO, scaled encoders, Simba with Kron, and scalable offline Q-learning, the interventions consistently improve stability and performance as models grow. The findings underscore that preserving gradient information is a crucial prerequisite for effective parameter scaling in deep RL, offering a practical pathway to scalable, robust agents in diverse environments.

Abstract

Scaling deep reinforcement learning networks is challenging and often results in degraded performance, yet the root causes of this failure mode remain poorly understood. Several recent works have proposed mechanisms to address this, but they are often complex and fail to highlight the causes underlying this difficulty. In this work, we conduct a series of empirical analyses which suggest that the combination of non-stationarity with gradient pathologies, due to suboptimal architectural choices, underlie the challenges of scale. We propose a series of direct interventions that stabilize gradient flow, enabling robust performance across a range of network depths and widths. Our interventions are simple to implement and compatible with well-established algorithms, and result in an effective mechanism that enables strong performance even at large scales. We validate our findings on a variety of agents and suites of environments.

Figures (24)

• Figure 1: Training dynamics under stationary and non-stationary supervised learning. (Left) In the stationary setting, both shallow and deep models fit the data effectively across widths. Under non-stationarity only shallow networks partially recover during training, while deeper ones collapse. (Right) This collapse correlates with degraded gradient flow. In stationary settings, gradient norms remains stable across all network scales (shaded boxes) while in non-stationary settings (solid-colored boxes), gradient magnitudes diminish with depth and width, suggesting poor adaptability.
• Figure 2: Mean episode returns and gradient norms across increasing MLP depths and widths on two ALE games using PQN. (Left) Only shallow networks achieve high episode returns; performance collapses for deeper networks. (Right) The collapse correlates with vanishing gradient norms, suggesting that deeper models fail to adapt to non-stationarity in deep RL.
• Figure 3: Training pathologies emerge as MLP depth increases. Deeper networks exhibit a higher fraction of inactive neurons, reduced representation rank (SRank), vanishing Hessian trace (loss curvature), and degraded learning performance (mean Q-values and episode returns). These trends indicate that scaling depth limits expressivity and plasticity, impairing policy quality.
• Figure 4: (Left) MLP architectures and (right) scaling strategies studied.
• Figure 5: Gradient-stabilizing interventions improve scalability in deep RL. (Left) Standard fully connected networks trained with PQN collapse at greater depths due to vanishing gradients. In contrast, multi-skip architectures maintain gradient flow and scale effectively. (Right) The default RAdam optimizer leads to instability in deep networks, while switching to the Kron optimizer preserves gradient signal and enables stable learning without architectural changes.