SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Hojoon Lee* Hanseul Cho* Hyunseung Kim* Daehoon Gwak Joonkee Kim Jaegul Choo Se-Young Yun Chulhee Yun
Kim Jaechul Graduate School of AI, KAIST {joonleesky, jhs4015, mynsng, daehoon.gwak, joonkeekim, jchoo, yunseyoung, chulhee.yun}@kaist.ac.kr
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.
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].
Figure 1: Scaling of our proposed method, PLASTIC, on the Atari-100k benchmark compared to state-of-the-art methods (Section 4.3).
Figure 2: Left: Performance of Synthetic Experiments. Layer Normalization (LN) and Sharpeness-Aware Minimization (SAM) considerably enhance input plasticity, while their effect on label plasticity is marginal. Conversely, Concatenated ReLU (CReLU) and periodic reinitialization (Reset) predominantly improve label plasticity with subtle benefits on input plasticity. Right: Performance of RL Benchmarks. PLASTIC consistently outperforms individual methods, highlighting the synergistic benefits of its integrated approach.
Another line of investigation emphasizes the role of gradient propagation, highlighting that as 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.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 sharpness-aware 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 Temporal-Difference (TD) error which is defined as:
$$
\mathcal {L} (\boldsymbol {w}, \boldsymbol {w} ^ {-}, \tau) = \left[ Q _ {\boldsymbol {w}} (s, a) - \left(r + \gamma \max  _ {a ^ {\prime}} Q _ {\boldsymbol {w} ^ {-}} \left(s ^ {\prime}, a ^ {\prime}\right)\right) \right] ^ {2} \tag {1}
$$
where  $\boldsymbol{w}$  denotes the model weights,  $\boldsymbol{w}^{-}$  is the weights of the target model, and  $\tau = (s,a,r,s^{\prime})$  is the transition sampled from the replay buffer  $\mathcal{B}$ .
Noisy Networks [22] introduce stochasticity into the model weights, reparameterizing them as:
$$
\boldsymbol {w} = \boldsymbol {\mu} + \boldsymbol {\sigma} \cdot \tilde {\varepsilon} \tag {2}
$$
where  $\mu$  and  $\sigma$  are the learnable weights, whereas  $\tilde{\varepsilon}$  is a random vector sampled from an isotropic Gaussian distribution, adding randomness to the layer. (We denote element-wise product by  $\cdot$ .)
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  $\pi$  and a critic model  $Q$ , each parameterized by weights  $\theta$  and  $w$  respectively.
In training, the critic  $Q_{w}$  is trained to minimize the following objective, defined as
$$
\mathcal {L} _ {Q} (\boldsymbol {\theta}, \boldsymbol {w}; \tau) = \left[ Q _ {\boldsymbol {w}} (s, a) - \left(r + \gamma \left(Q _ {\boldsymbol {w}} \left(s ^ {\prime}, a ^ {\prime}\right) - \alpha \log \pi_ {\boldsymbol {\theta}} \left(s ^ {\prime}, a ^ {\prime}\right)\right)\right) \right] ^ {2}, \quad a ^ {\prime} \sim \pi_ {\boldsymbol {\theta}} (\cdot | s ^ {\prime}) \tag {3}
$$
where  $\alpha$  is the entropy coefficient and  $\tau = (s, a, r, s')$  is the transition sampled from the buffer  $\mathcal{B}$ .
Subsequently, the policy  $\pi_{\theta}$  is jointly trained to maximize the following objective:
$$
\mathcal {L} _ {\pi} (\boldsymbol {\theta}, \boldsymbol {w}; s) = \mathbb {E} _ {a \sim \pi_ {\boldsymbol {\theta}}} [ (Q _ {\boldsymbol {w}} (s, a) - \alpha \log \pi_ {\boldsymbol {\theta}} (a | s) ]. \tag {4}
$$
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
Figure 3: Comparative performance improvements of various methods and their combinations on input and label adaptation tasks. The left panel showcases individual method performance, while the right panel delves into the synergistic benefits of combining multiple methods.
(SAM) optimizer for the former, and Concatenated Rectified Linear Unit (CReLU) activation and Reset for the latter. For the SAM optimizer, the perturbation parameter was tuned across the set  $\{0.1, 0.03, 0.01\}$ , and for the Reset method, the reset interval was tuned over  $\{5, 10, 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.
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).
Atari-100k. Following the standard evaluation protocol from [32, 58], we evaluate the performance of 26 games on Arcade Learning Environments [73], limited to  $100\mathrm{k}$  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  $\rho$  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  $(\lambda_{\mathrm{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  $\mathrm{HNS} = \frac{\text{agent\_score - random\_score}}{\text{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.
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 plasticity-preserving methods. For this purpose, we used the Atari-100k and DMC-Medium benchmarks.
Table 1: Performance on Atari-100k and DMC-M. The results are averaged over 10 random seeds.
Method	Atari-100k	DMC-M
IQM	Median	Mean	IQM	Median	Mean
Base	0.258 (0.224, 0.292)	0.277 (0.209, 0.295)	0.476 (0.432, 0.520)	213 (146, 304)	288 (223, 356)	293 (239, 347)
SAM [21]	0.325 (0.296, 0.354)	0.327 (0.284, 0.368)	0.501 (0.465, 0.537)	278 (196, 371)	341 (264, 404)	332 (277, 389)
LN [47]	0.259 (0.235, 0.293)	0.247 (0.218, 0.276)	0.463 (0.403, 0.522)	412 (332, 488)	415 (351, 491)	421 (366, 477)
CReLU [2]	0.256 (0.224, 0.287)	0.193 (0.176, 0.252)	0.498 (0.444, 0.553)	338 (133, 566)	398 (235, 549)	394 (267, 522)
Reset [52]	0.343 (0.314, 0.373)	0.291 (0.231, 0.369)	0.660 (0.611, 0.715)	514 (434, 590)	491 (430, 568)	498 (442, 555)
PLASTIC	0.421 (0.388, 0.457)	0.347 (0.266, 0.422)	0.933 (0.812, 1.067)	565 (498, 626)	540 (476, 599)	537 (487, 587)
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.
Figure 4: Left: The maximum eigenvalue of Hessian  $(\lambda_{\mathrm{max}}(\nabla^2\mathcal{L}))$ , representing the curvature of the loss landscape. Right: Fraction of active units in the head layers. Both metrics were derived from the DrQ algorithm's evaluation on the Atari-100k benchmark.
To delve deeper into the mechanisms driving these performance enhancements, we focused on two specific metrics: the maximum eigenvalue of the Hessian  $(\lambda_{\mathrm{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.