PLASTIC: Improving Input and Label Plasticity for Sample Efficient Reinforcement Learning
Authors: Hojoon Lee
Hanseul Cho
Hyunseung Kim
Daehoon Gwak
Joonkee Kim
Jaegul Choo
Se-Young Yun
Chulhee Yun
Venue: NeurIPS 2023 poster (2023)
Keywords: Reinforcement Learning
Sharpness Minimization
Generalization
Plasticity
Deep Learning
TL;DR: We introduce PLASTIC, a simple-to-use algorithm which address the loss of plasticity phenomenon and enhance sample efficiency in Reinforcement Learning.

Abstract: In Reinforcement Learning (RL), enhancing sample efficiency is crucial, particularly in scenarios when data acquisition is costly and risky. In principle, off-policy RL algorithms can improve sample efficiency by allowing multiple updates per environment interaction. However, these multiple updates often lead the model to overfit to earlier interactions, which is referred to as the loss of plasticity. Our study investigates the underlying causes of this phenomenon by dividing plasticity into two aspects. Input plasticity, which denotes the model's adaptability to changing input data, and label plasticity, which denotes the model's adaptability to evolving input-output relationships. Synthetic experiments on the CIFAR-10 dataset reveal that finding smoother minima of loss landscape enhances input plasticity, whereas refined gradient propagation improves label plasticity. Leveraging these findings, we introduce the **PLASTIC** algorithm, which harmoniously combines techniques to address both concerns. With minimal architectural modifications, PLASTIC achieves competitive performance on benchmarks including Atari-100k and Deepmind Control Suite. This result emphasizes the importance of preserving the model's plasticity to elevate the sample efficiency in RL. The code is available at https://github.com/dojeon-ai/plastic.

PLASTIC: Improving Input and Label Plasticity for Sample Efficient Reinforcement Learning
Hojoon Lee"
Hanseul Cho" Hyunseung Kim' Daehoon Gwak Joonkee Kim
Jaegul Choo Se-Young Yun Chulhee Yun
Kim Jaechul Graduate School AI, KAIST fjoonlesky, jhs4015, mynsng, daehoon gwak, joonkeekim, jchoo, yunseyoung, chulhee yun yunnakaisttteeee ac. kr
Abstract
In Reinforcement Learning (RL), enhancing sample efficiency is crucial, particu- larly in scenarios when data acquisition is costly and risky. In principle, off- off-policy RL algorithms can improve sample efficiency by allowing multiple updates per environment interaction. However, these multiple updates often lead the model overfit to earlier interactions, which referred to as the loss of plasticity. Our study investigates the underlying causes of this phenomenon by dividing plasticity into two aspects. Input plasticity, which denotes the model's 's adaptability to changing input data, and label plasticity, which denotes the model's adaptability to evolv- ing input-output relationships. Synthetic experiments on the CIFAR- 10 dataset reveal that finding smoother minima of loss landscape enhances input plasticity, whereas refined gradient propagation improves label plasticity. Leveraging these findings, we introduce the PLASTIC algorithm, which harmoniously combines techniques to address both concerns. With minimal architectural modifications, PLASTIC achieves competitive performance on benchmarks including Atak and Deepmind Control Suite. This result emphasizes the importance of preserving the model's plasticity to clevate the sample efficiency in RL. The code available at https: / /ttg/ggtttt///.t//// com/ dojedooooooooo www-aaappppaaa.aaaaaa
Introduction
In Reinforcement Learning (RL), achieving sample effi- ciency is crucial in various domains including robotics, au- tonomous driving, and healthcare, where data acquisition constrained and expensive [42, 36). In theory, off-policy RL algorithms promise increased sample efficiency by al- lowing multiple updates of policy or value functions from single data instance [26, 24, 43]. However, they tend suffer from the loss of plasticity, a phenomenon where the models overfit earlier interactions and fail adapt
Inter Quartl Mean (IQM)
0.60
PLASTIC" (RR8)
PLASTIC (RR8)
SR-SPR(RRB)
(RR4)
(RR4)
0.50
(RR4)
DreamerV3
0.45
MLR
(RR2)
0
(RR2)
PlayVirtual
new experiences [52, 19, 43].
The origins of the loss of plasticity are focus of contem- porary research. One avenue of study points is the role of the smoothness of the loss surface. Models seeking smoother minima of loss landscape tend to exhibit more stable learning patterns and enhanced plasticity [46, 47].
Runtime (GPU Hours)
Figure 1: Scaling of proposed method, PLASTIC, on the Atari-100k benchmark compared to state-of-the-aat methods (Section 4.3).
Equal contributions.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Synthetic Experiment
Benchmark
Input Plasticity
Label Plasticity
Atari 100k tasks
DMC Suite
tasks
80
70
9.40
500
70
68
0.30
65
SPAA
PLASTRC
Booe
Pt. PLASTIC
SAA
C.nel
200
SAH
CAelU
PLASTIC
611
Figure 2: Left: Performance of Synthetic Experiments. Layer Normalization (LN) and Sharpeness- Aware Minimization (SAM) considerably enhance imput plasticity, while their effect on label plasticity marginal. Conversely, Concatenated ReLU (CReLU) and periodic reinitialization (Reset) predomi- nantly improve label plasticity with subtle benefits input plasticity. Right: Performance of RI. Benchmarks. PLASTIC consistently outperforms individual methods, highlighting the synergistic benefits of its integrated approach.
Another line of investigation emphasizes the role of gradient propagation, highlighting that as train- ing continues, the saturation of active units can impede weight updates, diminishing the model's adaptability [18, 2].
To understand the interplay of these factors within the context of RL, we designed couple of synthetic experiments. The RL agent often confronts two adaptation scenarios, which involve adjusting to new data inputs and evolving input-label dynamics. In particular, action once considered to be positive but later perceived as negative. Therefore, throughout this work including synthetic experiments, categorize the model's plasticity into cases:
example of the latter might be
Input Plasticity: adaptability to shifts in input data distributions, p(z).
Label Plasticity: adaptability to evolving labels (returns) for given inputs (observations), p(y
For our exploration, we employ the CIFAR-10 dataset [37]. To assess input plasticity, we train the model from progressively growing dataset, where the data is segmented into 100 chunks and sequentially added to the buffer. On the other hand, label plasticity is examined by altering labels at intervals, reflecting the dynamic nature of input-output associations in RL. We also adopt strategies to enhance loss smoothness and gradient propagation in the tested models. The Sharpeness-Awatt Minimization (SAM) optimizer [21] and Layer Normalization (LN) [5] are employed for the former, while periodic reinitialization of the last few layers (Reset) [81, 52] and Concatenated ReLU (CReLU) activations [61, 2] used to enhance the latter.
Our synthetic experiments yield clear distinctions between the proposed two concepts of plasticity. As illustrated the left side of Figure 2, smoother minima of loss largely improve input plasticity, while maintaining the amount gradient propagation has a pronounced impact on label plasticity. Building these findings, we propose algorithm called PLASTIC, which combines SAM optimizer, LN, periodic Reset, and CReL.U activation to improve both input and label plasticity. The simplicity of PLASTIC facilitates its seamless integration into standard off-policy RL frameworks, requiring minimal code modifications. Notably, across the Atari and Deepmind Control Suite benchmarks, PLASTIC achieved competitive performance the leading methods, underscoring its potential in enhancing RL's sample efficiency.
In summary, main contributions are listed as follows:
Through synthetic experiments, we find out that loss smoothness and refined gradient propagation play separate roles in improving the model plasticity (Section 3).
We introduce the PLASTIC, a simple-to-use and efficient algorithm that improves the model's plasticity by seeking smooth region of loss surface and preserving gradient propagation.
Empirically, PLASTIC achieved competitive performance on challenging RL benchmarks, includ- ing Atari- 100k and Deepmind Control Suite (Section 4.2-4.5).

---

2 Preliminaries
2.1 Sample Efficient Reinforcement Learning
Achieving sample efficiency is crucial for applying RL to real-world problems. Constraints such as limited online data collection (c.g.. robotics, educational agents) safety considerations (c.g.. autonomous driving, healthcare) emphasize the need for efficient RL. algorithms [36, 42].
Off-policy RL. algorithms like Rainbow [26] and SAC [24] offer improved sample efficiency by permitting multiple policy or value function updates based on single data collection. However, this advantage may introduce a pitfall: the increased update rate can lead to overfitting, undermining the model's s generalizability and adaptability to new datasets [29, 43].
To overcome these challenges, various methods have been proposed:
Data Augmentation: Renowned for its effectiveness in computer vision [16, 80], data augmen- tation excels in visual RL algorithms [73, 72, 40, 55, 68), particularly when combined with self-supervised learning methods [39, 58, 77].
Diverse regularization strategies have also demonstrated their effectiveness, including L2 regularization [45], spectral normalization [23, 11], dropout [27], and both feature mixing [12] and swapping [10].
Regularization:
Self-Supervised Learning: By incorporating auxiliary learning objectives, self self-supervised learning has emerged as a potential solution. Objectives encompass pixel or latent space reconstruction [74, 77], future state prediction [58, 19, 41J, and contrastive learning focusing on either instance [39, 20] or temporal discrimination [53, 64, 48].
Despite these advancements, a key question remains. Why does an overfitted model face challenges when adapting to new datasets?
2.2 Understanding Plasticity in Reinforcement Learning
The question for understanding model plasticity RL driven by the inherent need for plasticity as agents consistently encounter new inputs and evolving input-output relationships.
seminal contribution by Lyle et al. [47] highlighted the importance of smoother loss landscapes. By applying layer normalization across all convolutional and fully connected layers, they managed to flatten the loss landscapes, subsequently enhancing performance emphasis on smooth loss landscapes is relatively new in RL, its importance has been substantiated in supervised learning, where empirical evidence suggests models converged on wider and smoother loss surface generalize better to unseen datasets [34, 30, 71J. Following this insight, the sharpness- aware minimization (SAM) optimizer [21] has recently gained attention in supervised learning [13, 44), aiming to minimize both training loss and its sharpness.
Atari environments. While the
Concurrently, various studies point out a progressive decline in network's active units as a probable cause for the loss of plasticity [2, 18]. As neural networks iteratively adjust their weights to minimize training losses, the number of active units tends to shrink, often culminating the dead ReLU phenomenon [28, 17]. This reduction in active units hampers the gradient propagation to upper layers, thus impairing network adaptability. Proposed remedies include periodic neuron reinitialization [52, 19, 63] or employing Concatenated ReLU activation [61, 2], both showing promise in RL.
With these observations in mind, our primary endeavor is delve into the complex interplay of these factors. We aim to analyze whether these factors synergistically affect the model's plasticity or operate as distinct, individual contributors.
2.3 Off-Policy Reinforcement Learning Algorithms
Rainbow. Rainbow [26] is a widely used off-policy algorithm for discrete control settings that
integrates six different extensions to the standard DQN algorithm [50]. The extensions include Double Q-learning [66], Prioritized Experience Replay [57], Dueling Networks [69], Multi-Step
Return, Distributional Q-function [8], and Noisy Networks [22]. Rainbow significantly improves the
performance and robustness of the standard DQN algorithm, thereby addressing shortcomings in function approximation and exploration-cxploitttioo trade-offs.

---

The Q-value updates in the Rainbow algorithm follow the principle of minimizing the Temporal- Difference (TD) error which is defined as:
((w, w . = 1Q (s, - (r max Qw as),
where w denotes the model weights, w- is the weights of the target model, and (8, a, s') is the transition sampled from the replay buffer B.
Noisy Networks [22] introduce stochasticity into the model weights, reparameterizing them as:
μ
where la and a are the learnable weights, whereas a random vector sampled from an isotropic Gaussian distribution, adding randomness to the layer. (We denote element-wiss product by *..)
Soft Actor-Critic. Soft Actor-Critic (SAC) [24], a prevalent off off-policy algorithm for continuous control, aims to maximize the expected return coupled with an entropy bonus. SAC consists of policy т and a critic model Q, each parameterized by weights and ur respectively.
In training, the critic Qw trained to minimize the following objective, defined as
(0, w; T) (s, a) (r + (Q (s', (s', a') log ( a')))?² d'~ (e-ss) (3)
where a is the entropy coefficient and r = (s, a, T, s') is the transition sampled from the buffer Subsequently, the policy me is jointly trained to maximize the following objective:
(0 w:
(Qcemol(@@uia (s. a) a al (a(a)))]
(4)
Throughout this paper, all of the integrated and evaluated on top of these standard off-policy RL algorithms.
methods (LN, SAM, Reset, CReLU) will be
3 Synthetic Experiments
Analyzing the loss of plasticity in RL is indeed intricate given RL's multifaceted nature. encompassing challenges such as credit assignments, noisy targets, and the exploration-exploiation tradcoff. To alleviate this complexity, we design synthetic supervised learning experiments within controlled framework using the CIFAR-10 10 dataset [37]. Our focus is to evaluate the model's '' adaptive capabiliites under two distinct adaptation scenarios:
Input A Adaptation: In RL, as the agent continually interacts with the enviromment, it constantly encounters new data through exploration. Effective adaptation to this data paramount for effective decision-making To simulate this scenario, we partitioned the training data into 100 chunks, sequentially adding them to a buffer. The training was conducted by sampling data from this progressively growing buffer. Overfitting to earlier data chunks and failing to adapt to newer ones would indicate performance degradation compared to a model trained on the entire dataset.
Label Adaptation: The RL domain often experiences shifts in the relationships between inputs and labels. Our synthetic experiment mirrored this dynamic by periodically altering labels during the training phase. The labels were randomly shuffled 100 times, with each class's 's labels uniformly reassigned, ensuring consistent reassignment within a class (e.g., all 'cat' images transition from class 3 to class 4). A model's overfitting to initial relationships would impede its capability to adeptly adapt evolving relationships.
Our experimental setup resonates with common design choices in RL [50, 26, 73], comprising three convolutional layers for the backbone and three fully connected layers for the head. Employing Stochastic Gradient Descent (SGD) with momentum [56, 54] as the optimizer, and a batch size of 128, the model underwent 50,000 updates, with 500 updates each alternation step (i.e., appending a chunk for input adaptation and alternating labels for label adaptation). An exhaustive sweep was carried out for the learning rate and weight decay range from {0.1, 0.01, 0.001, 0.0001, 0.00001); }. integrating weight decay to minimize variance across individual runs. Each methodology was trained over 30 random seeds, with results presented within 95% confidence interval.
To delve into the impact of pursuing a smooth loss surface and preserving gradient propagation, selected four distinct methods: Layer Normalization (LN) and Sharpness- Aware Minimization

---

Improvement individual methad
Improvement cambined methods
13
75
Inpus Adspttation
Adaptation Adiaptationn
20-
Labe Adaptation
CReLU-Restt
Figure 3: Comparative performance improvements various methods and their combinations on input and label adaptation tasks. The left panel showcases individual method performance, while the right panel delves into the synergistic benefits combining multiple methods.
(SAM) optimizer for the former, and Concatenated Rectified Linear Unit (CReLU) activation and Reset for the latter. For the SAM optimizer, the perturbation parameter was tuned across the set {0.1, 0.03, 0.01}, and for the Reset method, the reset interval was tuned over {5, 10, 203 data chunks. Our analysis extended to exploring the synergistic interplay between these methods. To this end, we examined all unique pair-wise combinations of the selected methods (LNN SAM, CReI.U + Reset) with combination of all methods. This exploration aimed to unveil synergetic effects between different methods in promoting smoother model landscapes and enhancing gradient propagation.
In Figure 3, we observed clear distinction between the effects of various methods on input and label adaptation capabilities. The left panel illustrates the performance applying individual methods. We found that LN and SAM excel in enhancing input adaptation, yet only offer marginal improvements for label adaptation. In contrast, CReLU and Reset yield modest enhancements input adaptation while exhibiting a significant impact on label adaptation.
Turning attention the right panel, the focus is the combined synergies of methods. The pairing of LN with SAM predominantly enhances input adaptation. When LN is merged with either CReL.U Reset, there's marked improvement across both adaptation types. Likewise, the combination of SAM with either CReLU or Reset benefits both input and label adaptation scenarios. Among these combinations, the CReLU and Reset pairing is skewed towards improving label adaptation. Notably, PLASTIC, which integrates all of these methods yielded the most significant enhancements.
In summary, while individual methods are tailored towards specific enhancement areas, the integration of these techniques can yield synergies across both adaptation landscapes. These findings reinforce notion of distinct yet complementary roles of secking smooth loss surface and enhancing gradient propagation to improve the model's plasticity.
Experiments
In this section, we evaluate the effectiveness of enhancing the model's plasticity towards achieving sample-efficient reinforcement learning. Our experiments span two widely recognized benchmarks: the Atok [9] for discrete control tasks and the DeepMind Control Suite (DMC) [65] for continuous control tasks. We have organized our experimental pursuits into four main subsections:
Assessing the performance of the PLASTIC algorithm across both domains (Section 4.2).
Analyzing the scaling behavior of PLASTIC. Specifically, we focus on the model's responsiveness with respect to the number of updates per environment interaction (Section 4.3).
Exploring the advantages of improving plasticity large pre-trained model (Section 4.4).
Ablation study for different combinations of plasticity plasticity-preserving preserving methods (Section 4.5).
4.1 Experimental Setup
Atari-100k. Following the standard evaluation protocol from [32, 58), we evaluate the performance of 26 games on Arcade Learning Environments [73], limited to 100k interactions. The results are
5

---

measured by averaging over 10 independent trials. We use Data- BataaRRegllriizdd Q-learning (DrQ) [73] as base algorithm of PLASTIC, which is built on top of the Rainbow algorithm [26] and utilize data augmentation techniques to alleviate overfitting.
Deepmind Control Suite Medium (DMC-M). To evaluate performance in the DMC benchmark, we adopt the evaluation protocol and network architecture identical to Nikishin et al. [52]. This benchmark consists of 19 continuous control tasks (i.e., 8 easy tasks and 11 medium tasks), where the agent controls its behavior from raw pixels. We selected 11 medium tasks and trained the agent for million environment steps. Within this setting, we use Data- Data-Regularized Q-learning (DrQ) algorithm [73], which incorporates data augmentation techniques in conjunction with the soft actor-critic algorithm [24]. The results are averaged over 10 random seeds.
Baselines. This section presents selection of plasticity-preservinn methods designed to seek smooth loss surfaces or improve gradient propagation in RL. To seek a smooth loss surface, we consider Layer Normalization (LN) [47] and Sharpness-Awann Aware Minimization Optimizer (SAM). For LN, we apply layer normalization after each convolutional layer in the backbone network. For SAM, we vary the perturbation parameter p from {0.03, 0.1, 0.3), selecting the model with the best performance.
For enhancing gradient propagation, consider the re-initialization of the part of the networks (Reset) [52] and Concatenated ReLU activations (CReLU) [2]. Reset involves periodically resetting the parameters of the head layers to their initial distribution. Following [52], reinitialize the parameters of the head layers' for every 40,000 gradient updates in Atari and 100,000 gradient updates in DMC. For CReLU, we replace ReLU activation with concatenated ReLU, which preserves both positive and negative features.
Metrics. To gain a deeper understanding of the underlying geometry of the loss landscape (i.e., smoothness), we employ the maximum eigenvalue of the Hessian (Amex) [33]. This metric provides insights into the curvature of the loss function, where larger value indicates a sharper and more intricate optimization landscape and smaller value corresponds to a smoother and potentially more favorable landscape for optimization. To quantify the gradient propagation, we periodically record the fraction of the active units in the feature map by processing 512 transitions [47, 61, 18], randomly sampled from the replay buffer. When using ReLU activation, values below zero become inactive, which does not contribute to the updates of incoming weights.
To evaluate the performance of the agent, we report a bootstrapped interval for the interquartile mean (IQM), median, mean, and optimality gap. following the guidelines from [3]. For each environment, we average the performance of the trajectories at the end of training, 100 for Atari and 10 for DMC. For Atari, we normalize these scores Human Normalized Score (HNS) as HNS agent score random score which quantifies the relative performance of the agent compared to human score - random score
human-level performance. More details of the experimental setup can be found in the Appendix.
4.2 Main Experiments
In this section, we present the experimental results aimed evaluating the sample efficiency of the PLASTIC algorithm. Our examination involved contrasting its efficacy against individual plasticity- preserving methods. For this purpose, we used the Atari-100k and DMC-Medium benchmarks.
Table 1: Performance on Atari-100k and DMC-M. The results are averaged over 10 random seeds.
Atari-100k
DMC-M
Method
IQM
Median
Mean
IQM
Median
Mean
Base
0.258 (0.244.0.2 0.292)
0.277 (0.209, 0.295)
0.476 (0.432, 0.520)
213 (146. 304)
288 (223, 356)
293 (219, 347)
SAM [21]
0.325 (0:206, 0.354)
0.327 (0.254, 0.368)
0.501 (0.465, 0.537)
278 (196, 371)
341 (264, 404)
332 (277, 500)
LN [47]
0.259 (0.235, 0.293)
0.247 (0.218, 0.266)
0.463 (0.403, 0.522)
412 485)
415 (951.491)
421 (366.477) 477)
CReLU [2]
0.256 0.227)
0.193
0.252)
0.498 (0.446, 0.553)
338
398
394 522)
Reset [52]
0.343 0.373)
0.291 (0.231, 0.369)
0.660 (0.611.
514
491 (430,
498 (442, 555)
PLASTIC
0.421 (0.388. 0.457)
0.347 (0.266, 0.422)
0,933 (0.812. 1.067)
565 (498. 626)
540 (476,
537 587)

---

Table I showcases the comparative outcomes of these methods. From the results, we observed that while individual plasticity presstving methods enhanced sample efficiency across both discrete and continuous control benchmarks, the PLASTIC algorithm outperformed them all. This integration of multiple methods within PLASTIC clearly demonstrates superior results over any single approach. comprehensive set of results for each environment is available in the Appendix.
Maximum Elgen Value Umna
30
1.0
25
(Head) Fraction of AU
0.8
None
SAM
Reset
PLASTIC
CRe
0.6
15
0.4
10
%.0
None
LN
SAM CReLU
Reset PLASTIC
0.2
0.4
0.6
0.8
1.0
Method
Environment Steps (x10f)
Figure Left:
The maximum eigenvalue Hessian (Amax (V2L), representing the curvature of the loss landscape. Right: Fraction of active units in the head layers. Both metrics were derived from the DrQ algorithm's 's evaluation on the Atari- benchmark.
To delve deeper into the mechanisms driving these performance enhancements, we focused on two specific metries: the maximum eigenvalue of the Hessian (Amax) and the fraction of active units in the head layers. As illustrated in Figure 4, methods such as LN and SAM encourage the model to converge on smoother loss landscape. On the other hand, CReLU and Reset exhibit a pronounced presence of active units. Notably, by integrating these methods, PLASTIC succeeded in converging smooth loss region while maintaining a large number of active units.
Our prior synthetic experiments provided insight that convergence on smoother loss landscape improves the model's input plasticity while enhanced gradient propagation enhances the model's label plasticity. Consequently, we conjecture that PLASTIC's superior performance is attributed to its capability to synergistically enhance both input and label plasticity.
4.3 Scaling Replay Ratio
Off Off-policy RL algorithms, in theory, promise enhanced sample efficiency by increasing the updates per environment interaction, commonly known as the replay ratio. Yet, a practical challenge arises: escalating updates often inversely affect sample efficiency due to the loss of the model's plasticity. This section delves into whether our PLASTIC algorithm can combat this decrement in plasticity, thereby improving sample efficiency with scaling replay ratio.
For our analysis, we scaled the replay ratio to 8 and evaluated the efficacy of PLASTIC on the Atari-100k benchmark. addition, we compared PLASTIC against notable ofate-of-the-art (SOTA) model-based and model-free algorithms.
To ensure consistency and comparability in our evaluation, we computed the GPU hours for every method by leveraging their official codebases. For instance, Play [76], MLR [78], SR-SPR [19], and PLASTIC can parallelize multiple runs single GPU. Thus, for a fair comparison, followed the protocol as outlined by [19], dividing the total execution time by the total number of executed runs.
Given that SAM requires an auxiliary gradient computation during optimization, and CReLU increases the number of parameters, introduced simplified alternative, PLASTICI, This simplified version only integrates Layer Normalization (LN) and the reset mechanism, ensuring minimal computational overhead. In addition, for each reset, we incorporate the Shrink & Perturb [4] method on the backbone network parameters. Shrink & Perturb softly reinitializes the model's parameters as d a + a) where é ~ initializer. While Shrink & Perturb slightly lower performance at a replay ratio of 2, it demonstrated enhanced sample efficiency when scaling the replay ratio upwards.
As depicted Table and Figure 1, PLASTIC establishes Pareto frontier between Inter-Quartia. Mean (IQM) and computational cost, illustrating its notable computational efficiency. While Ef-
7

---

Table 2: Comparison to the SOTA on Atari-100k. For IRIS. DreamerV3, EfficientZero, Play Virtual, MLR, SR SR-SPR, and PLASTIC the results are averaged over 5, 5, 32, 15, 3, 10, and 5 seeds respectively.
Type
Method
Search
Params (M)
RR
GPU hours
TQM
Median
Mean
OG
IRIS [49]
30.4
36.3
0.501
0.289
1.046
0.512
Model-Based
Dreamer v3 [25]
17.9
5,1
0,497
0,466
1.097
0.505
EfficientZero [75]
V
8.4
28.0
n/a
1.090
1.943
n/a
Play Virtual [76]
7.4
2
4.9
0.374
n/a
n/a
0.558
MLR [78]
161.7
2
3.7
0.432
n/a
n/a
0.522
2
2.7
0.444
0,336
0.910
0.516
SR-SPR [19]
7.3
4
5.3
0,544
0.523
1.111
0.470
Model-Free
8
10.2
0.589
0,560
1.188
0.452
2
0.5
0.396
0.425
0.702
0.541
PLASTIC
6.8
1.0
0.518
0.517
0.858
0.478
8
1.9
0.583
0.542
0.939
0.448
2
1.0
0.421
0.347
0,933
0.535
PLASTIC
7.2
1.9
0,545
0.407
1.002
0.475
3.8
0.571
0.494
0.968
0.461
ficientZero exhibits sfate-of-the-art performance. uniquely employs search algorithm coupled with domain-specific heuristic, namely early environment resets. When excluding these specific elements, PLASTIC stands out strong competitor against the other methods.
Through these results, we found that PLASTIC can effectively prevent the loss of plasticity by incorporating various plasticity-preservvvnn strategies. Furthermore, for practitioners, the construc- tion of the Pareto frontier by PLASTICt is particularly beneficial. By simply incorporating layer normalization and implementing resets in the underlying off-policy RL algorithms, they can achieve improved sample efficiency. This comes with the advantage of minimal computational overhead and only requires minor code adjustments.
In conclusion, our experiments underscore the potential of leveraging enhanced plasticity as a means to improve sample efficiency in RL.
4.4 Compatibility with Large Pretrained Model
Recent years have seen growing interest in leveraging large pretrained models to improve sample efficiency in RL [59, 60, 7]. We investigate how combining PLASTIC's principles with these models can counteract the loss of plasticity, a common obstacle to rapid adaptation.
For our evaluation, we selected the SimTPR model [41]. Using 30-layer convolutional network, SimTPR is pretrained via self-predictiv representation learning on suboptimal Atari video datasets. As SimTPR does not employ Layer Normalization (LN) and Concatenated ReL.U (CReLUU acti- vations, we employ SAM to seek smoother loss landscape and use Reset techniques to facilitate gradient propagation. Following the approach in section 4.3, for each reset, we applied Shrink Perturb to the backbone network. In the fine-tuning phase, we initialized an MLP-based head network on top of the frozen, pretrained backbone network. This head network underwent training for 100k environment steps, leveraging the Rainbow algorithm. See the Appendix section for more details.
Table provides a summary of applying either SAM or Reset to the pretrained model. Solely scaling the replay ratio during fine-tuning leads to noticeable decrement in its efficacy. However, we observed that the usage of Reset or SAM can counteract this decrease, leading to a pronounced enhancement performance. Furthermore, the integration of both techniques surpassed individual contributions, indicating the presence of synergy.
Table 3: Fine-tuning from pretrained model. Reset applies soft reset the backbone and hard reset to the head.
RR
SAM
Reset!
IQM
0.366 (0.324, 0.397)
0.709 (0.650, 0.745)
0.776 (0.703, 0.854)
These observations imply that elevating both input and label plasticity can play a pivotal role in enhancing sample efficiency, even in large, pre-trained models.
0.243 (0.214, 0,267)
0.780 (0.706, 0.865)
0.834 (0.769, 0.889)

---

4.5
Ablation Studies
To further analyze the interplay between seeking smooth loss surface and improving gradient propagation, we have conducted a series of ablation studies for different combinations of plasticity- preserving methods. Here, we averaged the performance over 10 random seeds.
Table showcases the results when applying vari- ous combinations to the Atari-100k benchmark. The first five and the last rows echo the results from Ta- ble 1. It becomes clear that the concurrent pursuit of a smooth loss surface and the enhancement of gradient propagation surpasses the individual applica- tion of either strategy (for instance, LN CReLU > LN ~ CReLU). The comprehensive integration of all methods, referred to as PLASTIC, demonstrates the highest performance.
Table 4: Performance comparison of various plasticity-prest prreserinn method combinations on the Atari-100k benchmark.
SAM
CReLU
Reset
IQM
(0.224, 0.292)
0.259 (0.235, 0.293)
0.325 (0.296, 0.354)
0.256 (0.224, 0.287)
0.343 (0.314, 0.373)
0,341 (0.325, 0.366)
0.284 (0.257, 0,314)
For practitioners considering our approach, it's worth highlighting its adaptability. The methodology does not necessitate the application of all methods across every plasticity category to reap substantial benefits. Instead, deploying just one technique from each cat- egory can yield notable improvements. For scenarios constrained by computational resources, using LN and Reset is recommended over SAM, due the latter's extra computational burden, and CReLU, given its added parameters (as noted in PLASTICI in Section 4.3). Additionally, when one is working with pre-trained network where altering the overall architecture is not feasible, LN and CReLU might not be practical since they inherently change the network's 's structure. In such instances, as detailed in Section 4.4, the combination of SAM and Reset emerges as potent solution to improve downstream performance.
(0.365, 0.430)
0.372 (0.357. 0,408)
0.411 (0.377, 0.447)
0.373 (0.344, 0.402)
0.421 (0.388, 0,457)
5 Conclusion, Limitations, and Future Work
In this paper, we aimed to address the prevalent issue of the "loss of plasticity" in off-policy RL algorithms. Through synthetic experiments on CIFAR-10, we discovered that finding a smoother point of loss landscape largely improves input plasticity, while maintaining effective gradient prop- agation enhances label plasticity. From these insights, we proposed PLASTIC, an algorithm that synergistically combines the SAM optimizer, LN, periodic Reset, and CReLU activation. Empirically, this combination substantially improves both forms of plasticity. Demonstrating robust performance on benchmarks including Atari- 100k, PLASTIC offers a promising avenue for advancing sample efficiency in RL.
Nevertheless, our study has certain constraints. Our empirical evaluations were primarily limited to the Atari and DMC environments. A compelling direction for future work would be to assess PLASTIC's is efficacy in more intricate settings, such MetaWorke [79], Procgen [15]. These environments introduce greater non-stationarities, challenging the model's 's adaptability.
While the smoothness of loss surfaces and gradient propagation are fundamental to our findings, they might not capture the full complexity of the "loss of plasticity" phenomenon. Notably, even though SAM was sufficient to find the smoothest loss surfaces, Layer Normalization's 's integration amplifies its efficacy. Similarly, while CReLU inherently mitigates the reduction of active units, introducing periodic resets offers pronounced enhancements. We hypothesize that the nuanced attributes like the smoothness of loss landscapes and the number of active units might only scratch the surface of deeper determinants influencing model plasticity. Therefore, in-depth understanding and measurement of network plasticity can pave the way for more foundational solutions.
In conclusion, while we acknowledge the constraints, our research offers practioncal insights for amplifying the plasticity of RI. agents We hope our findings open up diverse possibility for future exploration, potentially leading to more sample-efficient and adaptable algorithms in RL.

---

Acknowledgments and Disclosure of Funding
This work was supported by the Institute of Information Communications Technology Planning & Evaluation (ITPP) grant funded by the Korean government (MSIT) (No.2019-0-00075, Artifi- cial Intelligence Graduate School Program(KAIST). HL, HK, and JC acknowledge support from the National Supercomputing Center with supercomputing resources including technical support - HC and CY acknowledge support from the Institute of Information & communications Technology Planning & Evaluation (ITTP) grant (No. 2022-0-00188 144, Development and Study of Technologies to Inexpensively Conform to Evolving Policy on Ethics) and the National Research Foundation of Korea (NRF) grant (No. RS-2023-0021 1352), both funded by the
Korea government (MSIT).
References
[1] Momin Abbas, Quan Xiao, Lisha Chen, Pin-Yu Chen, and Tianyi Chen. Sharp-maml: Sharpness- aware model-agnostic meta learning. In Proc. the International Conference on Machine Learning (ICML.), 2022. 19
[2] Zaheer Abbas, Rosie Zhao, Joseph Modayil, Adam White, and Marlos C Machado. Loss of plasticity in continual deep reinforcement learning. arXiv preprint arXiv:2303.07507, 2023. 2,
[3] Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Mare Belle- mare. Deep reinforcement learning at the edge of the statistical precipice. Proc. the Advances in Neural Information Processing Systems (NeurIPS), 34:29304-29320, 2021. 6
[4] Jordan Ash and Ryan P Adams. On warm-starting neural network training. Proc. the Advances
in Neural Information Processing Systems (NeurIPS), 389884-8894, 2020. 7
[5] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint
167.055607.06450, 2016. 2
[6] Dara Bahri, Hossein Mobahi, and Yi Tay. Sharpness-aware minimization improves language model generalization. In Proc. the Annual Meeting of the Association for Computational Linguistics (ACL), pages 7360-7371, 2022. 19
[7] Bowen Baker, Ilge Akkaya, Peter Zhokhov, Joost Huizinga, Jie Tang, Adrien Ecoffet, Brandon Houghton, Raul Sampedro, and Jeff Clune. Video pretraining (vpt): Learning to act by watching unlabeled online videos. Proc. the Advances in Neural Information Processing Systems (NeurIPS), 2022. 8
[8] Mare G Bellemare, Will Dabney, and Rémi Munos. distributional perspective on reinforce- ment learning, In Proc. the International Conference on Machine Learning (ICML), pages 449 458. PMLR, 2017. 3
[9] Marc G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47:253-279, 2013. 5
[10] David Bertoin and Emmanuel Rachelson. Local feature swapping for generalization rein- forcement learning. Proc. the International Conference on Learning Representations (ICLR), 2022.
[11] Nils Bjorck, Carla Gomes, and Kilian Q Weinberger. Towards deeper deep reinforcement learning with spectral normalization. Proc. the Advances in Neural Information Processing Systems (NeurIPS), 34:8242-825 8255, 2022.
[12] Edoardo Cetin, Philip Ball, Steve Roberts, and Oya Celiktutan. Stabilizing off-policy deep 2022.
reinforcement learning from pixels. arXiv preprint
[13] Xiangning Chen, Cho-Jui Hsieh, and Boging Gong. When vision transformers outperform resnets without pre-training or strong data augmentations. Proc. the International Conference on Learning Representations (ICLR), 2022. 3
10

---

[14] Xiangning Chen, Cho-Jui Hsieh, and Boging Gong. When vision transformers outperform resnets without pre-training or strong data augmentations. In International Conference on Learning Representations, 2022. 19
[15] Karl Cobbe, Chris Hesse, Jacob Hilton, and John Schulman. Leveraging procedural generation to benchmark reinforcement learning. In Proc. the International Conference on Machine Learning (ICML), pages 1048-2056, PMLR, 2020. 9
[16] Terrance DeVries and Graham W Taylor. Improved regularization of convolutional neural networks with cutout. arXiv preprint arXiv: 04552, 2017. 3
[17] Shibhansh Dohare, Juan Srmmandez-Garciaa Parash Rahman, Richard Sutton, and A Rupam
Mahmood. Loss of plasticity in deep continual learning. 2023. 3
[18] Shibhansh Dohare, Richard S Sutton, and A Rupam Mahmood. Continual backprop: Stochastic gradient descent with persistent randomness. arXiv preprint arXiv:2108.06.3 06325, 2021. 2, 3, 6
[19] Pierluca D'Oro, Max Schwarzer, Evgenii Nikishin, Pierre-Luc Bacon, Mare G Bellemare, and Aaron Courville. Sample-efficient reinforcement learning by breaking the replay ratio barrier. In Proc. the International Conference Learning Representations (ICLR), 2023. 1,3, 7, 8, 16
[20] Jiameng Fan and Wenchao Li. Dribo: Robust deep reinforcement learning via multi-view information bottleneck. In Proc. the International Conference on Machine Learning (ICML), pages 6074-6102. PMLR, 2022. 3
[21] Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware min- imization for efficiently improving generalization. In Proc. the International Conference Learning Representations (ICLR), 2021. 2, 6, 19
[22] Meire Fortunato, Mohammad Gheshlaghi Azar, Bilal Piot, Jacob Menick, Ian Osband, Alex Graves, Vlad Mnih, Remi Munos, Demis Hassabis, Olivier Pietquin, al. Noisy networks for exploration. arXiv preprint arXiv: 10295, 2017. 3,4
[23] Florin Gogianu, Tudor Berariu, Mihaela C Rosca, Claudia Clopath, Lucian Busoniu, and Razvan Pascanu. Spectral normalisation for deep reinforcement learning: an optimisation perspective. In International Conference on Machine Learning, pages 3734-3744. PMLR, 2021.
[24] Tuomas Haarnoja, Aurick Zhou, Picter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In Proc. the International Conference on Machine Learning (ICML), pages 1861-1870. PMLR, 2018. 1, 3, 4,
[25] Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint 2r3i 004 104, 2023. 8
[26] Matteo Hessel, Joseph Modayil, Hado Van Hasselt, Tom Schaul, Georg Ostrovski, Will Dabney, Dan Horgan, Bilal Piot, Mohammad Azar, and David Silver. Rainbow: Combining improve- ments in deep reinforcement learning. In Proc. the AAAI Conference on Artificial Intelligence (AAAI), 2018 3, 3,4,6 6, 16
[27] Takuya Hiraoka, Takahisa Imagawa, Taisei Hashimoto, Takashi Onishi, and Yoshimasa Tsu- ruoka. Dropout q-functions for doubly efficient reinforcement learning. Proc. the International Conference on Learning Representations (ICLR), 2022. 3
[28] Minyoung Huh, Hossein Mobahi, Richard Zhang, Brian Cheung, Pulkit Agrawal, and Phillip Isola. The low-rank simplicity bias in deep networks. 2023.
[29] Maximilian Igl. Gregory Farquhar, Jelena Luketina, Wendelin Bochmer, and Shimon White- son. Transient non-stationarity and generalisation in deep reinforcement learning. Proc. the International Conference on Learning Representations (ICLR), 2021. 3
[30] Pavel Izmailov, Dmitrii Podoprikhin, T. Garipov, Dmitry P. Vetrov, and Andrew Gordon Wilson. Averaging weights leads to wider optima and better generalization. In Proc. the Conference on Uncertainty in Artificial Intelligence (UAI), 2018. 3. 19

---

[31] Jean Kaddour, Linqing Liu, Ricardo Silva, and Matt J Kusner. When do flat minima optimizers work? Proc. the Advances Neural Information Processing Systems (NeurIPS), 35:16577-
16595, 2022. 19
[32] Lukasz Kaiser, Mohammad Babacizadeh, Piotr Milos, Blažej Osiński, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, et al. Model based reinforcement learning for atari. In Proc. the International Conference on Learning Representations (ICLR), 2019.
[33] Simran Kaur, Jeremy Cohen, and Zachary Chase Lipton. On the maximum hessian eigenvalue and generalization. In Can't Believe It's Not Better Workshop: Understanding Deep Learning Through Empirical Falsification, 2023.
[34] Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On burge-batch training for deep learning: Generalization gap and sharp minima. In Proc. the International Conference on Learning Representations (ICI.R) 2017. 3, 19
[35] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. Proc. the International Conference on Learning Representations (ICLR), 2015. 19
[36] Jens Kober, J Andrew Bagnell, and Jan Peters. Reinforcement learning in robotics: A survey. The International Journal of Robotics Research, 2013. 1,3
[37] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009. 2,4 4
[38] Alex Krizhevsky, Geoffire Hinton, et al. Learning multiple layers of features from tiny images. 2009. 16
[39] Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Curl: Contrastive unsupervised represen- tations for reinforcement learning. Proc. the International Conference on Machine Learning (ICML), 2020. 3
[40] Misha Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Reinforcement learning with augmented data. Proc. the Advances in Neural Information Processing Systems (NeurIPS), 2020.
[41] Hojoon Lee, Koanho Lee, Dongyoon Hwang, Hyunho Lee, Byungkun Lee, and Jaegul Choo. On the importance of feature decorrelation for unsupervised representation learning in reinforcement learning. arXiv preprint 2=XXi:2300605637, 2023. 3, 8
[42] Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.0/6433 2020. 1,
3
[43] Qiyang Li, Aviral Kumar, Ilya Kostrikov, and Sergey Levine. Efficient deep reinforcement learning requires regulating overfitting. arXiv preprint arXiv:2304. 10466, 2023. 1,3
[44] Yong Liu, Siqi Mai, Xiangning Chen, Cho-Jui Hsieh, and Yang You. Towards efficient and scalable sharpness-awar minimization. In Proc. of the IEEE conference on computer vision and pattern recognition (CVPR), pages 12360-12370, 2022. 3
[45] Zhuang Liu, Xuanlin Li, Bingyi Kang, and Trevor Darrell. Regularization matters in policy optimization-an an empirical study on continuous control. Proc. the International Conference on Learning Representations (ICLR), 2021.
[46] Clare Lyle, Mark Rowland, and Will Dabney. Understanding and preventing capacity loss in reinforcement learning. Proc. the International Conference on Learning Representations (ICI.R), 2022.
[47] Clare Lyle, Zeyu Zheng, Evgenii Nikishin, Bernardo Avila Pires, Razvan Pascanu, and Will Dabney. Understanding plasticity neural networks. Proc. the International Conference on Machine Learning (ICML), 2023. 3,6
12

---

[48] Bogdan Mazoure, Remi Tachet des Combes, Thang Long Doan, Philip Bachman, and R Devon Hjelm. Deep reinforcement and infomax learning. Advances in Neural Information Processing Systems, 33:3686-3666 -3998, 2020. 3
[49] Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers sample efficient world models. arXiv preprini N0::2209900088. 2022. 8
[50] Volodymyr Mnih, Koray Kavukeuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. Nature, 2015. 3, 4, 16
[51] Clara Na, Sanket Vaibhav Mehta, and Emma Strubell. Train flat, then compress: Sharpness- aware minimization learns more compressible models. In Proc. of the Conference on Empirical Methods Natural Language Processing (EMNLP), pages 4909 4936, 2022.
[52] Evgenii Nikishin, Max Schwarzer, Pierluca D' O'Oro, Pierre-Luc Bacon, and Aaron Courville. The primacy bias in deep reinforcement learning. In Proc. the International Conference Machine Learning (ICML), pages 16828-16847. PMLR, 2022. 1, 2, 3, 6, 16, 17, 18
[53] Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint 1r0 007448 2018. 3
[54] Boris T Polyak. Some methods of speeding up the convergence of iteration methods. Ussr computational mathematics and mathematical physics, 4(5):1-17, 1964. 4, 19
[55] Roberta Raileanu, Max Goldstein, Denis Yarats, Ilya Kostrikov, and Rob Fergus. Auto- matic data augmentation for generalization in deep reinforcement learning. arXiv preprint arXiv:2006. 2262, 2020. 3
[56) Herbert Robbins and Sutton Monro. A Stochastic Approximation Method. The Annals of Mathematical Statistics, 22(3):400 407, 407, 1951. 4, 19
[57] Tom Schaul, John Quan, loannis Antonoglou, and David Silver. Prioritized experience replay. Proc. the International Conference Learning Representations (ICLR), 2016. 3
[58] Max Schwarzer, Ankesh Anand, Rishab Goel, Devon Hjelm, Aaron Courville, and Philip Bachman. Data-efficient reinforcement learning with self- predictieee representations. In Proc. the International Conference on Learning Representations (ICLR), 2020. 3,5
[59] Max Schwarzer, Nitarshan Rajkumar, Michael Noukhovitch, Ankesh Anand, Laurent Charlin, Devon Hjelm, Philip Bachman, and Aaron C Courville. Pretraining representations for data-efficient reinforcement learning. Proc. the Advances in Neural Information Processing Systems (NeurIPS), 2021.
[60] Younggyo Seo, Kimin Lee, Stephen I. James, and Pieter Abbeel. Reinforcement learning with action-free pre-training from videos. In Proc. the International Conference Machine Learning (ICML). pages 19561-19579. PMLR, 2022. 8
[61] Wenling Shang, Kihyuk Sohn, Diogo Almeida, and Honglak Lee. Understanding and improving convolutional neural networks via concatenated rectified linear units. In Proc. the International Conference Machine Learning (ICML), pages 22 221722225. PMLR, 2016. 2,
[62] Dongkuk Si and Chulhee Yun. Practical sharpness aware minimization cannot converge all the way to optima. Proc. the Advances in Neural Information Processing Systems (NeurIPS), 2023. 19
[63] Ghada Sokar, Rishabh Agarwal, Pablo Samuel Castro, and Utku Evci. The dormant neuron phenomenon in deep reinforcement learning. Proc. the International Conference Machine Learning (ICML), 2023.
[64] Adam Stooke, Kimin Lee, Pieter Abbeel, and Michael Laskin. Decoupling representation learning from reinforcement learning. In Proc. the International Conference on Machine Learning (ICML), 2021.
13

---

[65] Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. arXiv preprini arXiv: 1801000690 2018. 5
[66] Hado Van Hasselt, Arthur Guez, and David Silver. Deep reinforcement learning with double q-learning. In Proc. the AAAI Conference on Artificial Intelligence (AAAT), volume 30, 2016.
[67] Hado P Van Hasselt, Matteo Hessel, and John Aslanides. When to use parametric models in reinforcement learning? Proc. the Advances in Neural Information Processing Systems (NeurlPS), 32, 2019. 16
[68] Kaixin Wang, Bingyi Kang, Jie Shao, and Jiashi Feng. Improving generalization in reinforce- ment learning with mixture regularization. Proc. the Advances Neural Information Processing Systems (NeurIPS), 33:7968-7999 7978, 2020. 3
[69] Ziyu Wang, Tom Schaul, Matteo Hessel, Hado Hasselt, Marc Lanctot, and Nando Freitas. Dueling network architectures for deep reinforcement learning. In Proc. the International Conference Machine Leurning (ICML), pages 1995-2003. PMLR, 2016. 3
[70] Kaiyue Wen, Tengyu Ma, and Zhiyuan Li. How sharpness-aware minimization minimizes sharpness? In Proc. the International Conference on Learning Representations (ICLR), 2023. 19
[71] Dongxian Wu, Shu-Tao Xia, and Yisen Wang. Adversarial weight perturbation helps robust generalization. Advances in Neural Information Processing Systems, 11-2958-2969) 2020. 3,
19
[72] Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous control: Improved data tata-augmented reinforcement learning. arXiv preprint 1077v:2107.09645 2021. 3, 22
[73] Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. In Proc. the International Conference on Machine Learning (ICML), 2020. 3, 4, 5, 6, 16, 22
[74] Denis Yarats, Amy Zhang, Ilya Kostrikov, Brandon Amos, Joelle Pineau, and Rob Fergus. Improving sample efficiency in model-free reinforcement learning from images. In Proc. the AAAI Conference on Artificial Intelligence (AAAI), volume 35, pages 10674-10681, 2021.
[75] Weirui Ye, Shaohuai Liu, Thanard Kurutach, Pieter Abbeel, and Yang Gao. Mastering atari games with limited data. Proc. the Advances in Neural Information Processing Systems (NeurlPS), 34:25476-25488, 2021. 8
[76] Tao Yu, Cuiling Lan, Wenjun Zeng, Mingxiao Feng, Zhizheng Zhang, and Zhibo Chen. Playvir- tual: Augmenting sycle-consistent virtual trajectories for reinforcement learning. Proc. the Advances in Neural Information Processing Systems (NeurIPS), 34:5276-5221 5289, 2021. 7,
[77] Tao Yu, Zhizheng Zhang, Cuiling Lan, Yan Lu, and Zhibo Chen. Mask-based latent reconstruc- tion for reinforcement learning. Proc. the Advances in Neural Information Processing Systems (NeurIPS), 35:25117-25131, 2022. 3
[78] Tao Yu, Zhizheng Zhang, Cuiling Lan, Yan Lu, and Zhibo Chen. Mask-based latent reconstrue- tion for reinforcement learning. Proc. the Advances in Neural Information Processing Systems (NeurIPS), 35:25117-25131,, 2022. 7.
[79] Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on robol learning, pages 1094-1100. PMLR, 2020. 9
[80] Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz. mixup: Beyond 2017.
empirical risk minimization. arXiv preprint arXiv:
[81] Hattie Zhou, Ankit Vani, Hugo Larochelle, and Aaron Courville. Fortuitous forgetting in connectionist networks. Proc. the International Conference on Learning Representations (ICLR), 2022. 2
14

---

Appendix
A Implementation Details
16
A.I Computation
16
A.2 Synthetic Experiment
16
A.3 Arcade Learning Environment (Atari-100k)
16
A.4 Deepmind Control Suite Medium (DMC-M)
18
B Applying SAM to Deep RL
19
B.1 Preliminaries and Backgrounds SAM
19
B.2 Deep Q-Learning with SAM and Noisy Layers
20
B.3 Soft Actor-Critic with SAM
B.4
Ablations the usage of SAM
22
Per Environment Results
23
C.1 Atari-100k
23
C.2 DMC-M
23
Broader Impact
26
15

---

A Implementation Details
A.1
Computation
For all experiments, we used NVIDIA RTX 3090 GPU for neural network training and 32- core AMD EPYC 7502 for multi-threaded tasks, accelerating machine learning model training and inference.
For the Atari 100k benchmark, our computations were based on the DrQ algorithm with SAM (Sharpness-Aware Minimization) and Reset. We ran 3 experiments in parallel on single GPU where PLASTIC and PLASTIC! took and 1.5 hours, which effectively translates to an average of I and 0.5 hours per run.
For the DeepMind Control Suite (DMC), runtimes varied based the individual environment due the differences in action repeat (i.e.. the number of the same actions taken for each environmental interaction). On our experimental setup, training took roughly 8 hours to complete 2 million environment steps with an action repeat of 4, and around 15 hours with an action repeat of 2.
A.2 Synthetic Experiment
In our synthetic experiment, we utilized the CIFAR-10 dataset [38] to test both input adaptation and label adaptation scenarios. The dataset comprises 50,000 training images and 10,000 test images, evenly distributed across 10 classes. Each image in the dataset is color image of size 32x32.
the input adaptation experiment, the training data was segmented into 100 chunks, each consisting of 500 images. These chunks were sequentially added to a buffer, which progressively expanded during the course of the experiment. For every model update, data were randomly sampled from this progressively growing buffer.
Input Adaptation:
Label Adaptation: In the label adaptation experiment, label alterations were done 100 times during the training phase. The process of label reassignment was performed in a uniform manner across all classes. For example, if the original class labels ranged from 0 to 9, after each label alteration, all data points within a specific class would be consistently reassigned to new class within the same range.
We selected a model architecture that reflects common architectural designs utilized in reinforcement learning research [50, 26]. The architecture consists of three convolutional layers for the backbone and three fully-con colleceedd layers for the head. The three convolutional layers employed kernel sizes of 3x3, with strides of (3, 3, 1), and the number of channels was set as (32, 64, 64) respectively. Subsequent to this, the three fully-connected layers consisted of hidden dimensions of (512, 128, 10).
The model was trained for total of 50,000 updates, with a batch size of 128 for each update (i.e., equivalent to training 128 epochs). The training was conducted using Stochastic Gradient Descent (SGD) with momentum. To select the learning rates and weight decay rates, we vary the values range from (0.1, 0.01, 0.001, 0.0001, 0.00001). integrating weight decay to minimize variance across individual runs.
In the synthetic experiment, we utilized four distinct methods: Layer Normalization (LN), Sharpness- Aware Minimization, Concatenated Rectified Linear Unit (CReLU), and Reset. For the SAM optimizer, we vary the perturbation parameter across the set {0.1, 0.03, 0.01} and for the Reset method, we search the reset interval over (5, 10, 20) data chunks.
Finally, we reported the test accuracy, averaged over 30 random seeds with a 95% confidence interval, the end of each experiment.
A.3 Arcade Learning Environment (Atari-100k)
We implemented the Rainbow algorithm [26] while following design choices from DER [67] and DrQ [73]. Furthermore, we integrated modifications suggested by D' D'Oro et al [19], including the application momentum encoder for the target network and the use of this target network for exploration. While these adjustments had little impact under a low replay ratio regime (c=2), they played a significant role when scaling the replay ratio with the reset mechanism [52]. A detailed hyperparameter is described in Table 5.
16

---

In the case of LN, apply layer normalization after each convolutional and fully- colne-co layer. For CReLU, we replaced all ReLU activations with concatenated ReLU where the imput dimensions are doubled for each subsequent layers. For the methods using the Reset mechanism, we reset the head layer's parameters every 40,000 gradient updates, in accordance with the original paper [52].
Table 5: Hyperparameters on Atari 100k. The ones introduced by this work are at the bottom.
Hyperparameter
Value
State downsample size
(84, 84) True
Grey scaling
Data augmentation
Random Shifts and Intensity Jittering
Random Shifts
+ 4 pixels
Intensity Jiterring scale
0.05
Frame skip
4
Stacked frames
Action repeat
4
Training steps
100k
Update
Distributional Q
Dueling
True
Support of Q-distribution
51
Discount factor
0.99
Batch size
32
Optimizer (81, .2.
Adam (0.9, 0.999, 0.000015)
Learning rate
0.0001
Max gradient
10
Priority exponent
0.5
Priority correction
0.4 -> 1
EMA coefficient (r)
0.99
Exploration Network
Target
Exploration
Noisy nets
Noisy nets parameter
0.5
Replay buffer size Min buffer size for sampling
100k
2000
Replay ratio
2
Multi-step return length
10
Q-head hidden units
512
Q-head non-linearity
ReLU
Evaluation trajectories
100
SAM parameter
0.1
Reset interval
40, 000
17

---

A.4 Deepmind Control Suite Medium (DMC-M)
We utilized an open-source JAX implementation provided by Nikishin et al. [52] as the foundation for our work. We integrated the SAM optimizer implementation into this existing framework. To ensure consistency and comparability with prior studies, we strictly followed the architecture and hyperparameters outlined by Nikishin et al. [52], documented in Table 6.
In the case of the Base, 1.2, and Reset models described in Table 6, adopted the results from Nikishin et [52]. For the LN model, we incorporated layer normalization after each convolutional and fully-connected layer. Additionally, replaced all ReLU activations with concatenated ReLU for the CReLU baseline.
We utilized the same training steps of 2,000,000 and reset interval of 100,000 for the DMC medium task as used by Nikishin et al [52].
Table 6: Hyperparameters on DMC. The ones introduced by this work are at the bottom.
Parameter
Sctting
Gray-scaling
Truc
Observation down-sampling
(64, 64)
Frame stacked
3
Discount factor
0.99
Minibatch size
512
Learning rate
0,0003
Backbone: channels
32, 64, 128, 256
Backbone: stride
2, 2, 2, 2 50
Backhone: latent dim
Head: n. hidden layers
2
Head: hidden units
256
Target network update period
1 0.995
EMA coefficient т
Initial Temperature
0.1
Updates per step
Replay Buffer Size
1, 000, 000
Total training steps Evaluation trajectories
2, 000, 000 100
Quadruped 0.1 Others 0.01 100, 000
SAM parameter p
Reset interval
18

---

B Applying SAM to Deep RL
B.1 Preliminaries and Backgrounds on SAM
The capability to make a correct prediction when given unseen inputs, which we refer to as input plasticity, has been recognized as an important issue across the entirety of machine learning research, and numerous studies have been conducted to address it. Keskar et al. [34] propose a sharpness measure, arguing that finding wide and smooth minima is beneficial in terms of generalizability (i.e., performance on unseen datasets). Stochastic Weight Averaging (SWA) finds a flat minimum by averaging checkpoints in the model's training trajectory [30]. The work by [71] incorporates the addition of adversarial weight perturbation over the course of training. These works are based on the hypothesis by [34] that wide and flat minima result in a more plastic model under change of inputs.
Following a similar philosophy, Foret et al. [21] devised sharpness-aware minimization (SAM), which aims reduce not only the value of the training loss but also its sharpness. By reducing the sharpness of the loss and encouraging the model to get closer to flat minimum, SAM has been successfully applied in various settings including computer vision [14], natural language processing [6], meta- wheaa-eerningg [1], and model compression [51].
The aforementioned studies concern training in stationary data distributions. In contrast, our work involves experiments in reinforcement learning which has non- statioottonay data distribution in which the efficacy of SAM has not yet been verified. Indeed, empirically verified that the usage of SAM enhances the adaptability of deep RI. agents.
Now, we elaborate on what SAM is. Originally, SAM aims to train model which is robust to adversarially perturbed model weights, by solving the following bi-level optimization problem on the loss function L:
min {maxe:llll,sp C(w + +)]}
(5)
This problem implies minimization of loss value around an balball of radius p. which encourages a model parameter ur to find a smoother region of the loss landscape. Here, p 0 is hyperparameter, called SAM parameter, restricting the size of perturbation to make the bi-level problem (5) feasible. Since the exact solution of the inner maximization problem in (5) is not tractable yet, first-order Taylor approximation is applied as a remedy for finding an approximately optimal perturbation at w,
<,(w) == IV ((wl2
V C(w).
We call the step of computing (u)) as the perturbation step of SAM. As result of the perturbation step, we can approximately solve the problem (5) with any gradient update algorithm (i.e., base optimizer), such as Stochastic Gradient Descent (SGD) [56]. SGD with momentum [54], Adam [35], using the SAM gradient at current weight w,: v. w: + (w:)). We call the step of updating the weights with SAM gradient as the update step of SAM. In essence, SAM is an iterative algorithm alternating perturbation and update steps. Readers can check a more detailed derivation of the SAM gradient in [21].
It is both theoretically and empirically proven that such an optimization scheme implicitly prefers the weights in smoother (i.e., less sharp) region of loss landscape, in terms of, e.g., the maximum eigenvalue and/or trace of Hessian [21, 31, 70, 62].
Naturally, in (deep) RL,, we never have access to the full static loss function as described above. Hence, analogous most of the applications of SAM, it natural randomly sample a mini-batch B, of m transitions from replay buffer and apply SAM update (often called m-SAM) using the (wz)). where
stochastic gradient Wg, (w, + (,,,(ww))
(,)(1) = WwLC(w)ll VlLfs(r)
P
There are several considerations for applying SAM to RL agents. When applying SAM to the Rainbow agent, the random noises from noisy layers embedded the agent might hurt the SAM perturbation. However, we observe that the effect of regulating the noises is not significant. Moreover, there are options to apply SAM perturbation solely to the backbone or head. We find that applying SAM to the whole network is the most beneficial for enhancing generalization capability, yet
19

---

observe that solely applying SAM to the backbone is quite similar to the case of applying it to the whole. On the other hand, when applying SAM the SAC agent, there are multiple loss functions to be optimized, which share parameters. Again, applying SAM to both actor and critic is the most desirable option, yet we observe that the application critic is more critical for performance improvement. More detailed ablation studies appear in Section B.4.
B.2 Deep Q-Learning with SAM and Noisy Layers
We used the Rainbow agent and its variants learn discrete control problems of the Atari- benchmark. Although there are many other components inside the Rainbow agent, such as multi-step learning, distributional RL, and dueling Deep Q-Network (DQN), would be confusing integrate all details together in the pseudocode. Thus, for simplicity of the display, we provide pseudocode of applying SAM to vanilla DQN with noisy layers. See Algorithm 1.
Algorithm 1: Deep Q-Learning with Noisy Layers & SAM
Input: Learning rate SAM parameter Replay ratio R; Discount factor y; Polyak averaging factor :
Initialize: Replay memory D; Q-function weight w; Target weight +- w
3 foreach environment step / 1= 1,,,.. do
Collect trajectory (s, a. r, s', with e-greedy (d 1[s' is terminal))
Store the trajectory to D If d Reset environment state;
foreach optimization step i E {I...., .T....,R). R] do
{(s, a, r,,, d)} from D Compute TD target y(r, s', = 7(1 d) max Qu (s', aa);;
Sample minibatch of transitions
// Perturbation step Sample random noise & of noisy layers inside Q-function;
10
E;
11
Compute gradient g(w) = V
Σ
(u(r, s', Qw
(s,0, d)E8
12
Compute SAM perturbation w SAM +
P
g(w)
[g(w)
// Update step
13
Sample random noise g' of noisy layers inside Q-function;
S
SAM
SSA + E::
Compute SAM gradient g SAM = g (1 SAM
14
15
Gradient descent update w t- w ng // Could be modified as any gradient-bar based optimizer
16
Update target weight Tw + (1 T)w;
20

---

B.3 Soft Actor-Critic with SAM
We also provide pseudocode for applying SAM to the SAC algorithm. See Algorithm
Algorithm 2: Soft Actor-Critic with SAM
Input:
Learning rates IQQ 7m Tlo SAM parameters PQ; Pr. Pa; Replay ratio R; discount factor y; Polyak averaging factor minimum expected entropy Md;
Initialize: Replay memory D; Q-function weights w1, w2; Q-function weights w w22 Policy weight 0; Entropy regularization coefficient a
foreach environment step 1.2.... do
Observe state and sample/execute an action ~ [o(-|s); Observe reward r, next state s', and done signal 1[s is terminal];
Store the trajectory (s, r, 8', d) to D:
If d 1, Reset environment state; foreach optimization step i {I..... , R] do
Sample minibatch of transitions {(s, a, r,,, from D;
// Perturbation step Compute target
10
y(r, s', d: 0) = + 7(1-
111-d d)
min Qw; (s', log mo (d')')) (d ~ To(-1s));
1=1,2
Compute Q-function gradient:
(q.( V w:
Σ
(y(r, s', d: (w, (i 1,2)
(sa,ras
Compute policy gradient by sampling differentiable action äه (s To (11)):
9x (w) V Σ
(s, Qw. 1=1,2
min Qw.
log no (āo((s)sss)
13
Compute temperature gradient: (a) = V ELEB (=alogmo( a?);
14
SAM-perturbations:
w
SAM
wi +
PQ
(w:(;:;; (i = 1,2)
11
@SAM = 0 +
Pr
9 (0);
(0) l1
ā SAM a +
Pa
Dromp (a);
9 Itempp (a:)ll2
// Update step Compute target y(r, s' d; @SAM); Compute SAM gradients:
15
16
SAM = 9Q.( 9Q.1
SAM gz (ŌSAM),
SAM
gtemp (ā SAM))
g
17
Gradient descent updates:
SAM, 0 0 n1 wi w. nog
SAM, a a a 7a Ta
SAM,
SAM.
// Could be modified as any gradient- optimizer Update target weight w, + Tw, + -1-77 7
18
21

---

B.4 Ablations on the usage of SAM
We conduct some ablation studies aiming to answer the following questions:
1. How should we deal with noisy layers during training?
2. What if we solely apply SAM perturbation to the backbone or head of Rainbow agent?
3. What if we solely apply SAM perturbation to the actor or critic of SAC agent?
The first two questions are answered by experiments with the DrQ algorithm tested on the Atari-1 benchmark, whereas for the last question we test the DMC-M benchmark with the SAC agent.
Abli Neloy Luyer on/off
backbone choad
Which networks SAM
SAM
0.35
0.55
0.50
Rexet
0.50
aade HHSAM
Sdin
0.00
0.33
0000
0.00
No PuMt
Cnty Adtor Only OnSc
$42
(a) Noisy layers on/off.
(b) Backbonc/Head SAM.
(c) Actor//ritic SAM.
Figure 5: Ablation study on utilizing SAM.
Noisy Layers. For the case of the Rainbow agent, we compare three different schemes of SAM updates regarding random noises from noisy layers: (i) perturbation with independent noise, using independently sampled (e.g.. different) noise between perturbation and update step (default setting); (ii) update with reused noise, turning the noise on at the perturbation step and reusing that noise at the update step: (iii) noise-less perturbation, turning the noise off at the perturbation step and turning on the update step. Despite the variations, the overall impact on the agent's 's performance was relatively consistent. Note that option (i) is the most naïve approach which can be implemented by simply adding SAM optimizer code block into any RL implementation with noisy layers.
Backbone-SAM v.s. Head-SAM. We experiment with resectioning the parts to be updated with SAM the RL agent architecture. We temporarily turn off the gradient of either backbone or head at the perturbation step and then turn on that gradient back at the update step. The result of the experiment, presented in Figure 5b, says that updating with backbone-only perturbation more beneficial than using head-only perturbation, although applying both is much better. Indeed, SAM perturbation of the head alone hurts the performance without any other SAM perturbation or resetting.
We want to explain why these happen from the point of view of the number of active neurons in the backbone/head. Interestingly, we observe that the SAM perturbation of backbone parameters induces more active units in the head, whereas the SAM perturbation of head parameters results in less active units in the backbone. Recall that both tendencies are observed when we apply SAM perturbation to both the backbone and head at once. From this, we can say that the rise of active units in the head by applying SAM is due to the SAM perturbation of the backbone, whereas the reason for the backbone getting sparser is the SAM perturbation of the head.
Actor-SAM v.s. Critic-SAM. SAC is an actor-critic algorithm that employs three distinct sets of updating parameters, namely the actor, critic, and alpha. Since alpha has just one parameter, focus on investigating the impact of SAM on actor and critic parameters. We conduct an ablation study utilizing the identical configuration demonstrates that exclusively applying SAM to critic parameters is more critical than solely applying SAM to actor parameters while using both is much better.
the DMC experiment presented in Section 4.2. Figure 5c
In the context of SAC, the actor and critic networks are constructed to jointly utilize the backbone layer. In order to ensure stable learning, updates to the backbone layer are exclusively driven by critic losses, while actor losses do not influence the backbone layer [73, 72]. Consequently, if SAM is exclusively applied to the actor parameters, the backbone layer remains unaffected by SAM as
22

---

there is no gradient propagation. As a result, the performance of the SAM solely used on the actor parameters, which has no impact on the backbone, is inferior to other approaches incorporating SAM, which corroborates the aforementioned findings.
0.60
Ablation: SAM param (o)
0.50
Ablation: Batch size
SAM
0.53
0.55
0.50
Resst
0.50
0.45
0.40
0.35
0.30
0.25
0.20
b6-16
05:32
(a) SAM parameter p.
(b) Batch sizes.
Figure 6: More ablation studies on using SAM.
SAM parameter (p). Being the main hyperparameter of SAM, the p > O represents the intensity we do not tune
of SAM perturbation. In every other experiment with DER and the value of but fix them as 0.1. Here. however, we run two other values of p's. 0.03 and 0.3, to examine the effect of p. The result shown in Figure 6a. We definitely observe the trade-off due to the size of If the p is too small, the scale of SAM perturbation becomes negligible and it cannot drastically improve the generalization capability compared to the case without SAM. However, larger value of is not always beneficial; the generalization capability does not always scale with the size of SAM perturbation. We, however, claim that finding an appropriate value of p is not too difficult; our recommendation is to test a few values between 0.01 and 1.
Batch size. We further study robustness with respect to batch sizes. We investigate the batch sizes among 16, 32 (baseline), and 64. From the resulting plot Figure 6b, we observe that SAM consistently improves the performance regardless of the choice of batch size, by comparing solid lines (daM+Adam) and dashed lines (Adam). We also observe that the performance improvement between reset and non-reset becomes much larger when SAM is applied.
C Per Environment Results
C.1 Atari-100k
C.2 DMC-M
23

---

DrQ algorithm. We report the mean trajectory 26 Atari games, evaluated on top of DrQ. For each random seed, the results are averaged over 100 different trajectories the end of training. The results are then averaged over 10 random seeds.
Table 7: Mean trajectory
the
Game
Random
Human
Base
SAM
LayerNorm
CRef.
Reset
PLASTIC
Alien
227.8
7127.7
757.62
854.84
809.14
101 1.84
926.73 73
1032.0
Amidar
5.8
1719.5
164.12
187.76
176.12
146.69
158. 22
201.6
Assult
222.4
742.0
582.50
597.9
506.26
585.57
694.17
858.5
Asterix
210.0
8503.3
823.3
780.6
889.3
887.8
921.15
1066.0
BankHeist
14.2
753.1
48.36
185.25
56.16
79.88
141.14
161.2
BattleZone
2360.0
37187.5
4088.00 0
13999.0
14416.0
11306.0
4700.0
2099.0
Boxing
0.1
12.1
11.56
16.43
21.2
18.73
41.14
44.5
Breakout
1.7
30.5
13.16
12.64
13.71
10.56
18.46
21.0
ChopperComaand
811.0
7387.8
778.0
997.7
879.2
1150.6
777.5
891.2
CraryClimher
10780.5
35829.4 4
13182.2
17927.7 7
15852.8
13957.6 6.
19082.8
31223.8
DemonAAtack
152.1
1971.0
7U1,
T35.54
556.33
531.05
1074.29
2117.8
Freeway
0.0
29.6
22.32
27.79
30.49
24.54
26.92
27.1
Prostbite
65.2
13317
1856,61
2052.95
1408.31
763.64
1257.33
1802.3
Gopher
257.6
2412.5
374.28
476.02
433.92
445.88
711.14
839.4
Hero
1027.0
3U826.4
5086.67
6196.56
5645.95
63322899
6987.43
7OUT
Jameshond
29.0
302.8
348.8
333.3
287.2
239.7
369.2
461.1
Kangaroo
52.0
3033500
3773.2
34538
3589.0
1919.8
2830.8
1696.1
Krull
1588.0
2665.5
3612.06
3385.05
3483.26
3710.24
4379.46
5019.5
Kusstr
258.5
22736.3
204122.4
16380.3
8616..
15188,4
12737.4
16105.0
MsPacman
307.3
6951.6
1210.02
1471.33
1226.5
1191.2
1362.05
1245.6
Pong
-20.7
14.6
8.82
-6.42
-12.85
1.96
10.0
17.77
Qbent
11.5
7845.0
3776.3
$495.45
2648-4
2429-2 22
3492.18
398663
ReadRunner
68.4
42051.7
15896.6
11919.5
13045..6
137033.4
129399.
150733.8
Seaquest
24.9
89571.3
525.88
519.2
394.24
383.8
645.26
635.9
Privatetiye
163.9
13455,0
100.0
19..0
100.0
99.4
100.13
100,0
UpNDown
533.4
11693.2
3690.74
4963.46
5506.2
4565.04
15463.98
66473.0
IQM
0.0
1.0
0,258
0.325
0,259
0,256
0.343
0.421
Median
0.0
1.0
0.277
0.327
U.247
0.193
0.291
0.347
Mean
0.0
1.0
0.176
0.501
0.463
0.498
0.660
0.933
1.0
0.0
0.589
0.628
0.379
0.538
Table 8: Mean trajectory scores on State-of-the-art methods. We report the individual scores the 26 Atari games. For IRIS, D3e EfficientZero, Virtual MLR, SR-SPR, and PLASTIC the results are averaged over 5, 5, 32, 15, 3, 10, and 5 seeds respectively.
Game
Efficien/erro
IRIS
DreamerVa
Virtual
MLR
SR-SPR-4
SR-SPR88
PLASTIC
PLASTIC:
Alien
808.5
420.0
1095.8
947.8
964.4
1015.5
1368.6
Amidar
148.6
143.0
142.9
165.3
227.7
211.8
2113.1
186.2
2201.2
Assault
1263,1
15214
638.0
702.3
613.7
987.5 3
1069.5
821.5
1074.
25557.8
653.6
982.8
BankHeist
351.0
53.1
617.8
245.9
180.1
460.0
472.3
370.1
313.8
BattleZome
13871.2
13074.0
12800.0
13260.0
16060.0
17800.6
19398.
15728.0
16868.0
Boxing
52.7
70.1
67.8
38.3
42.0
41.8
Breakoot
414.1
83.7
18.9
20.6
16.8
26.1
28.8
21.9
20.9
Chopper Command
1117.3
1565.0
100.0
922.4
910.7
1933.7
22011.0
798.4
871.2
CrazyClimber
83940.2
59321.2
71620.0
23176.7
21633.3
3831117
13122.3
31652.6
17893.0
DemonAttack
13003.9
2034.4
545.0
11:1.7
854.6
3016.2
2898.1
1562.9
2460.9
Freeway
21.8
31.1
0.0
16.1
30.2
24.5
24.9
29.4
30.4
Frosibite
296.3
259.1
1108.4
1984.7
2381_1
1809.9
1752.8
2152.5
2195.4
Gopher
3261.3
2236. 1
5828.6
684.3
822.3
717.5
711.2
582.7
612.6
9315.9
109611.
9000.7
Jameshond
517.0
462.7
510.0
394.7
423.2
408.8
392.5
401.3
428.
Kangaroo
721.1
838.2
$550.0
23817
8316.0
2021.1
3251.9
6218.2
22
Krull
3863.3
3923.1
5364.3
5201.9
KungFuMastee
31911.8
21739.8 8
291200.0
14259.0
10652.0 0
17656.5
17188.6
20839. 6
195
MsPacman
1281.2
999.1
1388.5
1315.4
1481.3
1544.7
1522.6
1662.0
1292.4
Pong
20.1
14.6
18.5
3.0
4.9
5.6
-3.0
Qhert
13781.9
745.7
31.77.0
3620.11
3410.4
3809.8
3850.6
4372.8
4967.2
RoadRanner
17751.3
9614.6
13429.4
14287.3
13623.5
20709.0
Seaquest
1100.2
661.3
582.0
532.9
625.3
766.6
800.5
571.7
850.1
Privatellye
96.7
100.0
1124.0
91.9
100.0
95.8
95.8
100.0
100.0
UpNDown
17261.2
IQM
1/4
0.501
0.497
0.374
0.432
0.541
0.589
0.583
0.571
Median
0.227
0,280
0.466
1//a
n/a
0.523
0.560
0.542
0.404
Mean
0.562
0.
OG
%/a
0.512
0.505
0.558
0.522
0.170
0.152
0.418
0. 16161
24

---

Table Mean trajectory scores for each value of replay ratio (2, 4, 8). We report the individual scores on the 26 Atari games. Here, the results averaged over random seeds.
Game
PLASTIC :2
PLASTIC:2
PLASTIC'
PLASTIC4
PLASTIC!S
PLASTIC:8
1021.7
Amidar
193.2
2016
151.4
206.1
186.2
220.2
521.5
Asterix
1007.6
1066.0
1098 1
1058.0
1181.
1188.0
Banklleist
217.7
BettleZone
17767.0
2099.0
17762.0
W2W2.0
15728.0
16868.1
Boxing
39.1
45.9
Breakout
16.4
21.0
22 4
27.9
21.9
21.9
ChomperCommand
972.4
891.2
875.1
799.1
871.2
CrazyClimber
21967.5
31223.5
24615.0
42367 6
31652.6
478933.0
DemonAttack
1016.2
1630.6
210/2.2
1562.9
2160.9
Freeway
30.1
27.1
29.6
299.6
29.4
301.4
Frostbite
1582.0
1802.3
2119.1
1899.3
2152.5
2195.4
Gopher
616.4
839.4
635.2
511.6
582.7
612.6
Hero
6582.8
7007.2
8118.2
11195.5
9000.77
Jamesbond
397.5
461.1
3883.7
438.7
401.3
428.8
Kangaroe
2836.3
1636.1
6800.8
2386.6
6218.2
22/9.2
Krull
4483.6
5.119.5
4719.8
52666 4
52019
564176
KunsteMaster
11873.9
176134
19828
20839.6
19546.2
MsPacman
1300.6
1245.6
1401.9
1457.2
1662.0
1292.4
Pong
-63
-17.7
0.3
33.5
Qbert
4006
3086.3
3794.1
4736.4
4372.8
4067.2
RoadRunner
15883.0
15073.8
1625-1.0
20709.0
Seaquest
676. 1
635.9
692.8
537.0
574.7
839.
PrivateEye
16.6
100.0
100,0
100.0
UpNDown
15470.0
66473.0
11399.0
51764.5
34342.4
33203.3
IQM
0.396
0.421
0.518
0.545
0.583
0.571
Medlaan
0.425
0.347
0.517
0.407
0.542
0.494
Mean
0.702
0,933
0.858
1,012
0,1888
0.939
OG
0.541
0.535
0.478
0.475
0.448
0.461
the Table 10: Mean trajectory DrQ algorithm. We report the mean trajectory
DMC medium environments, evaluated on top of DrQ. For each random seed, the results are averaged over 10 trajectories at the end of training. The results are then averaged over 10 random seeds.
Game
Base
SAM
LayerNorm
CReLU
Reset
PLASTIC
swrubot-swingup
16.34
15.36
23.45
40.79
57.
68.24
parppol-swinnup_spee
79.03
197.91
164.57
321.83
740.0
691.16
cheetah-run
727.44
736.61
591.47
835.33
665.67
708.04
finger-tur__ed _esy
207.02
129.92
325.21
221.0
232.76
379.69
finger-turn haard
80.95
60.07
485.70
117.5
89.11
546.56
hopper-hop
284.78
255.84
277.48
293.52
284.69
254.62
quadruped-rum
155.06
261.03
266.84
78.89
466.93
426.57
quadruped-walk
88.08
130.85
173.13
80.02
632.33
528.12
reacher-easy
505.08
738.80
926.94
914.3
957.96
879.71
reacher-haud
450.96
495.32
746.55
795.77
794.53
798.10
walker-run
632.93
639.94
652.41
643.24
564.95
632.55
IQM
278
338
565
Median
398
Mean
293
332
394
498
25

---

D Broader Impact
Our research possesses broader impacts in different aspects. The first impact lies in the sample efficiency of Reinforcement Learning (RL) algorithms. We incorporate principles of preventing the loss of plasticity, leading to more efficient and adaptive solutions. This advance can improve RL's sample efficiency, opening new avenues for future research in this domain.
The second impact is inclusivity and accessibility in RL. By decreasing data and computational demands, this work can aid underprivileged communities with fewer resources to participate in RI research and benefit from its many applications. This approach encourages diverse range of perspectives and experiences in the field, enriching the community.
However, while research has these positive implications, it is important to also acknowledge potential risks associated with RL technologies, particularly in robotics. We must continue to uphold ethical standards and prioritize safety to prevent misuse harm that could arise from these advancements.
26


Reviews:
[{"review_id": "Y0mS2xW4KF", "reviewer": "Reviewer_1mez", "summary": "Sample efficiency in RL is desirable to reduce computational and data collection costs, and is particularly critical to data-limited domains. \nWhile off-policy methods can improve sample efficiency by training multiple passes over the same data, it faces challenges due to overfitting, which makes it harder for the model to adapt to new data.\nThis paper argues that to address the problem, it is important to tackle both the generalization and plasticity of the model, and proposes a method that achieves this.\nThe method uses Sharpness-Aware Minimization (SAM) to improve the model’s generalization, and a reset mechanism that periodically reinitializes the final layers of the model to inject plasticity.\nUsing a synthetic supervised learning experiment, the authors show that SAM helps the model better adapt to new inputs, and the reset mechanism helps it to adapt better to new labels.\nThese improvements in adaptability enables off-policy methods to better utilize multiple updates on data, improving sample efficiency.\n\nThe authors evaluate their method on Atari 100k and DMC-M, applying them to the DER and DrQ learning algorithms that are designed for sample efficient learning. \nThe results show that both DER and DrQ gain significant performance improvements when equipped with SAM + reset, compared to other methods that aim to improve generalization or plasticity alone. \nThe authors further perform ablation studies. \nNotably, they find that SAM + CRelu performs almost as well as SAM + reset, which suggest that it is the combination of improving generalization and adaptability that leads to the performance gains, rather than specific synergies between SAM and reset.\nThis highlights that the value of their contribution, which does not lie in merely combining two known techniques, but in discovering a synergistic dynamic between generalization and plasticity for sample efficient RL.", "questions": "- Have you tested other combinations of methods to improve generalization and plasticity? For me, the main contribution is not SAM + reset, but the insight that generalizability and plasticity play separate roles in improving adaptability, and combining methods that improve each of them can lead to further performance gains; SAM + reset is just one instantiation used to verify your hypothesis. Ideally, I'd like to see that the same pattern holds on other combinations beyond SAM + reset and SAM + CRelu.\n- Why did you select DER and DrQ as your base learning algorithms? Have you experimented with applying your method to stronger-performing algorithms, and would the same performance gains you achieved on DER and DrQ also hold there?\n", "limitations": "The authors are forthcoming with the main limitations of this work, and explained them clearly.", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- The writing is clear, and figures are intuitive and well-presented. The authors' claims are modulated and supported by the experiments. The logic of the paper flows very well; many questions that I had while reading were quickly addressed in subsequent sections.\n- The paper studies an important topic that has broad-ranging technical and environmental impacts. Sample efficient learning is not only an interesting technical problem, but also directly contributes to decreasing the environmental footprint of our field. In the age of increasingly large models trained with increasingly large amounts of data and compute, this is a critical issue.\n- The solution that the authors propose is simple yet elegant, and can be widely applied as it requires little change to model architecture. Furthermore, beyond showing strong empirical results for their approach, the authors also provide a hypothesis on how the underlying principles of generalization and plasticity achieve a synergistic dynamic, and support it with synthetic experiments.", "weaknesses": "- Both SAM and the reset mechanism are known techniques, and applying them in conjunction to achieve good results is not by itself sufficient for novelty. However, in the ablation studies, the authors also demonstrate that SAM + CRelu performs almost as well as SAM + reset. This strengthens their claim that it is the synergistic relationship between generalization and plasticity, not the specific method to improve these properties, that leads to the performance gains, which is novel. Yet, just testing one alternative combination does not seem sufficient to support the claim.\n- The experiments that study the impact of generalization and plasticity on the model's learning dynamics is synthetic and on supervised learning, which may not be an accurate approximation of the dynamics faced during RL. For example, uniform-randomly changing labels may not accurately simulate the moving targets problem faced in deep Q-learning.\n- The base learning algorithms that the authors evaluate with, DER and DrQ, are not state-of-the-art for data-limited RL. For example, EfficientZero reports achieving a normalized median of 1.090 and normalized mean of 1.943 on Atari 100k, which is better than the results shown in this paper.\n- Several fields appear to be missing in Table 1?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer 1mez,  \n\nWe appreciate your insightful questions and positive support. We have provided a detailed response to the comments which includes additional experiments on finding different synergetic combinations and integration of SAM + Reset on the advanced algorithms. Please let us know if you have any further comments or feedback. We will do our best to address them.  \n\n> **Question 1.1:**\nHave you tested other combinations of methods to improve generalization and plasticity? For me, the main contribution is the insight that generalizability and plasticity play separate roles in enhancing adaptability. \n\nThank you for highlighting our paper's key insight. To relieve the reviewer’s concern, we did investigate various method combinations. As seen in Figure A.1 and Table A.1 (attached pdf), LN(B) and SAM(B+H) consistently excel in input adaptation, while CReLU(H) and Reset(H) stand out for label adaptation.\n\nWhen combining these methods, we made some insightful observations in both synthetic and reinforcement learning experiments:\n\n- CReLU + Reset showed moderate improvements.\n- SAM + LN also yielded moderate improvements.\n- Yet, pairing any input adaptation method (either SAM or LN) with any label adaptation method (CReLU or Reset) consistently produced robust results.\n- A comprehensive combination of all these methods further enhanced performance.\nThese findings reinforce our notion of the distinct yet complementary roles of generalizability and plasticity in enhancing adaptability.\n\n> **Question 1.2:**\nFor synthetic experiments, uniform-randomly changing labels may not accurately simulate the moving targets problem faced in deep Q-learning.\n\nDue to the inherent complexities of learning dynamics in RL, we simplified our experiments using a supervised learning setup. We acknowledge your concerns regarding the alignment of our \"Label Adaptation\" scenario with RL's dynamics, especially in deep Q-learning.\n\nIn Q-learning, the agent undergoes two primary label adaptation scenarios: 1) the continual change of the best action as new data is received (akin to changing labels), and 2) the alteration of target Q-values (resembling noisy labels). We recognize the importance of both scenarios but decided to focus on the former scenario to keep the synthetic experiment manageable.\n\nWe concur that our current design might not entirely capture the dynamics of RL. Exploring a more realistic synthetic experiment is indeed an exciting direction for future research. We will point out such limitations in our manuscript.\n\n\n> **Question 1.3:**\nHave you experimented with applying your method to stronger-performing algorithms?\n\n\nWe have explored the applicability of our method beyond just DER and DrQ. Specifically, our experiments encompass BBF, a state-of-the-art algorithm that learns from scratch on the Atari-100k benchmark, and SimTPR, which leverages pre-trained representations for Atari. In both instances, combining SAM and Reset consistently outperformed using either method in isolation. Detailed empirical results can be found in Table A.2 (attached pdf). We appreciate your insightful suggestion.\n\n\n> **Question 1.4:**\nSeveral fields appear to be missing in Table 1.\n\nThank you for pointing out the placeholders. The updated values for DER + CReLU and DER + LayerNorm, previously marked as “0.xxx,” were provided in Supplementary Section 6, highlighted in red. We apologize for any inconvenience this caused.\n\n\n\n"}, {"review_id": "TFDyEKkwqo", "reviewer": "Reviewer_GE2W", "summary": "This work studies the role of generalization and plasticity in sample-efficient deep RL. This paper proposes sharpness-aware minimization (SAM) to improve generalization in RL. And it provides details on how to use SAM with deep RL algorithms like SAC and Rainbow. Empirical evaluation shows that combined usage of SAM and periodically resetting the last few layers of the network improves sample efficiency in the Atari-100k benchmark. On the other hand, adding SAM does not provide any benefit over resetting in Deepmind Control Suite.", "questions": "1. The paper uses two terms, \"adaptability\" and \"plasticity.\" To me, these words have very similar meanings. Plasticity means the ability to learn from new data, and adaptability means being able to fit new data. Can you please explain in what sense are you using these words and what is the difference between the two? If they have a similar meaning, using just one word in the paper might be better. Currently, the use of these two words causes confusion in section 4. Maybe in the conclusion of Section 4, it is better to say that SAM improves plasticity when there is input non-stationarity and resetting the last few layers improves plasticity when the labels change.\n\n2. What is the confidence interval reported in Figure 2? Is it the 95% bootstrapped confidence interval? ", "limitations": "The key limitation of the paper is that the experiments did not test the central claims of the paper. To claim that both plasticity and generality are important for improving adaptability, methods that inject plasticity into the whole network should be tested. \nIncluding the following baselines in Section 4 with proper hyper-parameter turning and more random seeds will significantly improve the completeness of the claims.\n1. CReLU\n2. LN\n3. L2 (properly tuned) + Reset\n4. CReLU + Reset\n5. LN + Reset\n\nSimilarly, the following baselines in Section 5.2 need to be added.\n1. CReLU + Reset\n2. LN + Reset\n3. L2 + Reset\n4. L2 + best combination of CReLU, LN, Reset\n5. SAM + best combination of CReLU, LN, Reset\n\n\nThe paper contains interesting ideas and promising results. However, the current empirical evaluation is incomplete and does not support the central claim of the paper. Unfortunately, I can not recommend accepting the paper in its current form. But, I'm willing to change my score if the authors can include all the relevant baselines for experiments in Sections 4 and 5.2 and show that the conclusions still hold.\n\nEDIT: I have update my score based on the new experiments and changes in the main message of the paper.", "rating": 7, "confidence": 5, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "This work tackles the critical problem of sample efficient reinforcement learning. Sample efficiency is particularly important for applications where exploration is risky and expensive, such as healthcare.\n\nThe paper proposes using SAM to improve the generalization of deep RL algorithms. SAM has not been used in RL before, so the details regarding the usage of SAM in RL are valuable.\n\nThe paper attempted to study the role of generalization and plasticity in a simple continual supervised learning problem. This problem does not have the confounders found in RL, so it can improve our understanding of the underlying phenomena.\n\nThe results on the Atari-100k benchmark show that combining SAM with resets improves the performance over just using SAM or resets. This exciting result shows that SAM can improve the sample efficiency of deep RL.", "weaknesses": "This paper contains interesting ideas and promising experimental results. The paper's central claim is that generalizability and plasticity constitute different roles in improving performance (lines 6-7). However, the experiments performed in the paper do not adequately test this claim. There are significant issues with the experiments in Sections 4 and 5.2.\n\nThe paper claims that \"we find that incorporating both generalization and plasticity is crucial for improving the model's ability to adapt to new data (Section 4)\" (lines 56-57). Unfortunately, the experiments in Section 4 have the following major problems:\n\n1. The paper claims that generalization and plasticity constitute different roles. To represent generalization, it used SAM, and to represent plasticity, it periodically reset the last few layers. Then two experiments were performed, one where the input changes over time and the other where the labels change. The results show that resetting the last few layers does not help when the input changes. But that is not surprising. Resetting the last few layers only injects plasticity into the last few layers. However, the conclusions drawn from this result are too general. The paper says that plasticity does not help when the input distribution changes, but we can only draw this conclusion if the paper contains results for methods that inject plasticity into the whole network.\n\nThe experiments in Section 4 should include methods that inject plasticity into the full network. This means including methods that reset the last few layers and use CReLUs in the full network, or reset the last few layers and use LayerNorm after all the layers, or using selective reinitialization methods like Continual Backprop or ReDo. From the current results, it is unclear if SAM provides any benefits over plasticity injecting methods.\n\n2. The experiments were only performed for five random seeds (line 73 of the appendix), which raises questions about the statistical significance of these results. These experiments are performed on CIFAR, so they are probably not computationally expensive. I think it should be possible to perform more runs, say 30, to increase the confidence in these results. The paper also needs to report how the confidence interval is calculated in Figure 1.\n\n3. The hyperparameters are not properly tuned for the input adaptation experiment. The experiment used a learning rate of 0.001, which was the smallest one that was tested. Even lower learning rates need to be tested. \n\nSimilar problems exist with the experiments in Section 5.2.\n1. Again, the resetting methods used in these experiments only inject plasticity in the last few layers. The effect of injecting plasticity in the full network needs to be tested before including SAM. \n\n2. The paper proposes L2 regularization as a baseline for SAM. However, it is not tested in the experiments. Experiments with DER show that L2 performs better than SAM, so it is puzzling why L2+Reset is not compared to SAM+Reset. L2+plasticity-preserving-methods should be compared with SAM+Reset.\n\nBesides these major concerns, the clarity of the writing can be improved. For example, the sentence on line 128 needs to be completed. Spending more space to explain section 3.2 will also be valuable for the community. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer GE2W,\n\nWe appreciate your constructive guidance. Based on your feedback,\n- We've refined our Synthetic experiments.\n- We explored the synergies of input and label adaptation techniques in RL experiments.\n- Line 128 has been corrected for clarity.\n- We plan to provide a detailed explanation of SAM in Section 3.\n- We clarified the terminology focusing on “plasticity”.\n\nPlease let us know if you have any further comments or feedback. We will do our best to address them.  \n\n> **Question 2.1:**\nThe experiments in Section 4 have the following major problems:\n- The hyperparameters are not properly tuned. \n- The experiments in Section 4 should include methods including  {CReLU, LN, L2 (properly tuned) + Reset, CReLU + Reset, LN + Reset}.\n- Only five random seeds were used. \n- Needs to report confidence intervals.\n\nThank you for your observations. We've undertaken a rigorous revision of our synthetic experimental setup:\n\n- Tuning Learning Rate & Weight Decay: Recognizing the criticality of tuning for learning rate and weight decay in non-stationary datasets, we've executed an exhaustive search across {0.1, 0.01, 0.001, 0.0001, 0.00001} for both. We decided to incorporate weight decay as a base setup as we found it was important to reduce the variance of each individual run.\n- Broadening Baselines: We broadened our benchmarks to include {LN (B+H), LN (B), SAM (B+H), ReDO (B+H), CReLU (B+H), CReLU (H), Reset (H)}, with B and H denoting backbone and head network parts. For ReDO we tuned the dormant threshold from {0.2, 0.1, 0.05, 0.02, 0.01}. In addition, we further investigated the synergistic interactions of these baselines.\n- Increasing Random Seeds: We increased the number of random seeds from 5 to 30.\n- Reporting Confidence Interval: For each method, we report the 95% confidence interval.\n\nFigure A.1 (attached pdf), reveals a clear bifurcation of algorithms excelling in either input adaptation or label adaptation. For input adaptation, LN(B) and SAM(B+H) demonstrated prominent efficacy. Conversely, for label adaptation, Reset (H) and CReLU (H) were effective. Exploring the synergies of combined approaches offered interesting insights. The best performance across both scenarios was achieved either by strategically blending methods that exhibited proficiency in both input and label adaptation or by incorporating all of the methodologies.\n\nWe hope these comprehensive experiments address your concerns and illustrate the robustness of our message.\n\n\n> **Question 2.2:**\nSimilar problems exist with the experiments in Section 5.2. \n\n- The baselines used in Synthetic Experiments need to be added. \n- L2+plasticity-preserving-methods should be compared with SAM+Reset.\n\nWe extended our experiments in the context of a sample-efficient RL setup, specifically with DrQ on the Atari-100k benchmark. From Table A.1(attached pdf), we observed a pronounced synergy when mixing input adaptation with label adaptation techniques. On the other hand, when we concentrated exclusively on either input or label adaptation, we only observed marginal enhancements.\n\nFor L2 regularization, we conducted careful tuning both individually and in combination with other methods, but its performance was found to be suboptimal. Therefore, we prioritize other methods for exploration. Regarding the specific combinations like Reset(B+H) and LN(B+H) suggested by the reviewer: given the time constraints, we didn't delve into these. However, we are willing to explore these avenues in our revised manuscript.\n\nTo conclude, our extended experiments reaffirm our initial hypotheses, and we're excited to further refine our work based on the insights provided.\n\n\n> **Question 2.3:**\nThe sentence on line 128 needs to be completed. \n\nWe found out that the sentence was inadvertently truncated. The complete sentence should be as follows.  \n\n“... which encourages a model parameter $w$ to find a smoother region of the loss landscape.”\n\nWe will ensure it is corrected in our revised manuscript.\n\n\n> **Question 2.4:**\nSpending more space to explain section 3.2 will also be valuable for the community.\n\nThank you for your suggestion. We understand that Sharpness Aware Minimization (SAM) might be less familiar to the reinforcement learning community. We will expand on Section 3.2 in the revised manuscript to ensure clarity of the methodology.\n\n\n> **Question 2.5:**\nTo me, \"adaptability\" and \"plasticity\" have very similar meanings. If they have a similar meaning, using just one word in the paper might be better. \n\nWe concur with your observation that \"plasticity\" has been ambiguously defined in the reinforcement learning literature. However, upon reviewing the definitions from various literature from continual learning [1,2] and neuroscience [3,4], it does appear that using \"plasticity\" in the same context as \"adaptability\" would be more natural.\n\nTo clarify the definitions:\n- We redefine \"Plasticity\" as the model's ability to adapt.\n- \"Input Plasticity\" refers to the model's capability to adjust to input non-stationarity or changes in $p(x)$.\n- \"Label Plasticity\" refers to the model's capability to adjust to label non-stationarity or shifts in $p(y|x)$.\n- We plan to retitle our paper as \"Enhancing Input and Label Plasticity for Sample Efficient Reinforcement Learning.\"\n\nWe believe this change in terminology and definition will provide more clarity and prevent confusion throughout the paper.\n\n[1] A study on the plasticity of neural networks. Berariu et al., arXiv 2021\n\n[2] Continual backprop: Stochastic gradient descent with persistent randomness. Dohare et al., arXiv 2022.\n\n[3] Neuroplasticity. Costandi, Moheb., MIt Press 2016.\n\n[4] Neuroplasticity: New biochemical mechanisms. Reznikov et al., Springer Healthcare, 2020.\n\n\n> **Question 2.6:**\nWhat is the confidence interval reported in Figure 2?\n\nIt is the 95% bootstrapped confidence interval. We will clarify this in the revised manuscript.\n"}, {"review_id": "zaGmIRoMoy", "reviewer": "Reviewer_dLhS", "summary": "This paper presents a new method to enhance sample efficiency in reinforcement learning, by integrating two existing techniques: sharpness-aware minimization (SAM) and weight resetting. It shows that SAM and resetting work in a complementary way where SAM addresses input adaptability and resetting addresses label adaptability. Experiments on Atari100K and DM control demonstrate the effectiveness of the proposed method.", "questions": "Apart from the points raised above, I have some questions and comments, and would like to hear the authors' feedback:\n- Table 1: Some entries missing (\"xxx\" in the table).\n- Line 128: The sentence (\"..., which encourages to Here, ...\") is confusing.\n- I notice the authors have cited [10] in the related works. I am wondering if the authors have tried the soft resetting strategy on Atari, since it has been shown to be effective.\n- It would be better to add error bars in Figure 1(right), Figure 4, and Figure 5.", "limitations": "The limitations of the proposed method have been adequately discussed in Section 6.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- This paper is clear, well-written, self-contained, and enjoyable to read.\n- While the proposed method simply combines two existing approaches, such a combination is novel and well justified by the synthetic experiments and ablation studies.\n- The experiments are comprehensive and well executed.", "weaknesses": "- The authors mention that the relatively small improvement in DMC might be due to the reduced visual variety in DMC. The [DMControl Generalization Benchmark](https://github.com/nicklashansen/dmcontrol-generalization-benchmark) augments DMC with rich visual variety, which could be a good testbed to validate the claim.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer dLhs,\n\n\nThank you for your constructive feedback. To address your concerns, \n- We explored the DMControl Generalization Benchmark (DMC-GB). Our findings from this extended study suggest that DMC exhibits a large degree of input non-stationarity than we initially presumed. \n- We rectified discrepancies in Table 1, improved clarity in various sections, and enhanced our figures.\n- In a broader perspective, we've also conducted rigorous tests in synthetic environments and experimented with SAM + Reset on more advanced algorithms. Further detail is described in our general response.\n\nPlease let us know if you have any further comments or feedback. We will do our best to address them. \n\n\n> **Question 3.1:**\nThe authors mention that the relatively small improvement in DMC might be due to the reduced visual variety in DMC. The DMControl Generalization Benchmark augments DMC with rich visual variety, which could be a good testbed to validate the claim.\n\nThank you for bringing up the DMControl Generalization Benchmark (DMC-GB). In contrast to the traditional approach of DMC-GB—training on clear inputs and testing on noisy ones—we opted to train and test on noisy inputs. This approach aligns with our intention of illustrating the strength of different generalization techniques against input non-stationarity.\n\nOur experimental results in Table A.3 (attached pdf) indicate that when paired with a reset, SAM performs comparably to Adam. Without the reset, however, SAM surpasses Adam. An intriguing observation was that the DMC environment displayed a higher degree of input non-stationarity compared to Atari, even without the background noise introduced in DMC-GB. This challenges our initial hypothesis, suggesting that handling severe input non-stationarity might be better suited for DMC.\n\nThen, if DMC inherently demands addressing input non-stationarity, why does the reset strategy perform so well? The context lies in the different reset strategies between environments. In the original reset paper referenced as [1], DMC employs an aggressive reset strategy, reinitializing nearly 90% of its network for DrQ and 100% of its network on SAC. Instead, in environments like Atari, 50% of the network is reinitialized. This aggressive approach, highlighted in [1] effectively manages input non-stationarity, potentially overshadowing the benefits of SAM when paired.\n\nOur research primarily focuses on resetting the final few layers. However, the efficacy of broad resets in tackling input non-stationarity is evident in our synthetic experiments, where Reset (B, H) showcases effectiveness in Figure A.1 (attached pdf) While these aggressive resets worked well for DMC, we believe this strategy is not always optimal as evidenced by the results of Reset (B, H) from Atari in Table A.1. Extensive reinitialization forces the model to relearn, and as networks grow larger, this becomes more challenging. \n\nThus, we believe blending different strategies to combat input non-stationarity is essential to further enhance the performance of both DMC and DMC-GB.\n\n[1] The Primacy Bias in Deep Reinforcement Learning., Nikishin et al., ICML 2022.\n\n> **Question 3.2:**\nTable 1: Some entries missing (\"xxx\" in the table).\n\nThe updated values for DER + CReLU and DER + LayerNorm, previously marked as “0.xxx,” are provided in Supplementary Section 6, highlighted in red. We apologize for any inconvenience this caused.\n\n\n> **Question 3.3:**\nLine 128: The sentence (\"..., which encourages to Here, ...\") is confusing.\n\nThank you for pointing this out. We found out that the sentence was inadvertently truncated. The correct sentence should read:\n\n\"... which encourages a model parameter $w$ to find a smoother region of the loss landscape.”\n\nWe will ensure it is corrected in the revised manuscript.\n\n\n> **Question 3.4:**\nI noticed the authors have cited [10] in the related works. I am wondering if the authors have tried the soft resetting strategy on Atari since it has been shown to be effective.\n\nThank you for highlighting the soft resetting strategy referenced in [10].\n\nIndeed, we have experimented with the soft reset approach on the Atari benchmarks, as detailed in our appendix section 5. Our findings indicate that when the replay ratio is set at 2 (the standard setup in our primary experiments), the use of the soft reset (i.e., Shrink & Perturb) tends to degrade performance. However, as we increment the replay ratio, the advantages of the soft reset became notably evident, even surpassing the results of well-known state-of-the-art algorithms.\n\n\n> **Question 3.5:**\nIt would be better to add error bars in Figure 1(right), Figure 4, and Figure 5.\n\nWe will incorporate error bars in Figures 1(right), 4, and 5 in the revised manuscript.\n"}, {"review_id": "cJ2nRRxZaN", "reviewer": "Reviewer_mzhA", "summary": "\nThe paper studies the generalization and plasticity of deep reinforcement learning (RL) agents. The paper's primary contribution is a combination of the existing Sharpness-Aware Minimization (SAM) and Resetting mechanisms. The main insight is that combining SAM+Resets provides additive benefits: the former induces feature sparsity (and hence could be argued to improve generalization), while the latter yields more active units for the last layers (and hence could be argued to increase network plasticity). Moreover, the authors conduct an analysis on synthetic data and arrive at the conclusion that SAM improves adaptability to input changes, while resets mostly improve the ability to adapt to new targets. When SAM+Resets are applied to Data-Efficient Rainbow (DER) and Data-regularized Q (DrQ) on Atari 100k, the combination gives consistent improvements.\n", "questions": "The paper doesn’t study continuous control from proprietary states — what was the motivation for excluding the setting?\n\nThe combination of preliminaries in Section 3 and the experimental details in Section 5.1 is confusing: first, the authors give background about the Rainbow and Soft Actor-Critic algorithms, while later use their modifications, DER and DrQ, for actual experimentation. The reviewer suggests describing directly DER and DrQ.\n", "limitations": "While the paper presents evidence about the effects of SAM+Resets on network adaptability, an explanation is lacking: why SAM and Resets enhance input and label adaptability and, for example, not visa versa?\nLikewise, the reviewer finds intriguing the result in Figure 2a (bottom left): on their own, both resets and SAM decrease the fraction of active units, while their combination yields a strong increase. Why?\nDiving deeper into explaining the phenomena would increase the significance of the paper.\n", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "The main strength of the paper is a thorough empirical investigation:\n- The majority of claims are supplemented with supporting evidence: for example, the authors report sharpness metrics (such as the trace of the Fisher matrix and the largest Hessian eigenvalue) to demonstrate that SAM increases the flatness of the solution (the result is non-trivial since, unlike in supervised learning, the loss surface is always changing in RL).\n- The paper presents ablations for the design choices when applying SAM to RL algorithms (for example, whether to apply SAM for the whole network or for a subset of layers in DER or whether to apply SAM for both actor and critic networks in DrQ or individually).\n- The SAM+Resets combination is thoroughly compared to other methods that were proposed to address plasticity loss (such as layer normalization and concatenated ReLUs).\n", "weaknesses": "There are several weak spots that prevent the reviewer from assigning a higher score:\n- The SAM+resets combination is applied for the DER algorithm in Atari 100k, which is a weaker baseline than the SPR algorithm used in the original resets paper; the choice makes the reviewer wonder if the addition from SAM would be smaller for SPR and other more advanced algorithms.\n- Figure 4 contains an encouraging demonstration of improvements in replay ratio scalability for DrQ on Atari-100k from adding SAM on top of resets. However, the paper does not have scalability results for the other settings (DrQ on DMC / DER on Atari 100k / SAC on DMC from proprietary states) making the reviewer question the generality of the insight.\n\nThere are also several minor weaknesses:\n- Table 1 contains placeholder performance values 0.xxx for DER + CReLU\n- Putting in bold the results for SAM + Reset in Table 2 might mislead a limited-attention reader. The difference between SAM + Reset 516 (441, 590) and Reset 514 (434, 590) is insignificant and should not be highlighted.\n\nAddressing the outlined limitations might increase the score.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer mzhA,\n\nThank you for your constructive feedback. To address your concerns, \n- We've experimented with SAM on advanced algorithms. \n- We provide an explanation of the generality of our insight. \n- Rectified .xxx in Table 1. \n- The presentation in Table 2 will be revised to avoid misinterpretations.\n- Explained the rationale behind excluding state-based environments. \n- We plan to improve clarity between DER and DrQ. \n\nPlease let us know if you have any further comments or feedback. We will do our best to address them. \n\n\n> **Question 4.1:**\nThe reviewer wonders if the addition from SAM would be smaller for SPR and other more advanced algorithms.\n\nTo clarify, our method primarily builds upon the DrQ, as detailed in Supplementary Section 5. Contrasting with SR-SPR, which couples DrQ with SSL and Reset, our approach couples DrQ with SAM and Reset. We have tried SR-SPR + SAM and found that the performance gains were indeed marginal. This can likely be attributed to the inherent generalization enhancements from the SSL objectives already present in SPR.\n\nAdditionally, our SAM + Reset combination has been rigorously tested with advanced algorithms like BBF and SimTPR on the Atari-100k benchmark. The results, showcased in Table A.2 (attached pdf), confirm the robustness and efficacy of this synergy.\n\nIn our revised manuscript, we'll accentuate these findings, emphasizing the distinctions and comparative benefits of our approach.\n\n> **Question 4.2:**\nThe paper does not have scalability results for the other settings (DER on Atari 100k / DrQ, SAC on DMC) making the reviewer question the generality of the insight.\n\nFor DER on Atari-100k: DER utilizes shallow and wide architecture and focuses on amplifying generalizable representations. Consequently, testing scalability on DER might detract from its foundational goal of enhancing the generalization of larger networks.\n\nFor DrQ and SAC on DMC: Initially, we surmised that DMC's limited generalization advantages stemmed from lower input non-stationarity. Yet, our DMC-GB experiments on Table A.3 (attached pdf) and subsequent comparisons between Atari and DMC revealed that DMC indeed possesses a greater input non-stationarity.\n\nIn addition, DMC employs rigorous reset protocols: DrQ resets almost 50% of its network, and SAC resets it entirely, as detailed in [1]. These comprehensive resets inherently counter input non-stationarities, consistently restoring the system to an initial state. As [1] has demonstrated the efficacy and scalability of such resets, we believe this further validates the relevance and breadth of our insights.\n\nOur further tests included scaling SAM+Reset on SimTPR, a 30-layer CNN model. As presented in Table A.1, the synergy of SAM and Reset is evident at a replay ratio of 2 and retains its potency when increased to 4. Recent findings from BBF also highlight the harmonious interplay between generalization (via strong L2 regularization and Self-Supervised objectives) and plasticity (via resets), across replay ratios from 2 to 8.\nIn light of these, we remain confident in the broader applicability of our insights.\n\n[1] Sample-Efficient RL by Breaking the Replay Ratio Barrier, D’Oro et al., ICLR 2023.\n\n> **Question 4.3:**\nTable 1 contains placeholder performance values 0.xxx for DER + CReLU\n\nThank you for pointing out the placeholders. The updated values for DER + CReLU and DER + LayerNorm, previously marked as “0.xxx,” are provided in Supplementary Section 6, highlighted in red. We apologize for any inconvenience this caused.\n\n> **Question 4.4:**\nPutting in bold the results for SAM + Reset in Table 2 might mislead a limited-attention reader.\n\nWe agree that emphasizing the results for SAM + Reset could mislead readers given the minimal difference with just the Reset. We will present the results without bolding.\n\n> **Question 4.5:**\nThe paper doesn’t study continuous control from proprietary states — what was the motivation for excluding the setting?\n\nWe emphasized image-based environments due to their inherent non-stationarity and the pronounced distribution shifts. In these tasks, even minor agent actions can drastically alter pixel distributions, demanding robust generalization. Conversely, state-based scenarios generally present more consistent input distributions, making them less challenging from a generalization perspective.\n\nWe'll elucidate this choice further in our revised manuscript for clarity.\n\n> **Question 4.6:**\nThe combination of preliminaries in Section 3 and the experimental details in Section 5.1 is confusing. The reviewer suggests describing directly DER and DrQ.\n\nThank you for pointing out the potential confusion. DER and DrQ, while rooted in similar foundational algorithms, differ in aspects like model architecture and data augmentation. In our revised manuscript, we'll refine the presentation, explicitly highlighting these differences to ensure a clear understanding.\n\n> **Question 4.7:**\nWhy do SAM and Resets enhance input and label adaptability and not visa versa? In Figure 2a, both resets and SAM decrease the fraction of active units in the head, while their combination yields a strong increase. Why? \n\nOn the distinction in adaptability enhancements: SAM, by its design, results in sparser features, which may not readily allow for flexibility in adapting to shifting label dynamics. Reset, on the other hand, appears to naturally assist in label adaptability by preventing overfitting to evolving label relationships. However, intuitively, its precise contribution to input adaptability is less evident.\n\nRegarding the combined effect observed in Figure 2a: SAM's sparsity might necessitate a more active head unit to decipher diverse feature combinations, and when paired with Resets, this effect amplifies, leading to an overall rise in active units for the head. However, we acknowledge that this is a hypothesis, and further investigations are needed to solidify this understanding."}]

Author rebuttal:
We sincerely appreciate all four reviewers for their constructive and insightful comments.

The reviewers recognized the strengths of our paper as follows:
- A fresh insight into the dissection of generalization and plasticity (Reviewer 1mez, GE2W).
- The introduction of a synergistic solution: SAM + Reset, which seamlessly integrates without any architectural modifications (Reviewer dlhS, mzhA).

We recognize that reviewers have highlighted several key points for us to address:
- Ensuring the robustness of our synthetic experiments (Reviewer GE2W).
- Delving deeper into explored synergies (Reviewer 1mez, GE2W).
- Investigating the approach's applicability to advanced algorithms (Reviewer 1mez, mzhA).
- Extension to various domains (Reviewer dlhS).

Below, we address each of these points in the following sections. Please refer to the attached PDF for detailed figures and tables.

> **Refinements to the synthetic experiments**

In response to feedback, we have:

- Adjusted learning rates and weight decay from a wide range of values.
- Broadened our experimental baselines, including ReDO [1]. For clarity: "B" signifies the network's backbone and "H" indicates its head. Our findings spotlight the efficacy of LN(B) and SAM(B+H) for input adaptation, and the efficacy of CReLU(H) and Reset(H) for label adaptation. Moreover, the interplay between these methods has been thoroughly investigated.
-Increased random seeds from 5 to 30.

As illustrated in Figure A.1, we found that 

- Combining SAM(B+H) with LN(B), and pairing CReLU(H) with Reset(H), resulted in incremental improvements.
- Merging input-focused methods like SAM(B+H) or LN(B) with label-focused counterparts, such as CReLU(H) or Reset(H), consistently delivered impressive outcomes.
- Integrating all the methods together yielded the most significant enhancements.

These results vividly underscore the synergy between generalization and plasticity, reinforcing our research's relevance.

[1] The Dormant Neuron Phenomenon in Deep Reinforcement Learning. Sokar et al., ICML 2023



> **Deeper Exploration of Synergies**

Responding to the feedback received, we've intensified our investigation into the potential combinations of techniques within both synthetic and RL settings, as described in Figure A.1 and Table A.1.

Our in-depth analysis revealed that while methods tailored for input and label adaptability operate distinctively, they produce a remarkable synergy when fused. We believe this comprehensive exploration solidifies the importance of our contributions.

> **Extension to Advanced Algorithms**

Moving beyond the realms of DER and DrQ, we have ventured into more advanced algorithms:

- BBF [2]: Recognized as a state-of-the-art algorithm for the Atari-100k.
- SimTPR [3]: An algorithm that leverages pretrained representations for Atari.

Table A.2 summarizes our findings:

- BBF: Utilizing SSL + Reset or SAM + Reset distinctly outperforms isolated implementations.
- SimTPR: The amalgamation of SAM and Reset magnifies results compared to their standalone use. Importantly, a replay ratio of 4 eclipses the results at a ratio of 2, highlighting the approach's scalability.

Collectively, these experiments reiterate our method's generality and its alignment with modern algorithms.

[2] Bigger, Better, Faster: Human-level Atari with human-level efficiency. Schwarzer et al., ICML 2023.

[3] On the Importance of Feature Decorrelation for Unsupervised Representation Learning for RL., ICML 2023.


> **Extension to DMC Generalization Benchmark**

Diverging from the conventional DMC-GB-500k benchmark [4], we opted for a consistently noisy environment both for training and testing. This choice was taken to evaluate generalization techniques under persistent input non-stationarity. Here, we used 5 different environments, {walker-walk, cart pole-swing-up, finger-spin, walker-stand, ball-in-cup-catch}.

Table A.3 encapsulates our findings:

SAM, when paired with reset, exhibits a performance closely mirroring that of Adam. However, without the influence of a reset, SAM surpasses Adam.

Contrary to our initial presumptions, we found that the DMC environment manifested a pronounced degree of input non-stationarity in contrast to Atari, even when excluding DMC-GB's intrinsic noise. This observation leads to the insight that "DMC inherently calls for methods addressing input non-stationarity."

Such a realization prompted a deeper introspection into the role of reset strategy. As detailed in the original reset paper [5], DMC employs a notably aggressive reset strategy — engaging up to 90% reset for the DrQ algorithm and a comprehensive reset for SAC. Meanwhile, Atari adopts a milder reset approach, reinitializing around half of its network. DMC's rigorous reset approach serves as an effective antidote to its inherent input non-stationarity. This might explain why the benefits of SAM appear subdued when combined with such an assertive reset. Intriguingly, our evaluations on synthetic and Atari platforms revealed that Reset (B, H) strategies faltered in performance. This highlights the necessity for a more tailored approach, especially as network sizes expand.

To summarize, we believe effectively managing input non-stationarity in RL is a key to improving performance for both DMC and DMC-GB.

[4] Generalization in Reinforcement Learning by Soft Data Augmentation., Hansen et al., ICRA 2021
[5] The Primacy Bias in Deep Reinforcement Learning., Nikishin et al., ICML 2022

> **Conclusion**

In conclusion, we deeply appreciate the constructive feedback from the reviewers. Their insights have refined our research, and we believe our revisions now thoroughly address the raised concerns. We eagerly await further feedback.

Warm regards,
Authors of Submission 12455.


Meta-review: Clear accept. 

This paper investigates improving sample and computational efficiency in RL via a new method that combines periodic resets and Sharpness-Aware Minimization (two know ideas, but SAM is new to RL) and showing both are needed. The reviewers found the method simple and appealing, and the general thesis interesting and well motivated.

Experiments show the new method helps when applied in supervised learning and when combined with two recent, popular RL methods in Atari 100k benchmark and Deepmind Control Suite. Reviewer GE2W pointed out major issues with the experiments, and the authors heroically made the requested changes (including increasing # of seeds from 5 to 30) and GE2W was happy; although this changed how the results should be discussed (more on that below). Reviewer mzhA had similar requests for additional results and clarifications and the authors provided them. 

All reviewers voted to accept, and two domain experts voted clear accept. The reviewers noted the paper was well written and nicely polished with measured claims and well executed experiments.

Please ensure the following changes are complete for camera ready:
- do not bold results in tables where the errorbars (+/-) overlap. This is misleading
- Reviewers dLhS and GE2W gave advice on how to adjust the text based on the new results (e.g., some results are now a statistical tie which is ok if accurately and honestly discussed). Please implement these writing adjustments to the abstract, intro and results discussion.
Decision: Accept (poster)