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
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 Θ(k1) 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
Figure 1: Aggregate Reliable metrics ( Agarwal et al., 2021 ) with 95% Stratified Bootstrap CIS.
Figure 2: Empirical validation of SWD across TD3 in MuJoCo environments (mean ± std over 5 runs). SWD consistently improves sample efficiency and performance.
Figure 3: Empirical validation of SWD across Double DQN in ALE environments (mean ± 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.
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.
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
O(M)
Independent of Buffer
Exact SWD
O(N+M)
Linear w.r.t Buffer
Approximate SWD
O(B+M)
Linear w.r.t Buckets
Appendix
Table 1: Computational complexity comparison. N denotes the buffer size ( 106 ), M the batch size, and B the number of buckets ( 2000 ).
Method
Wall-Clock Time
Episode Return
Uniform
8.65 h
190.46±7.99
Exact SWD
10.43 h
229.01±37.43
Approximate SWD
8.70 h
224.93 ± 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) → (256)
(state_dim + action_dim) → (256)
Activation
ReLU
ReLU
Fully Connected Layer
(256) → (128)
(256) → (128)
Activation
ReLU
ReLU
Output Fully Connected Layer
(128) → (action_dim)
(128) → (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×8 / 4
32
ReLU
Conv2
32
4×4 / 2
64
ReLU
Conv3
64
3×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×84 .
Layer
Configuration
Activation
Input (Flatten)
3136 units ( 7×7×64 )
-
FC1
Linear( 3136→512 )
ReLU
Output
Linear( 512→∣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) → (hidden_dim)
None
Fully Connected (Expansion)
(hidden_dim) → (4 × hidden_dim)
ReLU
Fully Connected (Compression)
(4 × hidden_dim) → (hidden_dim)
None
Residual Connection
Input ⊕ Block Output ∗
None
Appendix
Table 6: Architecture of the SimBa Residual Block
Component
Structure & Dimension Flow
Input Projection (Fully Connected)
(input_dim) → (hidden_dim)
Residual Block Stack
× num_blocks † (each block follows Table 6 )
Final Layer Normalization
(hidden_dim) → (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) → (action_dim)
(512) → (1)
Output Activation
Tanh ‡
None
Appendix
Table 8: Network Structures of the SimBa-SAC Framework
Hyperparameter
TD3 Configuration
Actor Learning Rate
10−4
Critic Learning Rate
10−3
Discount Factor
0.99
Batch Size
128
Replay Buffer Size
106
SWD -Specific Hyperparameters
Appendix
Table 9: Hyperparameters of the TD3 Algorithm
Hyperparameter
Value
General Training
Optimizer
Adam
Learning Rate
1×10−4
Discount Factor ( γ )
0.99
Buffer Size
1×106
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×10−4
Q-Network (Critic) Learning Rate
1×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)
Figure 9: SAC leanrning curve on DMC tasks
Figure 10: Sensitivity analysis regarding the UTD. Data represents the mean ± std of five experimental runs conducted on the Humanoid Run.
Decay Steps
Minimum Weight Threshold ( wmin )
( Tdecay )
0.02
0.04
0.06
0.08
0.10
20,000
229.7 ±26.4
240.9 ±37.1
234.9 ±15.0
217.9 ±38.6
226.1 ±23.7
40,000
231.4 ±44.4
224.5 ±34.4
231.3 ±30.8
227.0 ±23.1
225.5 ±22.5
60,000
217.4 ±40.2
231.2 ±29.9
240.5 ±55.0
215.7 ±27.9
240.7 ±35.3
80,000
233.6 ±42.9
231.8 ±35.7
225.2 ±42.6
220.9 ±17.6
231.3 ±54.4
100,000
224.0 ±32.1
201.8 ±31.9
217.0 ±48.5
241.6 ±38.4
229.2 ±29.5
Appendix
Table 12: Parameter Sensitivity Analysis. Grid search results on Humanoid Run (Mean ± Std). The best performance is marked in bold .
Decay Strategy
Episode Return
vs. Linear SWD
Linear Decay (Ours)
229.01 ± 37.43
–
SAC (Baseline)
190.46 ± 7.99
−16.8%
Exponential Decay ( τ=1 )
187.04 ± 29.85
−18.3%
Polynomial Decay ( p=2 )
132.91 ± 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.
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.
Zhiqiang He
The University of Electro-Communications, Tokyo, Japan
Deep reinforcement learning (RL) algorithms have been observed to lose the ability to adaptively learn from new data, a phenomenon termed Loss of Plasticity (LoP). While many metrics have been shown to correlate with it, the causal mechanisms and the geometric conditions from which it can arise have not been clearly identified. We analyze LoP in Proximal Policy Optimization with isotropic Gaussian policies, the state-of-the-art formulation for many continuous action spaces in RL. We report multiple metrics which characterize the effect that the policy standard deviation σ has on LoP, and we also suggest related explanations for the effectiveness of widely adopted LoP mitigation techniques. We observe a significant effect of σ on the extent of LoP, and accompany this with both theoretical and experimental supporting arguments. Our experimental validation consists of LoP benchmark evaluations against standard baselines as well as real-world policy deployment on a 7-DoF robotic arm.
Andrea Vasco Grieco, Evelyn D'Elia, Edoardo Caldarelli +1
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
In deep reinforcement learning (DRL), an agent is trained from a stream of experience. In a continual learning setting, such agents can suffer from plasticity loss: their ability to learn new skills from new experiences diminishes over training. Recently, Mixture-of-Experts (MoE) networks have been reported to enable scaling laws and facilitate the learning of diverse skills. However, in continual reinforcement learning settings, their performance can degenerate as learning proceeds, indicating a loss of plasticity. To address this, building on Neural Tangent Kernel (NTK) theory, we formalize the plasticity loss in MoE policies as a loss of spectral plasticity. We then derive a tractable proxy for spectral plasticity, one expressible in terms of individual expert feature matrices. Leveraging this proxy, we introduce SPHERE, a practical Parseval penalty tailored for MoE-based policies that alleviates the loss of spectral plasticity. On MetaWorld and HumanoidBench, SPHERE improves average success under continual RL by 133% and 50% over an unregularized MoE baseline, while maintaining higher spectral plasticity throughout training.
Lirui Luo, Guoxi Zhang, Hongming Xu +2
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