SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Evgenii Nikishin* Junhyuk Oh Georg Ostrovski Clare Lyle Razvan Pascanu Will Dabney André Barreto DeepMind
A growing body of evidence suggests that neural networks employed in deep reinforcement learning (RL) gradually lose their plasticity, the ability to learn from new data; however, the analysis and mitigation of this phenomenon is hampered by the complex relationship between plasticity, exploration, and performance in RL. This paper introduces plasticity injection, a minimalistic intervention that increases the network plasticity without changing the number of trainable parameters or biasing the predictions. The applications of this intervention are two-fold: first, as a diagnostic tool — if injection increases the performance, we may conclude that an agent's network was losing its plasticity. This tool allows us to identify a subset of Atari environments where the lack of plasticity causes performance plateaus, motivating future studies on understanding and combating plasticity loss. Second, plasticity injection can be used to improve the computational efficiency of RL training if the agent has to re-learn from scratch due to exhausted plasticity or by growing the agent's network dynamically without compromising performance. The results on Atari show that plasticity injection attains stronger performance compared to alternative methods while being computationally efficient.
"You cannot teach an old dog new tricks" an old proverb says. While the common wisdom is not necessarily a source of absolute truth, neuroscientists recognized a long time ago that biological agents indeed gradually lose adaptability with age [Livingston, 1966]. This phenomenon is referred to as loss of plasticity in brains [Nelson, 1999, Mateos-Aparicio and Rodríguez-Moreno, 2019] and happens for multiple reasons, including natural degradation of neurons and their connections [Mahncke et al., 2006, Kolb and Gibb, 2011].
Since the biological causes of loss of plasticity do not apply to artificial agents, in principle there is no reason to expect that this phenomenon also happens in the context of machine learning. Surprisingly, several recent works show that reinforcement learning (RL) agents that use neural networks may gradually lose the ability to learn from new experiences [Dohare et al., 2021, Lyle et al., 2022, Nikishin et al., 2022].
The precise mechanisms causing loss of plasticity in RL are not well understood. The problem is particularly challenging to study in this context because performance in RL is influenced by many factors. For example, an agent without plasticity issues may still struggle to learn if it fails to properly explore the environment [Taïga et al., 2019]. Past literature focused on using and controlling proxy measures of plasticity such as the number of saturated rectified linear units [Nair and Hinton, 2010] and feature rank [Kumar et al., 2021], but it is unclear how well these measures manage to capture the underlying phenomenon [Gulcehre et al., 2022].
This paper complements past evidence about the existence of plasticity loss in deep RL and introduces plasticity injection, an intervention that augments plasticity of the agent's neural network. The
*Work done during the internship; currently at Mila, Université de Montréal.
conceptual idea is simple: at any point in training, one can freeze the current network and create a new one that is going to be learning a change to the predictions, whilst ensuring that the change is initially zero. Crucially, plasticity injection does not increase the number of trainable parameters and does not affect the network's predictions when it is applied. Because of these properties, the intervention enables careful analysis of the plasticity loss phenomenon in RL while keeping other confounding factors aside.
We suggest two uses of plasticity injection, one as an analytic tool and another as a practical algorithmic technique. For analysis, we propose an experimental protocol that uses plasticity injection for diagnosing the problem of plasticity loss: for example, if an agent that was struggling to improve its behavior escapes a performance plateau after the intervention, we can conclude that the agent had been experiencing problems with its network plasticity. Using this protocol in the Arcade Learning Environment [Bellemare et al., 2013], we identify scenarios where loss of plasticity hinders the learning process. Furthermore, based on the intervention-enabled analysis, we provide recommendations for controlling the degree of plasticity loss.
We also propose to use plasticity injection as a way to improve computational efficiency of RL training in the following scenarios. First, when the agent loses its plasticity because its network turns out to be too small, plasticity injection can be dynamically used to increase the capacity of the agent without having to re-train the agent with a larger network from scratch. We empirically show that our method improves the aggregate score across 57 Atari games by  $20\%$  compared to other methods for dynamically addressing plasticity loss. Second, plasticity injection can be used to minimize computation by switching from a small network to a larger network in the middle of training without compromising the performance compared to using the larger network from scratch; we also empirically verify it on Atari games.
1. A minimalistic intervention called plasticity injection that increases plasticity of the agent while preserving the number of trainable parameters and not affecting its predictions;
2. Complementary evidence about the existence of the plasticity loss phenomenon in deep RL;
3. An experimental protocol for diagnosing loss of plasticity using the intervention;
4. A way to improve computational efficiency of RL training by dynamically expanding a neural network used by the agent.
Plasticity in Continual Learning. Discussions about plasticity of neural networks date back (at least) to a seminal paper by McCloskey and Cohen [1989] outlining the plasticity-stability dilemma, a trade-off between preserving performance on previous tasks and maintaining adaptability to future ones. The continual learning community historically put a higher emphasis on the stability aspect, addressing catastrophic forgetting of past behaviors [French, 1999]. Recently, several works raised awareness of difficulties with learning on future tasks too. Ash and Adams [2020] demonstrated an instance of loss of generalization, when pre-training a network might unrecoverably damage generalization even if pre-training was done on a uniform subsample of the same dataset. Berariu et al. [2021] deepened the study and conjectured that the phenomenon might happen because of the reduction of gradient noise when warm-starting the network. Dohare et al. [2021] explicitly study the network plasticity in continual learning and demonstrate the reduced ability to minimize even the training error as the number of tasks increase. These works build an understanding of the problem by studying simplified settings that isolate different aspects of learning capabilities in continual learning, whereas our work aims at tackling the deep RL setting in its whole.
Loss of Plasticity in Deep RL. Issues with plasticity and related phenomena have been recently highlighted in deep RL under a plethora of different names. Lyle et al. [2022] show loss of capacity for fitting targets in online RL and Kumar et al. [2021] demonstrate a related implicit underparameterization phenomenon caused by bootstrapping with more emphasis on the offline RL case. Both of these works use the feature rank as a proxy measure for plasticity but later Gulcehre et al. [2022] question the reliance on such measure by demonstrating a weak correlation between the rank and the agent's performance, partially motivating our study that focuses directly on agent's performance to reason about plasticity. Works of Sokar et al. [2023] and Abbas et al. [2023] focus on
Figure 1: Demonstration of plasticity loss in a sequence of policy evaluation tasks. The task (a policy to evaluate) changes every 1000 iterations. The reset every task setting shows that newly-initialized parameters are able to fit each task, whereas the reset never setting shows the diminishing capability to fit the data when using the trained parameters from one task as the initialization for another task.
saturation of neurons over the course of training, but Lyle et al. [2023] demonstrate that the saturation alone cannot fully characterize the plasticity loss phenomenon. Nikishin et al. [2022] discuss the primacy bias in deep RL, a tendency to excessively train on early data damaging further learning progress, and propose to periodically reset a part of the network to address the issue while relying on the replay buffer as a knowledge transfer mechanism. Earlier, Igl et al. [2021] had observed that deep RL agents can lose the ability to generalize due to non-stationarity and proposed to use distillation as a mitigation mechanism. Plasticity injection closely relates to these approaches by leveraging newly initialized weights, but does not require re-training and directly continues learning.
Architectures. The works above mainly discuss algorithmic aspects with less focus on the network architecture, although it is also an important component of the agent's design [Mirzadeh et al., 2022]. The closest work in this space is about progressive networks [Rusu et al., 2016] that considers a setting with multiple environments and adds a new network with cross-connections to the layers of previous networks. A network after plasticity injection can be viewed as a simplified version of the architecture with a motivation of increasing plasticity within a single task without affecting agent's predictions. A line of work on the mixture of experts [Shazeer et al., 2017] and modular networks [Andreas et al., 2016] is also related, but the focus in these papers typically is compositionality or handling multiple modalities. The idea of growing network layers or neurons has also been investigated [Fahlman and Lebiere, 1989, Chen et al., 2015]; plasticity injection belongs to a family of these methods up to the differences in the growing strategy and in explicitly controlling for the number of trainable parameters. In the context of language modeling, Hu et al. [2022] explored a similar idea of freezing a pre-trained model and fine-tuning a low-rank addition to the weight matrices on a downstream task. Lastly, plasticity injection can be conceptually viewed as an instance of residual learning [He et al., 2016] and boosting [Schapire, 1990].
3 An Illustration of Plasticity Loss
Plasticity of a neural network is broadly defined as the ability to learn from new experiences. To provide intuition on how this ability can decrease over time, we present a didactic example before investigating the case with deep RL. Figure 1 shows the mean-squared error (MSE) on a sequence of supervised policy evaluation problems derived from the Up n Down Atari environment. We first trained an agent on this environment for 200M frames and stored the policies occurring at every 10M frames. Then, for each stored policy, we sampled states from the corresponding stationary distribution and computed Monte-Carlo estimates of the value function for each state, resulting a training set composed of states and their values. We then trained a network to solve the resulting sequence of prediction problems. This sequence of related prediction problems differing in the input and target distribution aims to reproduce the scenario faced by an online RL agent [Dabney et al., 2021]. The curve labeled "reset never" corresponds to starting each prediction problem using the
final parameters from the previous one, while "reset every task" corresponds to randomly initializing the network parameters at every prediction problem.
The conventional wisdom about transfer learning suggests that, if two tasks are related, pre-training on the first might accelerate learning on the second [Pan and Yang, 2009]. Here we observe the opposite trend: it takes longer and longer for the network to decrease training error on the subsequent policy evaluation problems if its parameters are not re-initialized. This example gives a simple demonstration of how plasticity loss can occur; we refer to the work by Dohare et al. [2021] for an in-depth study of the phenomenon in the continual setting.
After building intuition about loss of plasticity, we turn our attention to its analysis in deep RL. The key distinctive feature of RL is the presence of an exploration confounder: in contrast to the continual setting with a fixed sequence of datasets, an RL agent influences the future data it learns from. Thus, a failure of an RL system can be attributed not only to loss of plasticity but also to inability to explore. The next section presents a strategy to increase plasticity of an agent that addresses this difficulty with the analysis.
4 Plasticity Injection
Before describing the experimental design in detail, we list the motivating desiderata:
- Unaffected predictions: the agent's predictions should stay the same after the intervention to avoid abrupt changes. This criterion allows isolating confounding factors related to exploration;
- Preserving the trainable parameter count: the intervention should not affect the number of trainable parameters to minimize confounding factors related to representational capacity.
We now present the proposed intervention to increase plasticity of an RL agent. First, let us denote the neural network approximator employed by the agent (for example, used for action-value prediction) as  $h_{\theta}(x)$ , where  $\theta$  indicates the parameters. At some point in training, where the network might have started losing plasticity, we are going to freeze the parameters  $\theta$  and introduce a new set of parameters  $\theta'$  sampled from random initialization. The key idea is to keep two copies of  $\theta'$ , which we denote by  $\theta_1'$  and  $\theta_2'$ ; while  $\theta_1'$  are free parameters used to learn a residual to the old network outputs,  $\theta_2'$  remains frozen throughout. The agent's predictions after plasticity injection will be calculated using the following expression:
$$
\underbrace {h _ {\theta} (x)} _ {\text {f r o z e n}} + \underbrace {h _ {\theta_ {1} ^ {\prime}} (x)} _ {\text {t r a i n e d}} - \underbrace {h _ {\theta_ {2} ^ {\prime}} (x)} _ {\text {f r o z e n}}. \tag {1}
$$
Since initially  $\theta_1' = \theta_2'$ , immediately after plasticity injection the predictions of the neural network remain unaltered. As learning progresses,  $\theta_1'$  deviates from  $\theta_2'$  and  $h_\theta(x) - h_{\theta_2'}(x)$  serves as a bias term for predictions.
Note that if we apply plasticity injection to all parameters of the network, the new network will have to re-learn the representations encoded in  $h_{\theta}(\cdot)$  from scratch. Thus, we apply our intervention to only a subset of the parameters and explain the idea further with a slight abuse of notation. We schematically split the network into an encoder  $\phi(\cdot)$ , that denotes a mapping induced by first  $k$  layers of the network, and a head  $h_{\theta}(\cdot)$  where  $\theta$  now refers to parameters of the remaining layers of the network. After this relabelling, we can apply the intervention to  $h_{\theta}(\cdot)$  as outlined above. Appendix B later presents an ablation of sharing the encoder.
Figure 2 illustrates the strategy to apply plasticity injection. Note that gradients from the frozen heads affect the encoder too, i.e. we do not stop the gradient propagation from any of the components of the output. It is worth noting that the proposed intervention increases the total number of parameters of the network (but keeps the same number of trainable parameters), which in turn may result in an increase of training time. However, we later discuss in Section 5.3 how plasticity injection can save computational resources.
The idea of learning with newly-initialized last layers has been explored by Nikishin et al. [2022], who suggested resetting the corresponding parameters of the network at fixed intervals and used the replay buffer [Lin, 1992] to re-learn after resets. Their experimental evidence supports the hypothesis that resets mitigate plasticity loss. However, resetting parameters of the network abruptly changes
Figure 2: An illustration of the architecture before and after plasticity injection. Before the intervention, the network is schematically separated into an encoder  $\phi(\cdot)$  and a head  $h_\theta(\cdot)$ , both parts are learning. After plasticity injection, we freeze the parameters  $\theta$  of the head (we use red to indicate parameters that are not updated in the illustration) and create two copies of a randomly initialized parameters  $\theta'$ : one frozen and one unfrozen. The output of the agent is obtained by first passing the input  $x$  to the encoder  $\phi(\cdot)$ , next passing  $\phi(x)$  to all three heads, and finally combining the heads' outputs according to Expression (1).
its predictions, which results in a temporary decrease in performance and induces an exploration effect. From an analysis perspective, these abrupt changes make it more difficult to isolate the effect of additional plasticity on the agent's performance. From a practical perspective, plasticity injection does not rely on the buffer; Section 5.3 demonstrates how this difference can be critical.
This section presents results for two main applications of plasticity injection: as a tool for diagnosing plasticity loss and as a way to dynamically grow the network to efficiently use computations.
The baseline agent is Double DQN [Van Hasselt et al., 2016] learning for 200M interactions on a standard set of 57 Atari games from the Arcade Learning Environment benchmark [Bellemare et al., 2013]. The choice of Double DQN is motivated by the relative robustness and stronger performance with double Q-learning [Van Hasselt, 2010] compared to the vanilla DQN agent [Mnih et al., 2015] as well as simplicity compared to later DQN-based agents such as Rainbow [Hessel et al., 2018].
The majority of the experiments use a single plasticity injection after 50M frames; otherwise, we explicitly specify the number and timesteps of injections. Appendix B discusses ablations on the design choices when using plasticity injection. A convolutional neural network employed by the Double DQN agent consists of 5 layers. The encoder corresponds to the first three of them (hence  $k = 3$ ), while the head refers to the last two. Since DQN-based agents employ a target copy of the network parameters, we perform the same interventions on them.
For reliable evaluation of the performance across environments, we adopt the protocol of Agarwal et al. [2021] with a focus on the interquartile mean (IQM). All experiments use 3 random seeds.
5.2 Plasticity Injection as a Diagnostic Tool
Consider the task of improving a deep RL system when an agent performs suboptimally. Practitioners know how non-trivial is the process of pinpointing exact reasons why an agent might be struggling to improve the behavior. One of the reasons, as we discussed, can be loss of network plasticity throughout training.
We view the proposed intervention as a tool that can provide insight when analyzing deep RL systems. The procedure for using it is as follows: when an agent is on a performance plateau or has a slower
Figure 3: A demonstration of diverse effects from plasticity injection applied to the Double DQN agent after 25M, 50M, and 100M frames on a selection of Atari games comprising two examples where the intervention improves the performance and two examples where it does not. The baseline in Space invaders and Phoenix demonstrates the diminishing performance improvements and the performance plateau respectively, whilst the agent after the injection is capable of achieving higher returns. The stalled performance in Assault is due to exploration challenges (see Appendix D for further details): adding plasticity could not alleviate them. If the agent does not show signs of the diminishing ability to learn, like in Robotank, the injection would not lead to improved performance. Varying the injection timestep allows identifying the moment when plasticity loss occurs. Results for all 57 environments are available in Figure 7.
learning progress, take a saved copy of the agent, perform plasticity injection, and compare the training curves with and without the intervention. This way we answer a counterfactual question: what could have been the agent's performance if the network had more plasticity?
Figure 3 gives a set of example behaviors after following the procedure. In Space invaders, the baseline agent keeps learning but the post-injection agent improves at a faster rate towards the end of learning; we might interpret the observation as an indication of decreasing network plasticity over the course of training. In Phoenix, we see a completely stalled performance and the intervention allows doubling the final returns; such an observation point at possible catastrophic loss of plasticity, where additional interactions do not translate to improved behavior. In Assault, on the other hand, the agent has plateaued but the injection does not make a difference. Further inspection revealed that around a score of 2800, the environment transitions to a new regime where an agent needs to start using an action that was not relevant before (see Appendix D for a visualization). This observation suggests that performance stagnation is related to exploration. In Robotank, the learning progress shows no signs of pathologies, giving evidence that the agent does not have problems with plasticity.
Plasticity injection can also demonstrate when loss of plasticity occurs. The post-intervention performance in Space invaders does not differ for varying injection timestep, suggesting that the agent might not start experiencing consequences of the lost plasticity until around 100M frames. On the other hand, in Phoenix, plasticity injection improves the performance earlier, implying that the agent lost its plasticity around 25M frames. Varying the injection timestep in Assault and Robotank does not change the performance significantly, supporting our previous conclusion about these games.
Figure 4 summarizes when and to which extent the Double DQN agent benefits from plasticity injection across 57 Atari games. The observations about improvements from injection complement evidence of the existence of plasticity loss in deep RL [Kumar et al., 2021, Lyle et al., 2022, Nikishin et al., 2022]. We note that the argument here is nuanced: since the notion of plasticity is defined broadly and is challenging to measure, it is our best interpretation that the post-intervention agent can learn further because it addressed plasticity issues. But because of an experimental design that strived to be careful, we believe that it is the most likely explanation.
What should we do after using the tool and observing loss of plasticity? Dohare et al. [2021], Nikishin et al. [2022], Gogianu et al. [2021] provide evidence that plasticity loss is strongly affected by the learning rate, the replay ratio (the number of gradient steps per an environment step), the network size $^{1}$ , and normalizations (such as spectral norm [Miyato et al., 2018]). We measure the sensitivity of the aggregate improvements of the final score from plasticity injection at 50M with respect to these choices of the agent specification. Results in Figure 5 are consistent with observations from previous
To make a network two times larger, we multiply the width of all hidden layers by  $\sqrt{2}$ .