The Rank and Gradient Lost in Non-stationarity: Sample Weight Decay for Mitigating Plasticity Loss in Reinforcement Learning

Zihao Wu; Hongyao Tang; Yi Ma; Jiashun Liu; Yan Zheng; Jianye Hao

cs.LGApr 2, 2026

The Rank and Gradient Lost in Non-stationarity: Sample Weight Decay for Mitigating Plasticity Loss in Reinforcement Learning

Authors: Zihao Wu, Hongyao Tang, Yi Ma, Jiashun Liu, Yan Zheng, Jianye Hao

Organizations: School of Computer Software, Tianjin University · International Joint Institute of Tianjin University, Fuzhou · The Hong Kong University of Science and Technology

Abstract

Deep reinforcement learning (RL) suffers from plasticity loss severely due to the nature of non-stationarity, which impairs the ability to adapt to new data and learn continually. Unfortunately, our understanding of how plasticity loss arises, dissipates, and can be dissolved remains limited to empirical findings, leaving the theoretical end underexplored.To address this gap, we study the plasticity loss problem from the theoretical perspective of network optimization. By formally characterizing the two culprit factors in online RL process: the non-stationarity of data distributions and the non-stationarity of targets induced by bootstrapping, our theory attributes the loss of plasticity to two mechanisms: the rank collapse of the Neural Tangent Kernel (NTK) Gram matrix and the $Θ(\frac{1}{k})$ decay of gradient magnitude. The first mechanism echoes prior empirical findings from the theoretical perspective and sheds light on the effects of existing methods, e.g., network reset, neuron recycle, and noise injection. Against this backdrop, we focus primarily on the second mechanism and aim to alleviate plasticity loss by addressing the gradient attenuation issue, which is orthogonal to existing methods. We propose Sample Weight Decay -- a lightweight method to restore gradient magnitude, as a general remedy to plasticity loss for deep RL methods based on experience replay. In experiments, we evaluate the efficacy of \methodName upon TD3, \myadded{Double DQN} and SAC with SimBa architecture in MuJoCo, \myadded{ALE} and DeepMind Control Suite tasks. The results demonstrate that \methodName effectively alleviates plasticity loss and consistently improves learning performance across various configurations of deep RL algorithms, UTD, network architectures, and environments, achieving SOTA performance on challenging DMC Humanoid tasks.

Figures & tables

Panel of Figure 1 — Figure 1: Aggregate Reliable metrics ( Agarwal et al., 2021 ) with 95% Stratified Bootstrap CIS.

Panel of Figure 2 — Figure 2: Empirical validation of SWD across TD3 in MuJoCo environments (mean $\pm$ std over 5 runs). SWD consistently improves sample efficiency and performance.

Panel of Figure 3 — Figure 3: Empirical validation of SWD across Double DQN in ALE environments (mean $\pm$ std over 5 runs). SWD consistently improves sample efficiency and performance.

Figure 4: Performance comparison between SWD and PER based on SAC. Aggregate Reliable metrics ( Agarwal et al., 2021 ) with 95% Stratified Bootstrap CIS in DMC tasks.

Panel of Figure 5 — Figure 5: Experiments conducted in the humanoid-run environment demonstrate that SWA exhibits a lower gradient magnitude, GraMa, and inferior performance, which validates our hypothesis.

Panel of Figure 6 — Figure 6: GraMa Metric in Humanoid Locomotion: Run, Walk, and Stand: The results clearly demonstrate that SWD effectively mitigates the loss of plasticity in humanoid robots across these key locomotor states.

Figure 7: Performance comparison across different UTD ratios (1, 2, 5) in Humanoid Run. SWD consistently outperforms uniform sampling across all UTD settings, with improvements ranging from 17.3% to 30.1%.

Figure 8: Performance comparison between SWD+S&P and other methods designed to address plasticity issues. Aggregate Reliable metrics ( Agarwal et al., 2021 ) with 95% Stratified Bootstrap CIS in Humanoid run.

Appendix figures & tables15 assets

Supplementary material from the paper’s appendix.

Appendix

Method	Complexity	Scale Dependency
Uniform Sampling	$\mathcal{O}(M)$	Independent of Buffer
Exact SWD	$\mathcal{O}(N+M)$	Linear w.r.t Buffer
Approximate SWD	$\mathbf{\mathcal{O}(B+M)}$	Linear w.r.t Buckets

Appendix

Table 1: Computational complexity comparison.

N

denotes the buffer size (

10^{6}

M

the batch size, and

B

the number of buckets (

2000

Method	Wall-Clock Time	Episode Return
Uniform	8.65 h	$190.46\pm 7.99$
Exact SWD	10.43 h	$229.01\pm 37.43$
Approximate SWD	8.70 h	224.93 $\pm$ 17.47

Appendix

Table 2: Runtime and performance comparison on Humanoid Run . The approximate method retains performance while significantly reducing training time.

Network Component	Actor Network	Critic Network †
Fully Connected Layer	(state_dim) $\to$ (256)	(state_dim + action_dim) $\to$ (256)
Activation	ReLU	ReLU
Fully Connected Layer	(256) $\to$ (128)	(256) $\to$ (128)
Activation	ReLU	ReLU
Output Fully Connected Layer	(128) $\to$ (action_dim)	(128) $\to$ (1)
Activation	Tanh ‡	None

Appendix

Table 3: Network Structures of the Twin Delayed Deep Deterministic Policy Gradient (TD3)

Layer	Input Channels	Kernel Size / Stride	Output Channels	Activation
Conv1	4	$8\times 8$ / 4	32	ReLU
Conv2	32	$4\times 4$ / 2	64	ReLU
Conv3	64	$3\times 3$ / 1	64	ReLU

Appendix

Table 4: Architecture of the Nature CNN Encoder used in Double DQN. The input consists of 4 stacked frames of size

84\times 84

Layer	Configuration	Activation
Input (Flatten)	$3136$ units ( $7\times 7\times 64$ )	-
FC1	Linear( $3136\to 512$ )	ReLU
Output	Linear( $512\to\|\mathcal{A}\|$ )	-

Appendix

Table 5: Architecture of the Double DQN Q-Network. The input is the flattened feature vector from the Encoder.

Layer/Operation	Input/Output Dimensions	Activation Function
Layer Normalization	(hidden_dim) $\to$ (hidden_dim)	None
Fully Connected (Expansion)	(hidden_dim) $\to$ (4 $\times$ hidden_dim)	ReLU
Fully Connected (Compression)	(4 $\times$ hidden_dim) $\to$ (hidden_dim)	None
Residual Connection	Input $\oplus$ Block Output ∗	None

Appendix

Table 6: Architecture of the SimBa Residual Block

Component	Structure & Dimension Flow
Input Projection (Fully Connected)	(input_dim) $\to$ (hidden_dim)
Residual Block Stack	$\times$ num_blocks † (each block follows Table 6 )
Final Layer Normalization	(hidden_dim) $\to$ (hidden_dim)

Appendix

Table 7: Architecture of the SimBa Encoder

Component	Actor Network	Critic Network
Input Dimension	(state_dim)	(state_dim + action_dim)
SimBa Encoder	hidden_dim=128; num_blocks=1	hidden_dim=512; num_blocks=2
Fully Connected	(128) $\to$ (action_dim)	(512) $\to$ (1)
Output Activation	Tanh ‡	None

Appendix

Table 8: Network Structures of the SimBa-SAC Framework

SWD -Specific Hyperparameters
Hyperparameter	TD3 Configuration
Actor Learning Rate	$10^{-4}$
Critic Learning Rate	$10^{-3}$
Discount Factor	$0.99$
Batch Size	$128$
Replay Buffer Size	$10^{6}$

Appendix

Table 9: Hyperparameters of the TD3 Algorithm

General Training
Hyperparameter	Value
Optimizer	Adam
Learning Rate	$1\times 10^{-4}$
Discount Factor ( $\gamma$ )	$0.99$
Buffer Size	$1\times 10^{6}$
Batch Size	$32$

Appendix

Table 10: Hyperparameters of Our Double DQN Implementation

Hyperparameter	SAC (with SimBa Encoder)
Optimizer	AdamW (weight decay = $10^{-2}$ )
Policy (Actor) Learning Rate	$1\times 10^{-4}$
Q-Network (Critic) Learning Rate	$1\times 10^{-4}$
Discount Factor	$0.99$
Batch Size	$256$
Warmup Steps (for Policy Update)	$5000$

Appendix

Table 11: Hyperparameters of Our SAC Implementation (with SimBa Encoder)

Panel of Figure 9 — Figure 9: SAC leanrning curve on DMC tasks

Panel of Figure 10 — Figure 10: Sensitivity analysis regarding the UTD. Data represents the mean $\pm$ std of five experimental runs conducted on the Humanoid Run.

Decay Steps	Minimum Weight Threshold ( $w_{\min}$ )
( $T_{\text{decay}}$ )	0.02	0.04	0.06	0.08	0.10
20,000	229.7 $\pm 26.4$	240.9 $\pm 37.1$	234.9 $\pm 15.0$	217.9 $\pm 38.6$	226.1 $\pm 23.7$
40,000	231.4 $\pm 44.4$	224.5 $\pm 34.4$	231.3 $\pm 30.8$	227.0 $\pm 23.1$	225.5 $\pm 22.5$
60,000	217.4 $\pm 40.2$	231.2 $\pm 29.9$	240.5 $\pm 55.0$	215.7 $\pm 27.9$	240.7 $\pm 35.3$
80,000	233.6 $\pm 42.9$	231.8 $\pm 35.7$	225.2 $\pm 42.6$	220.9 $\pm 17.6$	231.3 $\pm 54.4$
100,000	224.0 $\pm 32.1$	201.8 $\pm 31.9$	217.0 $\pm 48.5$	241.6 $\pm 38.4$	229.2 $\pm 29.5$

Appendix

Table 12: Parameter Sensitivity Analysis. Grid search results on Humanoid Run (Mean

\pm

Std). The best performance is marked in bold .

Decay Strategy	Episode Return	vs. Linear SWD
Linear Decay (Ours)	229.01 $\pm$ 37.43	–
SAC (Baseline)	190.46 $\pm$ 7.99	$-16.8\%$
Exponential Decay ( $\tau=1$ )	187.04 $\pm$ 29.85	$-18.3\%$
Polynomial Decay ( $p=2$ )	132.91 $\pm$ 11.12	$-42.0\%$

Appendix

Table 13: Performance comparison of different decay strategies on Humanoid Run . The relative difference is calculated with respect to our Linear Decay method.

Loss of Plasticity

Loss of plasticity is the declining ability of trained neural networks to learn new tasks or data after continued optimization.

See all 20 papers

Rethinking Plasticity in Deep Reinforcement Learning
Mar 22, 2026Zhiqiang HeLoss of Plasticity Non-Stationary RL
The University of Electro-Communications, Tokyo, Japan
Bridging Exploration and Loss of Plasticity in On-Policy Reinforcement Learning
Oct 26, 2026Andrea Vasco Grieco, Evelyn D'Elia, Edoardo Caldarelli +1Loss of Plasticity Proximal Policy Optimization
DAUIN, Politecnico di Torino, 10129 Turin, Italy · SAIL Research Unit, Istituto Italiano di Tecnologia, 16163 Genoa, Italy · Rehab Technologies Lab, Istituto Italiano di Tecnologia, 16163 Genoa, Italy
SPHERE: Mitigating the Loss of Spectral Plasticity in Mixture-of-Experts for Deep Reinforcement Learning
May 6, 2026Lirui Luo, Guoxi Zhang, Hongming Xu +2Loss of Plasticity Reinforcement Learning
State Key Lab of General AI, School of Intelligence Science and Technology, Peking University · State Key Laboratory of General Artificial Intelligence, BIGAI · Institute for Artificial Intelligence, Peking University

Mar 22, 2026J/K move · Enter open · S save

Rethinking Plasticity in Deep Reinforcement Learning

Zhiqiang He

The University of Electro-Communications, Tokyo, Japan

Loss of Plasticity Non-Stationary RL Neural Network Training Dynamics Neural Network Optimization

Open paper AskPDF arXiv

This paper investigates the fundamental mechanisms driving plasticity loss in deep reinforcement learning (RL), a critical challenge where neural networks lose their ability to adapt to non-stationary environments. While existing research often relies on descriptive metrics like dormant neurons or effective rank, these summaries fail to explain the underlying optimization dynamics. We propose the Optimization-Centric Plasticity (OCP) hypothesis, which posits that plasticity loss arises because optimal points from previous tasks become poor local optima for new tasks, trapping parameters during task transitions and hindering subsequent learning. We theoretically establish the equivalence between neuron dormancy and zero-gradient states, demonstrating that the absence of gradient signals is the primary driver of dormancy. Our experiments reveal that plasticity loss is highly task-specific; notably, networks with high dormancy rates in one task can achieve performance parity with randomly initialized networks when switched to a significantly different task, suggesting that the network's capacity remains intact but is inhibited by the specific optimization landscape. Furthermore, our hypothesis elucidates why parameter constraints mitigate plasticity loss by preventing deep entrenchment in local optima. Validated across diverse non-stationary scenarios, our findings provide a rigorous optimization-based framework for understanding and restoring network plasticity in complex RL domains.

Abstract

Loss of Plasticity

Explore similar work

Rethinking Plasticity in Deep Reinforcement Learning

Bridging Exploration and Loss of Plasticity in On-Policy Reinforcement Learning

SPHERE: Mitigating the Loss of Spectral Plasticity in Mixture-of-Experts for Deep Reinforcement Learning