PLASTIC: Improving Input and Label Plasticity for Sample Efficient Reinforcement Learning

Summary

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

Strengths

The main strength of the paper is a thorough empirical investigation: - 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). - 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). - The SAM+Resets combination is thoroughly compared to other methods that were proposed to address plasticity loss (such as layer normalization and concatenated ReLUs).

Weaknesses

There are several weak spots that prevent the reviewer from assigning a higher score: - 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. - 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. There are also several minor weaknesses: - Table 1 contains placeholder performance values 0.xxx for DER + CReLU - 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. Addressing the outlined limitations might increase the score.

Questions

The paper doesn’t study continuous control from proprietary states — what was the motivation for excluding the setting? The 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.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

2 fair

Limitations

While the paper presents evidence about the effects of SAM+Resets on network adaptability, an explanation is lacking: why SAM and Resets enhance input and label adaptability and, for example, not visa versa? Likewise, 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? Diving deeper into explaining the phenomena would increase the significance of the paper.

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.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

The limitations of the proposed method have been adequately discussed in Section 6.

Summary

This work studies the role of generalization and plasticity in sample-efficient deep RL. This paper proposes sharpness-aware minimization (SAM) to improve generalization in RL. And it provides details on how to use SAM with deep RL algorithms like SAC and Rainbow. Empirical evaluation shows that combined usage of SAM and periodically resetting the last few layers of the network improves sample efficiency in the Atari-100k benchmark. On the other hand, adding SAM does not provide any benefit over resetting in Deepmind Control Suite.

Strengths

This work tackles the critical problem of sample efficient reinforcement learning. Sample efficiency is particularly important for applications where exploration is risky and expensive, such as healthcare. The 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. The 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. The results on the Atari-100k benchmark show that combining SAM with resets improves the performance over just using SAM or resets. This exciting result shows that SAM can improve the sample efficiency of deep RL.

Weaknesses

This paper contains interesting ideas and promising experimental results. The paper's central claim is that generalizability and plasticity constitute different roles in improving performance (lines 6-7). However, the experiments performed in the paper do not adequately test this claim. There are significant issues with the experiments in Sections 4 and 5.2. The 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: 1. 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. The 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. 2. 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. 3. 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. Similar problems exist with the experiments in Section 5.2. 1. 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. 2. 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. Besides 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.

Questions

1. The paper uses two terms, "adaptability" and "plasticity." To me, these words have very similar meanings. Plasticity means the ability to learn from new data, and adaptability means being able to fit new data. Can you please explain in what sense are you using these words and what is the difference between the two? If they have a similar meaning, using just one word in the paper might be better. Currently, the use of these two words causes confusion in section 4. Maybe in the conclusion of Section 4, it is better to say that SAM improves plasticity when there is input non-stationarity and resetting the last few layers improves plasticity when the labels change. 2. What is the confidence interval reported in Figure 2? Is it the 95% bootstrapped confidence interval?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

2 fair

Presentation

2 fair

Contribution

3 good

Limitations

The key limitation of the paper is that the experiments did not test the central claims of the paper. To claim that both plasticity and generality are important for improving adaptability, methods that inject plasticity into the whole network should be tested. Including the following baselines in Section 4 with proper hyper-parameter turning and more random seeds will significantly improve the completeness of the claims. 1. CReLU 2. LN 3. L2 (properly tuned) + Reset 4. CReLU + Reset 5. LN + Reset Similarly, the following baselines in Section 5.2 need to be added. 1. CReLU + Reset 2. LN + Reset 3. L2 + Reset 4. L2 + best combination of CReLU, LN, Reset 5. SAM + best combination of CReLU, LN, Reset The 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. EDIT: I have update my score based on the new experiments and changes in the main message of the paper.

Summary

Sample efficiency in RL is desirable to reduce computational and data collection costs, and is particularly critical to data-limited domains. While 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. This 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. The 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. Using 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. These improvements in adaptability enables off-policy methods to better utilize multiple updates on data, improving sample efficiency. The 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. The 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. The authors further perform ablation studies. Notably, 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. This 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.

Strengths

- The writing is clear, and figures are intuitive and well-presented. The authors' claims are modulated and supported by the experiments. The logic of the paper flows very well; many questions that I had while reading were quickly addressed in subsequent sections. - 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. - The solution that the authors propose is simple yet elegant, and can be widely applied as it requires little change to model architecture. Furthermore, beyond showing strong empirical results for their approach, the authors also provide a hypothesis on how the underlying principles of generalization and plasticity achieve a synergistic dynamic, and support it with synthetic experiments.

Weaknesses

- Both SAM and the reset mechanism are known techniques, and applying them in conjunction to achieve good results is not by itself sufficient for novelty. However, in the ablation studies, the authors also demonstrate that SAM + CRelu performs almost as well as SAM + reset. This strengthens their claim that it is the synergistic relationship between generalization and plasticity, not the specific method to improve these properties, that leads to the performance gains, which is novel. Yet, just testing one alternative combination does not seem sufficient to support the claim. - 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. - 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. - Several fields appear to be missing in Table 1?

Questions

- Have you tested other combinations of methods to improve generalization and plasticity? For me, the main contribution is not SAM + reset, but the insight that generalizability and plasticity play separate roles in improving adaptability, and combining methods that improve each of them can lead to further performance gains; SAM + reset is just one instantiation used to verify your hypothesis. Ideally, I'd like to see that the same pattern holds on other combinations beyond SAM + reset and SAM + CRelu. - 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?

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

Limitations

The authors are forthcoming with the main limitations of this work, and explained them clearly.

Feedback

Thank the authors for the detailed reply and additional experiments. I have one comment regarding Question 3.1 and expect the authors can take it into consideration when preparing the next version. The results show that "SAM + Resets" performs similarly to or worse than "SAM" on DMC and DMC-GB. The authors now give a new explanation, attributing it to the "aggressive reset strategy". If the benefits of SAM are potentially overshadowed by resetting on DMC, then the authors should make it clear in the paper (especially abstract) and modify the claims like "Extensive empirical studies on the Atari-100k and DeepMind Control Suite benchmarks demonstrate that this combined usage yields sparse, generalizable features and a dense, plastic policy.".