SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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  $\theta_{t} = \alpha \theta_{t-1} + (1 - \alpha)\phi$ , where  $\phi \sim$  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
Table 2: Comparison to the SOTA on Atari-100k. For IRIS, DreamerV3, EfficientZero, PlayVirtual, MLR, 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	IQM	Median	Mean	OG
Model-Based	IRIS [49]	-	30.4	-	36.3	0.501	0.289	1.046	0.512
DreamerV3 [25]	-	17.9	-	5.1†	0.497	0.466	1.097	0.505
EfficientZero [75]	✓	8.4	-	28.0	n/a	1.090	1.943	n/a
Model-Free	PlayVirtual [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
8	10.2	0.589	0.560	1.188	0.452
PLASTIC†	-	6.8	2	0.5	0.396	0.425	0.702	0.541
4	1.0	0.518	0.517	0.858	0.478
8	1.9	0.583	0.542	0.939	0.448
PLASTIC	-	7.2	2	1.0	0.421	0.347	0.933	0.535
4	1.9	0.545	0.407	1.002	0.475
8	3.8	0.571	0.494	0.968	0.461
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  $\mathrm{PLASTIC}^{\dagger}$  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.
Table 3: Fine-tuning from a pretrained model. Reset† applies a soft reset to the backbone and a hard reset to the head.
RR	SAM	Reset†	IQM
2	✓	0.366 (0.324, 0.397)
✓	✓	0.709 (0.650, 0.745)
0.776 (0.703, 0.854)
4	✓	0.243 (0.214, 0.267)
✓	✓	0.780 (0.706, 0.865)
0.834 (0.769, 0.889)
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 plasticity-preserving 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,  $\mathrm{LN} + \mathrm{CReLU} > \mathrm{LN} \approx \mathrm{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
Table 4: Performance comparison of various plasticity-preserving method combinations on the Atari-100k benchmark.
LN	SAM	CReLU	Reset	IQM
0.258 (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)
✓	✓	0.396 (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)
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.
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 Progen [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 practical 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.
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).
[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, 6
[3] 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
[4] 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
[5] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.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] 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
[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 in reinforcement learning. Proc. the International Conference on Learning Representations (ICLR), 2022. 3
[11] 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
[12] 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
[13] 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
A Implementation Details 16
A.1 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 on SAM 19
B.2 Deep Q-Learning with SAM and Noisy Layers 20
B.3 Soft Actor-Critic with SAM 21
B.4 Ablations on the usage of SAM 22
C Per Environment Results 23
C.1 Atari-100k 23
C.2 DMC-M 23
D Broader Impact 26
A.1 Computation
For all experiments, we used an NVIDIA RTX 3090 GPU for neural network training and a 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 a single GPU where PLASTIC and PLASTIC $\dagger$  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  $32 \times 32$ .
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  $3 \times 3$ , 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), 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, 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].
Table 5: Hyperparameters on Atari 100k. The ones introduced by this work are at the bottom.
Hyperparameter	Value
State downsample size	(84, 84)
Grey scaling	True
Data augmentation	Random Shifts and Intensity Jittering
Random Shifts	± 4 pixels
Intensity Jiterring scale	0.05
Frame skip	4
Stacked frames	4
Action repeat	4
Training steps	100k
Update	Distributional Q
Dueling	True
Support of Q-distribution	51
Discount factor γ	0.99
Batch size	32
Optimizer (β1, β2, ε)	Adam (0.9, 0.999, 0.000015)
Learning rate	0.0001
Max gradient norm	10
Priority exponent	0.5
Priority correction	0.4 → 1
EMA coefficient (τ)	0.99
Exploration Network	Target
Exploration	Noisy nets
Noisy nets parameter	0.5
Replay buffer size	100k
Min buffer size for sampling	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
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].
Table 6: Hyperparameters on DMC. The ones introduced by this work are at the bottom.
Parameter	Setting
Gray-scaling	True
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
Backbone: latent dim	50
Head: n. hidden layers	2
Head: hidden units	256
Target network update period	1
EMA coefficient τ	0.995
Initial Temperature	0.1
Updates per step	1
Replay Buffer Size	1,000,000
Total training steps	2,000,000
Evaluation trajectories	100
SAM parameter ρ	Quadruped: 0.1
Others: 0.01
Reset interval	100,000
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  $\mathcal{L}$ :
$$
\min  _ {\boldsymbol {w}} \left\{\max  _ {\boldsymbol {\epsilon}: \| \boldsymbol {\epsilon} \| _ {2} \leq \rho} \mathcal {L} (\boldsymbol {w} + \boldsymbol {\epsilon}) \right\}. \tag {5}
$$
This problem implies minimization of loss value around an  $\ell_2$ -ball of radius  $\rho$ , which encourages a model parameter  $\pmb{w}$  to find a smoother region of the loss landscape. Here,  $\rho > 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  $\pmb{w}$ ,
$$
\pmb {\epsilon} _ {\rho} ^ {*} (\pmb {w}) := \frac {\rho}{\| \nabla_ {\pmb {w}} \mathcal {L} (\pmb {w}) \| _ {2}} \nabla_ {\pmb {w}} \mathcal {L} (\pmb {w}).
$$