Title: PLASTIC: Improving Input and Label Plasticity for Sample Efficient Reinforcement Learning
Venue: NeurIPS 2023
Accepted: yes
Keywords: Reinforcement Learning

Sharpness Minimization

Generalization

Plasticity

Deep 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

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 .

1 Introduction

In Reinforcement Learning (RL), achieving sample efficiency is crucial in various domains including robotics, autonomous driving, and healthcare, where data acquisition is constrained and expensive [42, 36]. In theory, off-policy RL algorithms promise increased sample efficiency by allowing multiple updates of policy or value functions from a single data instance [26, 24, 43]. However, they tend to suffer from the loss of plasticity , a phenomenon where the models overfit to earlier interactions and fail to adapt to new experiences [52, 19, 43].

The origins of the loss of plasticity are a focus of contemporary 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].

PLASTIC (RR8)

(RR4)

(RR2)

(RR2)

(RR4)

SR-SPR(RR8)

MLR

PlayVirtual

(RR4)

PLASTIC<<(RR8)

†

CReLU

Atari&100k

26#tasks

SAM

LN

Base

0.25

Reset

PLASTIC

0.20

0.35

0.30

0.40

CReLU

DMC&Suite

11#tasks

SAM

LN

Base

200

Reset

PLASTIC

100

400

300

500

RL#Benchmark

CReLU

Input&Plasticity

SAM

LN

Base

65

Reset

PLASTIC

60

75

70

80

CReLU

Label&Plasticity

SAM

LN

Base

67

Reset

PLASTIC

66

69

68

70

Synthetic#Experiment

Another line of investigation emphasizes the role of gradient propagation, highlighting that as training 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 a 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, an example of the latter might be an action once considered to be positive but later perceived as negative. Therefore, throughout this work including our synthetic experiments, we categorize the model's plasticity into two cases:

Input Plasticity: adaptability to shifts in input data distributions, p ( x ) .

Label Plasticity: adaptability to evolving labels (returns) for given inputs (observations), p ( y | x ) .

For our exploration, we employ the CIFAR-10 dataset [37]. To assess input plasticity, we train the model from a 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-Aware 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] are used to enhance the latter.

Our synthetic experiments yield clear distinctions between the proposed two concepts of plasticity. As illustrated on the left side of Figure 2, smoother minima of loss largely improve input plasticity, while maintaining the amount of gradient propagation has a pronounced impact on label plasticity. Building on these findings, we propose an algorithm called PLASTIC , which combines SAM optimizer, LN, periodic Reset, and CReLU 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-100k and Deepmind Control Suite benchmarks, PLASTIC achieved competitive performance to the leading methods, underscoring its potential in enhancing RL's sample efficiency.

In summary, our 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 a smooth region of loss surface and preserving gradient propagation.

Empirically, PLASTIC achieved competitive performance on challenging RL benchmarks, including 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 (e.g., robotics, educational agents) or safety considerations (e.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 a single data collection. However, this advantage may introduce a pitfall: the increased update rate can lead to overfitting, undermining the model'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 augmentation excels in visual RL algorithms [73, 72, 40, 55, 68], particularly when combined with self-supervised learning methods [39, 58, 77].

Regularization: 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].

Self-Supervised Learning: By incorporating auxiliary learning objectives, self-supervised learning has emerged as a potential solution. Objectives encompass pixel or latent space reconstruction [74, 77], future state prediction [58, 19, 41], 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 in RL is driven by the inherent need for plasticity as agents consistently encounter new inputs and evolving input-output relationships.

A 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 in Atari environments. While the 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 a wider and smoother loss surface generalize better to unseen datasets [34, 30, 71]. Following this insight, the sharpnessaware minimization (SAM) optimizer [21] has recently gained attention in supervised learning [13, 44], aiming to minimize both training loss and its sharpness.

Concurrently, various studies point out a progressive decline in a 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 in 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 to 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-exploitation trade-offs.

The Q-value updates in the Rainbow algorithm follow the principle of minimizing the TemporalDifference (TD) error which is defined as:

where w denotes the model weights, w -is the weights of the target model, and τ = ( s, a, r, s ′ ) is the transition sampled from the replay buffer B .

Noisy Networks [22] introduce stochasticity into the model weights, reparameterizing them as:

where µ and σ are the learnable weights, whereas ˜ ε is a random vector sampled from an isotropic Gaussian distribution, adding randomness to the layer. (We denote element-wise product by ' · '.)

Soft Actor-Critic. Soft Actor-Critic (SAC) [24], a prevalent off-policy algorithm for continuous control, aims to maximize the expected return coupled with an entropy bonus. SAC consists of a policy π and a critic model Q , each parameterized by weights θ and w respectively.

In training, the critic Q w is trained to minimize the following objective, defined as

where α is the entropy coefficient and τ = ( s, a, r, s ′ ) is the transition sampled from the buffer B .

Subsequently, the policy π θ is jointly trained to maximize the following objective:

Throughout this paper, all of the plasticity-preserving methods (LN, SAM, Reset, CReLU) will be integrated and evaluated on top of these standard off-policy RL algorithms.

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-exploitation tradeoff. To alleviate this complexity, we design synthetic supervised learning experiments within a controlled framework using the CIFAR-10 dataset [37]. Our focus is to evaluate the model's adaptive capabilities under two distinct adaptation scenarios:

Input Adaptation: In RL, as the agent continually interacts with the environment, it constantly encounters new data through exploration. Effective adaptation to this data is 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 labels uniformly reassigned, ensuring a 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 to 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 at 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 a 95% confidence interval.

To delve into the impact of pursuing a smooth loss surface and preserving gradient propagation, we selected four distinct methods: Layer Normalization (LN) and Sharpness-Aware Minimization

LN

SAM

CReLU

Reset

5

0

5

10

15

Input Adaptation

Improvement of individual method

LN+SAM

LN+CReLU

LN+Reset

SAM+CReLU

SAM+Reset CReLU+Reset

PLASTIC

5

0

5

10

15

Input Adaptation

Improvement of combined methods

1

0

1

2

3

4

5

Label Adaptation

1

0

1

2

3

4

5

Label Adaptation

(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 , 20 } 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 (LN + SAM, ..., CReLU + Reset) with a 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 a clear distinction between the effects of various methods on input and label adaptation capabilities. The left panel illustrates the performance of 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 in input adaptation while exhibiting a significant impact on label adaptation.

Turning our attention to the right panel, the focus is on the combined synergies of methods. The pairing of LN with SAM predominantly enhances input adaptation. When LN is merged with either CReLU or Reset, there's a 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 our notion of distinct yet complementary roles of seeking smooth loss surface and enhancing gradient propagation to improve the model's plasticity.

4 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 Atari-100k [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 on a large pre-trained model (Section 4.4).

Ablation study for different combinations of plasticity-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

measured by averaging over 10 independent trials. We use Data-Regularized Q-learning (DrQ) [73] as a 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 2 million environment steps. Within this setting, we use 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 a selection of plasticity-preserving 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-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 ρ from { 0 . 03 , 0 . 1 , 0 . 3 } , selecting the model with the best performance.

For enhancing gradient propagation, we 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], we 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 ( λ max ) [33]. This metric provides insights into the curvature of the loss function, where a larger value indicates a sharper and more intricate optimization landscape and a 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 to a Human Normalized Score (HNS) as HNS = agent_score - random_score human_score - random_score , which quantifies the relative performance of the agent compared to 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 at evaluating the sample efficiency of the PLASTIC algorithm. Our examination involved contrasting its efficacy against individual plasticitypreserving methods. For this purpose, we used the Atari-100k and DMC-Medium benchmarks.

Table 1 showcases the comparative outcomes of these methods. From the results, we observed that while individual plasticity-preserving 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. A comprehensive set of results for each environment is available in the Appendix.

None

LN

SAM

CReLU

Reset PLASTIC

Method

0

5

10

15

20

25

30

Maximum Eigen Value (

)

0.0

0.2

0.4

0.6

0.8

1.0

Environment Steps (

10 5

)

0.0

0.2

0.4

0.6

0.8

1.0

(Head) Fraction of AU

None

LN

SAM

Reset

CReLU

PLASTIC

To delve deeper into the mechanisms driving these performance enhancements, we focused on two specific metrics: the maximum eigenvalue of the Hessian ( λ max ) 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 a 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 to a smooth loss region while maintaining a large number of active units.

Our prior synthetic experiments provided insight that convergence on a 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-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 a scaling replay ratio.

For our analysis, we scaled the replay ratio up to 8 and evaluated the efficacy of PLASTIC on the Atari-100k benchmark. In addition, we compared PLASTIC against notable state-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, PlayVirtual [76], MLR [78], SR-SPR [19], and PLASTIC can parallelize multiple runs on a single GPU. Thus, for a fair comparison, we 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, we introduced a simplified alternative, PLASTIC † . 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 θ t = αθ t -1 + (1 -α ) ϕ , 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 in Table 2 and Figure 1, PLASTIC † establishes a Pareto frontier between Inter-Quartile Mean (IQM) and computational cost, illustrating its notable computational efficiency. While Ef-

ficientZero exhibits state-of-the-art performance, it uniquely employs a search algorithm coupled with a domain-specific heuristic, namely early environment resets. When excluding these specific elements, PLASTIC stands out as a strong competitor against the other methods.

Through these results, we found that PLASTIC can effectively prevent the loss of plasticity by incorporating various plasticity-preserving strategies. Furthermore, for practitioners, the construction of the Pareto frontier by PLASTIC † 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 a 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 a 30-layer convolutional network, SimTPR is pretrained via self-predictive representation learning on suboptimal Atari video datasets. As SimTPR does not employ Layer Normalization (LN) and Concatenated ReLU (CReLU) activations, we employ SAM to seek a 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 3 provides a summary of applying either SAM or Reset to the pretrained model. Solely scaling the replay ratio during fine-tuning leads to a noticeable decrement in its efficacy. However, we observed that the usage of Reset or SAM can counteract this decrease, leading to a pronounced enhancement in performance. Furthermore, the integration of both techniques surpassed individual contributions, indicating the presence of synergy.

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.

4.5 Ablation Studies

To further analyze the interplay between seeking a smooth loss surface and improving gradient propagation, we have conducted a series of ablation studies for different combinations of plasticitypreserving methods. Here, we averaged the performance over 10 random seeds.

Table 4 showcases the results when applying various combinations to the Atari-100k benchmark. The first five and the last rows echo the results from Table 1. It becomes clear that the concurrent pursuit of a smooth loss surface and the enhancement of gradient propagation surpasses the individual application of either strategy (for instance, LN + CReLU > LN ≈ CReLU). The comprehensive integration of all methods, referred to as PLASTIC, demonstrates the highest performance.

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 category can yield notable improvements. For scenarios constrained by computational resources, using LN

and Reset is recommended over SAM, due to the latter's extra computational burden, and CReLU, given its added parameters (as noted in PLASTIC † in Section 4.3). Additionally, when one is working with a 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 structure. In such instances, as detailed in Section 4.4, the combination of SAM and Reset emerges as a potent solution to improve downstream performance.

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 propagation 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 efficacy in more intricate settings, such as MetaWorld [79], or Procgen [15]. These environments introduce greater non-stationarities, challenging the model'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 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, an 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 RL agents We hope our findings open up a 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 (IITP) grant funded by the Korean government (MSIT) (No.2019-0-00075, Artificial Intelligence Graduate School Program(KAIST)). HL, HK, and JC acknowledge support from the National Supercomputing Center with supercomputing resources including technical support (KSC-2023-CRE-0074). HC and CY acknowledge support from the Institute of Information & communications Technology Planning & Evaluation (IITP) grant (No. 2022-0-00184, Development and Study of AI Technologies to Inexpensively Conform to Evolving Policy on Ethics) and the National Research Foundation of Korea (NRF) grant (No. RS-2023-00211352), both funded by the Korea government (MSIT).

References

Momin Abbas, Quan Xiao, Lisha Chen, Pin-Yu Chen, and Tianyi Chen. Sharp-maml: Sharpnessaware model-agnostic meta learning. In Proc. the International Conference on Machine Learning (ICML) , 2022. 19

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, 6

Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Marc Bellemare. Deep reinforcement learning at the edge of the statistical precipice. Proc. the Advances in Neural Information Processing Systems (NeurIPS) , 34:29304-29320, 2021. 6

Jordan Ash and Ryan P Adams. On warm-starting neural network training. Proc. the Advances in Neural Information Processing Systems (NeurIPS) , 33:3884-3894, 2020. 7

Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.06450 , 2016. 2

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

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

Marc G Bellemare, Will Dabney, and Rémi Munos. A distributional perspective on reinforcement learning. In Proc. the International Conference on Machine Learning (ICML) , pages 449-458. PMLR, 2017. 3

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

David Bertoin and Emmanuel Rachelson. Local feature swapping for generalization in reinforcement learning. Proc. the International Conference on Learning Representations (ICLR) , 2022. 3

Nils Bjorck, Carla P Gomes, and Kilian Q Weinberger. Towards deeper deep reinforcement learning with spectral normalization. Proc. the Advances in Neural Information Processing Systems (NeurIPS) , 34:8242-8255, 2022. 3

Edoardo Cetin, Philip J Ball, Steve Roberts, and Oya Celiktutan. Stabilizing off-policy deep reinforcement learning from pixels. arXiv preprint arXiv:2207.00986 , 2022. 3

Xiangning Chen, Cho-Jui Hsieh, and Boqing Gong. When vision transformers outperform resnets without pre-training or strong data augmentations. Proc. the International Conference on Learning Representations (ICLR) , 2022. 3

Xiangning Chen, Cho-Jui Hsieh, and Boqing Gong. When vision transformers outperform resnets without pre-training or strong data augmentations. In International Conference on Learning Representations , 2022. 19

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 2048-2056. PMLR, 2020. 9

Terrance DeVries and Graham W Taylor. Improved regularization of convolutional neural networks with cutout. arXiv preprint arXiv:1708.04552 , 2017. 3

Shibhansh Dohare, Juan Hernandez-Garcia, Parash Rahman, Richard Sutton, and A Rupam Mahmood. Loss of plasticity in deep continual learning. 2023. 3

Shibhansh Dohare, Richard S Sutton, and A Rupam Mahmood. Continual backprop: Stochastic gradient descent with persistent randomness. arXiv preprint arXiv:2108.06325 , 2021. 2, 3, 6

Pierluca D'Oro, Max Schwarzer, Evgenii Nikishin, Pierre-Luc Bacon, Marc G Bellemare, and Aaron Courville. Sample-efficient reinforcement learning by breaking the replay ratio barrier. In Proc. the International Conference on Learning Representations (ICLR) , 2023. 1, 3, 7, 8, 16

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

Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. In Proc. the International Conference on Learning Representations (ICLR) , 2021. 2, 3, 6, 19

Meire Fortunato, Mohammad Gheshlaghi Azar, Bilal Piot, Jacob Menick, Ian Osband, Alex Graves, Vlad Mnih, Remi Munos, Demis Hassabis, Olivier Pietquin, et al. Noisy networks for exploration. arXiv preprint arXiv:1706.10295 , 2017. 3, 4

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. 3

Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In Proc. the International Conference on Machine Learning (ICML) , pages 1861-1870. PMLR, 2018. 1, 3, 4, 6

Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301.04104 , 2023. 8

Matteo Hessel, Joseph Modayil, Hado Van Hasselt, Tom Schaul, Georg Ostrovski, Will Dabney, Dan Horgan, Bilal Piot, Mohammad Azar, and David Silver. Rainbow: Combining improvements in deep reinforcement learning. In Proc. the AAAI Conference on Artificial Intelligence (AAAI) , 2018. 1, 3, 4, 6, 16

Takuya Hiraoka, Takahisa Imagawa, Taisei Hashimoto, Takashi Onishi, and Yoshimasa Tsuruoka. Dropout q-functions for doubly efficient reinforcement learning. Proc. the International Conference on Learning Representations (ICLR) , 2022. 3

Minyoung Huh, Hossein Mobahi, Richard Zhang, Brian Cheung, Pulkit Agrawal, and Phillip Isola. The low-rank simplicity bias in deep networks. 2023. 3

Maximilian Igl, Gregory Farquhar, Jelena Luketina, Wendelin Boehmer, and Shimon Whiteson. Transient non-stationarity and generalisation in deep reinforcement learning. Proc. the International Conference on Learning Representations (ICLR) , 2021. 3

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

Jean Kaddour, Linqing Liu, Ricardo Silva, and Matt J Kusner. When do flat minima optimizers work? Proc. the Advances in Neural Information Processing Systems (NeurIPS) , 35:1657716595, 2022. 19

Łukasz Kaiser, Mohammad Babaeizadeh, Piotr Miłos, Bła˙ zej Osi´ nski, 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. 5

Simran Kaur, Jeremy Cohen, and Zachary Chase Lipton. On the maximum hessian eigenvalue and generalization. In I Can't Believe It's Not Better Workshop: Understanding Deep Learning Through Empirical Falsification , 2023. 6

Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In Proc. the International Conference on Learning Representations (ICLR) , 2017. 3, 19

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. Proc. the International Conference on Learning Representations (ICLR) , 2015. 19

Jens Kober, J Andrew Bagnell, and Jan Peters. Reinforcement learning in robotics: A survey. The International Journal of Robotics Research , 2013. 1, 3

Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009. 2, 4

Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009. 16

Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Curl: Contrastive unsupervised representations for reinforcement learning. In Proc. the International Conference on Machine Learning (ICML) , 2020. 3

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. 3

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 arXiv:2306.05637 , 2023. 3, 8

Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643 , 2020. 1, 3

Qiyang Li, Aviral Kumar, Ilya Kostrikov, and Sergey Levine. Efficient deep reinforcement learning requires regulating overfitting. arXiv preprint arXiv:2304.10466 , 2023. 1, 3

Yong Liu, Siqi Mai, Xiangning Chen, Cho-Jui Hsieh, and Yang You. Towards efficient and scalable sharpness-aware minimization. In Proc. of the IEEE conference on computer vision and pattern recognition (CVPR) , pages 12360-12370, 2022. 3

Zhuang Liu, Xuanlin Li, Bingyi Kang, and Trevor Darrell. Regularization matters in policy optimization-an empirical study on continuous control. Proc. the International Conference on Learning Representations (ICLR) , 2021. 3

Clare Lyle, Mark Rowland, and Will Dabney. Understanding and preventing capacity loss in reinforcement learning. Proc. the International Conference on Learning Representations (ICLR) , 2022. 1

Clare Lyle, Zeyu Zheng, Evgenii Nikishin, Bernardo Avila Pires, Razvan Pascanu, and Will Dabney. Understanding plasticity in neural networks. Proc. the International Conference on Machine Learning (ICML) , 2023. 1, 3, 6

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-3698, 2020. 3

Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample efficient world models. arXiv preprint arXiv:2209.00588 , 2022. 8

Volodymyr Mnih, Koray Kavukcuoglu, 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

Clara Na, Sanket Vaibhav Mehta, and Emma Strubell. Train flat, then compress: Sharpnessaware minimization learns more compressible models. In Proc. of the Conference on Empirical Methods in Natural Language Processing (EMNLP) , pages 4909-4936, 2022. 19

Evgenii Nikishin, Max Schwarzer, Pierluca D'Oro, Pierre-Luc Bacon, and Aaron Courville. The primacy bias in deep reinforcement learning. In Proc. the International Conference on Machine Learning (ICML) , pages 16828-16847. PMLR, 2022. 1, 2, 3, 6, 16, 17, 18

Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXiv:1807.03748 , 2018. 3

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

Roberta Raileanu, Max Goldstein, Denis Yarats, Ilya Kostrikov, and Rob Fergus. Automatic data augmentation for generalization in deep reinforcement learning. arXiv preprint arXiv:2006.12862 , 2020. 3

Herbert Robbins and Sutton Monro. A Stochastic Approximation Method. The Annals of Mathematical Statistics , 22(3):400 - 407, 1951. 4, 19

Tom Schaul, John Quan, Ioannis Antonoglou, and David Silver. Prioritized experience replay. Proc. the International Conference on Learning Representations (ICLR) , 2016. 3

Max Schwarzer, Ankesh Anand, Rishab Goel, R Devon Hjelm, Aaron Courville, and Philip Bachman. Data-efficient reinforcement learning with self-predictive representations. In Proc. the International Conference on Learning Representations (ICLR) , 2020. 3, 5

Max Schwarzer, Nitarshan Rajkumar, Michael Noukhovitch, Ankesh Anand, Laurent Charlin, R 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. 8

Younggyo Seo, Kimin Lee, Stephen L James, and Pieter Abbeel. Reinforcement learning with action-free pre-training from videos. In Proc. the International Conference on Machine Learning (ICML) , pages 19561-19579. PMLR, 2022. 8

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 on Machine Learning (ICML) , pages 2217-2225. PMLR, 2016. 2, 3, 6

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

Ghada Sokar, Rishabh Agarwal, Pablo Samuel Castro, and Utku Evci. The dormant neuron phenomenon in deep reinforcement learning. Proc. the International Conference on Machine Learning (ICML) , 2023. 3

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. 3

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 preprint arXiv:1801.00690 , 2018. 5

Hado Van Hasselt, Arthur Guez, and David Silver. Deep reinforcement learning with double q-learning. In Proc. the AAAI Conference on Artificial Intelligence (AAAI) , volume 30, 2016. 3

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 (NeurIPS) , 32, 2019. 16

Kaixin Wang, Bingyi Kang, Jie Shao, and Jiashi Feng. Improving generalization in reinforcement learning with mixture regularization. Proc. the Advances in Neural Information Processing Systems (NeurIPS) , 33:7968-7978, 2020. 3

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 on Machine Learning (ICML) , pages 1995-2003. PMLR, 2016. 3

Kaiyue Wen, Tengyu Ma, and Zhiyuan Li. How sharpness-aware minimization minimizes sharpness? In Proc. the International Conference on Learning Representations (ICLR) , 2023. 19

Dongxian Wu, Shu-Tao Xia, and Yisen Wang. Adversarial weight perturbation helps robust generalization. Advances in Neural Information Processing Systems , 33:2958-2969, 2020. 3, 19

Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous control: Improved data-augmented reinforcement learning. arXiv preprint arXiv:2107.09645 , 2021. 3, 22

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

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. 3

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 (NeurIPS) , 34:25476-25488, 2021. 8

Tao Yu, Cuiling Lan, Wenjun Zeng, Mingxiao Feng, Zhizheng Zhang, and Zhibo Chen. Playvirtual: Augmenting cycle-consistent virtual trajectories for reinforcement learning. Proc. the Advances in Neural Information Processing Systems (NeurIPS) , 34:5276-5289, 2021. 7, 8

Tao Yu, Zhizheng Zhang, Cuiling Lan, Yan Lu, and Zhibo Chen. Mask-based latent reconstruction for reinforcement learning. Proc. the Advances in Neural Information Processing Systems (NeurIPS) , 35:25117-25131, 2022. 3

Tao Yu, Zhizheng Zhang, Cuiling Lan, Yan Lu, and Zhibo Chen. Mask-based latent reconstruction for reinforcement learning. Proc. the Advances in Neural Information Processing Systems (NeurIPS) , 35:25117-25131, 2022. 7, 8

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 robot learning , pages 1094-1100. PMLR, 2020. 9

Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. arXiv preprint arXiv:1710.09412 , 2017. 3

Hattie Zhou, Ankit Vani, Hugo Larochelle, and Aaron Courville. Fortuitous forgetting in connectionist networks. Proc. the International Conference on Learning Representations (ICLR) , 2022. 2

Appendix

A Implementation Details

A.1 Computation

For all experiments, we used an NVIDIA RTX 3090 GPU for neural network training and a 32core 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 a single GPU where PLASTIC and PLASTIC † took 3 and 1.5 hours, which effectively translates to an average of 1 and 0.5 hours per run.

For the DeepMind Control Suite (DMC), runtimes varied based on the individual environment due to 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 a color image of size 32x32.

Input Adaptation: In 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.

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 a 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-connected 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 a 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), SharpnessAware 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, at 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'Oro et al [19], including the application of a 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 (<=2), they played a significant role when scaling the replay ratio with the reset mechanism [52]. A detailed hyperparameter is described in Table 5.

In the case of LN, we apply layer normalization after each convolutional and fully-connected layer. For CReLU, we replaced all ReLU activations with concatenated ReLU where the input 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].

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, L2, and Reset models described in Table 6, we adopted the results from Nikishin et al [52]. For the LN model, we incorporated layer normalization after each convolutional and fully-connected layer. Additionally, we 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].

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 to 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 a flat minimum, SAM has been successfully applied in various settings including computer vision [14], natural language processing [6], meta-learning [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-stationary data distribution in which the efficacy of SAM has not yet been verified. Indeed, we empirically verified that the usage of SAM enhances the adaptability of deep RL agents.

Now, we elaborate on what SAM is. Originally, SAM aims to train a model which is robust to adversarially perturbed model weights, by solving the following bi-level optimization problem on the loss function L :

This problem implies minimization of loss value around an ℓ 2 -ball of radius ρ , which encourages a model parameter w to find a smoother region of the loss landscape. Here, ρ > 0 is a 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 ,

We call the step of computing ϵ ∗ ρ ( w ) as the perturbation step of SAM. As a 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], or Adam [35], using the SAM gradient at current weight w t : ∇ w L ( w t + ϵ ∗ ρ ( w t )) . 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 to most of the applications of SAM, it is natural to randomly sample a mini-batch B t of m transitions from replay buffer and apply SAM update (often called m -SAM) using the stochastic gradient ∇ w L B t ( w t + ϵ ∗ ρ, B t ( w t ) ) , where

There are several considerations for applying SAM to RL agents. When applying SAM to the Rainbow agent, the random noises from noisy layers embedded in 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 we

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 to 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 to 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 to learn discrete control problems of the Atari-100k benchmark. Although there are many other components inside the Rainbow agent, such as multi-step learning, distributional RL, and dueling Deep Q-Network (DQN), it would be confusing to integrate all details together in the pseudocode. Thus, for simplicity of the display, we provide a pseudocode of applying SAM to vanilla DQN with noisy layers. See Algorithm 1.

Algorithm 1: Deep Q-Learning with Noisy Layers & SAM

B.3 Soft Actor-Critic with SAM

We also provide a pseudocode for applying SAM to the SAC algorithm. See Algorithm 2.

B.4 Ablations on the usage of SAM

We conduct some ablation studies aiming to answer the following questions:

How should we deal with noisy layers during training?

What if we solely apply SAM perturbation to the backbone or head of Rainbow agent?

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-100k benchmark, whereas for the last question we test on the DMC-M benchmark with the SAC agent.

(

i

)

(

ii

)

(

iii

)

0.20

0.25

0.30

0.35

0.40

0.45

0.50

0.55

0.60

IQM of HNS

SAM

SAM+Reset

Base

Reset

No reset

Reset

0.20

0.25

0.30

0.35

0.40

0.45

0.50

0.55

0.60

IQM of HNS

Abl: backbone-/head-SAM

B:Adam | H:Adam

B:SAM   | H:Adam

B:Adam | H:SAM

B:SAM   | H:SAM

Base

Only Actor Only Critic

SAM

150

200

250

300

350

IQM of HNS

Abl: Which networks to use SAM

Base

SAM, Only Actor

SAM, Only Critic

SAM

(a) Noisy layers on/off.

(b) Backbone/Head SAM.

(c) Actor/Critic 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 it on at the update step. Despite the variations, the overall impact on the agent's performance was relatively consistent. Note that option (i) is the most naïve approach which can be implemented by simply adding a 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 in 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 is 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, we focus on investigating the impact of SAM on actor and critic parameters. We conduct an ablation study utilizing the identical configuration as the DMC experiment presented in Section 4.2. Figure 5c demonstrates that exclusively applying SAM to critic parameters is more critical than solely applying SAM to actor parameters while using both is much better.

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

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.03

ρ

=0.1

ρ

=0.3

0.20

0.25

0.30

0.35

0.40

0.45

0.50

0.55

0.60

IQM of HNS

Ablation: SAM param (

ρ

)

SAM

SAM+Reset

Base

Reset

(a) SAM parameter

ρ

.

bs=16

bs=32

bs=64

0.20

0.25

0.30

0.35

0.40

0.45

0.50

0.55

0.60

IQM of HNS

Ablation: Batch size

SAM

SAM+Reset

Base

Reset

(b) Batch sizes.

SAM parameter ( ρ ). Being the main hyperparameter of SAM, the ρ > 0 represents the intensity of SAM perturbation. In every other experiment with DER and DrQ(-Rainbow), we do not tune the value of ρ but fix them as 0.1. Here, however, we run two other values of ρ 's, 0.03 and 0.3, to examine the effect of ρ . The result is shown in Figure 6a. We definitely observe the trade-off due to the size of ρ . If the ρ 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, a 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 ρ 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 (SAM+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

D Broader Impact

Our research possesses broader impacts in two 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 on inclusivity and accessibility in RL. By decreasing data and computational demands, this work can aid underprivileged communities with fewer resources to participate in RL research and benefit from its many applications. This approach encourages a diverse range of perspectives and experiences in the field, enriching the community.

However, while our 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 or harm that could arise from these advancements.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "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\nSTRENGTHS:\nThe 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\nWEAKNESSES:\nThere 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\nQUESTIONS:\nThe 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\nLIMITATIONS:\nWhile 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-04"}

comments: 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.

STRENGTHS:
- This paper is clear, well-written, self-contained, and enjoyable to read.
- While the proposed method simply combines two existing approaches, such a combination is novel and well justified by the synthetic experiments and ablation studies.
- 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.

QUESTIONS:
Apart from the points raised above, I have some questions and comments, and would like to hear the authors' feedback:
- Table 1: Some entries missing ("xxx" in the table).
- Line 128: The sentence ("..., which encourages to Here, ...") is confusing.
- 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.
- 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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-05

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThis 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.\n\nWEAKNESSES:\nThis 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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nSample 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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?\n\nQUESTIONS:\n- 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\nLIMITATIONS:\nThe authors are forthcoming with the main limitations of this work, and explained them clearly.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-07"}