Technical Summary: Calibrated Partial Resets (CPR) for Continual Reinforcement Learning

1. Problem Statement

Deep neural networks often suffer from a loss of plasticity during prolonged non-stationary training, a critical issue in Continual Reinforcement Learning (RL). In these settings, policy updates alter the agent's future data distribution, creating feedback loops where small optimization instabilities can compound into catastrophic policy degradation (Juliani & Ash, 2024).

Existing remedies for plasticity loss, such as Continual Backpropagation (CBP), ReDo, and ReGraMa, rely on binary neuron resets. These methods identify low-utility neurons (based on activation or gradient statistics) and fully reinitialize them. While this restores gradient flow, the authors argue that binary resets conflate two distinct design choices: which units to refresh and how strongly to change them.

• Instability: The "all-or-nothing" nature of binary resets introduces abrupt parameter changes and gradient spikes, which destabilize long-horizon RL training and can lead to policy collapse.
• Brittleness: Full reinitialization zeros outgoing weights, causing distribution shifts in subsequent layers, particularly problematic without layer normalization.
• Uniform Decay Limitations: Alternative methods like Shrink & Perturb apply uniform decay to all units, which is smooth but indiscriminate, failing to target the specific units that need refreshing.

The core problem is the lack of a method that is both selective (targeting low-utility units) and smooth (avoiding abrupt parameter jumps) to prevent policy collapse while maintaining peak performance.

2. Methodology: Calibrated Partial Resets (CPR)

The authors propose Calibrated Partial Resets (CPR), an optimizer that periodically pulls neurons toward their initialization with a strength scaled by each neuron's utility. Unlike binary resets, CPR applies a continuous, partial adjustment.

2.1 Utility Estimation

CPR utilizes gradient-based utility scores, which recent work suggests scale more favorably (Liu et al., 2025).

• Raw Score ($S_i^\ell$): The mean magnitude of the gradient of the incoming weights for unit $i$ in layer $\ell$.
• Normalized Score ($u_i^\ell$): To ensure comparability across layers and time, the raw score is layer-normalized (divided by the layer mean) and smoothed using an Exponential Moving Average (EMA) with decay parameter $\beta$.

2.2 The Reset Mechanism

Every $f$ steps, CPR updates the weights of every neuron $i$ in layer $\ell$ according to:
$$W_i^{\ell, :} \leftarrow (1 - r_i^\ell) W_i^{\ell, :} + r_i^\ell \xi_i^{\ell, :}$$
$$W_{:, i}^{\ell+1} \leftarrow (1 - r_i^\ell) W_{:, i}^{\ell+1}$$

Where:

• $\xi$ is a weight drawn from the initialization distribution (e.g., Kaiming).
• $r_i^\ell$ is the reset coefficient, bounded by a maximum reset fraction $\rho \in (0, 1]$.
• The coefficient is determined by a monotonic decreasing function $\phi$ of the normalized utility $u_i^\ell$:
  $$r_i^\ell = \rho \cdot \phi(u_i^\ell)$$

2.3 The Transformation Function

The function $\phi$ maps utility to reset intensity. The authors use a sigmoid-based shape:
$$\phi(u_i^\ell) = \min\left(2\sigma[-\kappa(u_i^\ell - 1)], 1\right)$$

• Selectivity: Units with utility below the layer mean ($u < 1$) receive a reset fraction close to the maximum $\rho$. Units with high utility ($u > 1$) receive progressively smaller resets.
• Partiality: Even for the lowest utility units, the reset is partial (pulling toward initialization by $\rho$) rather than a full reinitialization. This avoids the gradient spikes associated with binary resets.
• Sharpness ($\kappa$): Controls how sharply the reset fraction drops as utility increases. The authors fix $\kappa=16$ based on ablation studies showing robustness across a wide range.

3. Key Contributions

1. Identification of Instability Source: The paper identifies binary neuron resets as a primary source of instability in long-horizon continual RL, distinguishing between the selection of units and the magnitude of their update.
2. CPR Optimizer: Introduction of a utility-scaled partial reset optimizer that avoids the brittleness of full reinitialization while maintaining the selectivity of reset methods.
3. Empirical Validation: Demonstration that CPR prevents policy collapse over 400 million training steps in the SlipperyAnt environment, a regime where binary reset methods and standard Adam fail.
4. Trade-off Control: Ablation studies reveal that the maximum reset fraction $\rho$ provides a tunable trade-off between plasticity (adaptation) and peak performance, allowing collapse prevention without requiring precise hyperparameter tuning.

4. Results

The authors evaluated CPR on four benchmarks: SlipperyAnt, SlipperyHumanoid, Continual MetaWorld, and Continual MinAtar.

• Policy Collapse Prevention: In SlipperyAnt (400M steps), CPR was the only method among 15 seeds to achieve zero policy collapses. In contrast, Adam collapsed, and binary reset methods (ReDo, ReGraMa) showed significant performance degradation or collapse.
• Performance: CPR achieved the highest Interquartile Mean (IQM) episodic return in SlipperyAnt and SlipperyHumanoid. It also outperformed baselines (CBP, ReDo, ReGraMa, Shrink & Perturb) in Continual MetaWorld and Continual MinAtar.
• Stability Metrics:
  • Dormant Neurons: CPR maintained a low, flat dormant-neuron ratio, comparable to CBP but without the late-training drift seen in Adam.
  • Gradient Norms: CPR exhibited smaller fluctuations in actor gradient norms compared to binary reset baselines, correlating with its ability to avoid collapses.
  • Parameter Norms: CPR constrained parameter growth, a known driver of plasticity loss, whereas baselines exhibited higher parameter growth.
• Ablations:
  • Sharpness ($\kappa$): Performance remained robust for $\kappa \in [2, 20]$. Low values ($\kappa \le 1$) degraded performance by making resets too uniform.
  • Reset Magnitude ($\rho$): A range of $\rho \in [0.01, 0.05]$ consistently prevented collapse. Lower $\rho$ preserved stability but hindered adaptation; higher $\rho$ increased plasticity but risked over-resetting useful structures.

5. Significance and Claims

The paper claims that CPR represents a qualitative shift in training reliability for long-horizon RL. By replacing all-or-nothing updates with utility-scaled partial resets, CPR injects necessary plasticity without the destabilizing sharp parameter updates that cause policy collapse.

The authors emphasize that while CPR introduces implementation complexity (calculating per-neuron utility and storing running averages), the computational overhead is marginal (~6% increase in wall-clock time for high-dimensional tasks) and the method is modular, requiring minimal changes to standard training pipelines (e.g., Optax).

The work concludes that utility-scaled reinitialization is a promising direction for continual learning, offering a tunable mechanism to balance the stability-plasticity trade-off. The authors acknowledge limitations, noting that CPR focuses on task dynamics and identity shifts and does not claim to resolve all forms of continual RL instability (e.g., reward non-stationarity or morphology changes). Future work is suggested to isolate the roles of utility estimation, partial reinitialization, and optimizer state handling across wider architectures.