Paper: The Role of Forgetting in Fine-Tuning Reinforcement Learning Models (ICLR-2024)

Reviewer comment (Writing):
[Minor] the term ‘realistic RL algorithms’ is confusing.

Author response:
We agree that the term is not precise and we get rid of it in the revised version of the paper. Thanks for pointing this out.

Reviewer comment (Theory):
How is non-stationary enironment different from data shifts in stationary environment? Is it not the same underlying data shift problem?

Author response:
On a high level view, the problem which we study can be framed as a ‘data shift’ problem. The difference from the standard online RL setup is the ‘direction of changes’ the training experiences FAR states only after seeing CLOSE ones. In the pre-training/fine-tuning scenario considered in this work, it might be otherwise, and, as we argue (see the General response), this will be more likely as researchers progress with building foundational models in RL. In this work, we highlight that this revered order may cause challenges. We hope this answers your question, please let us know, if you have any suggestions or further questions. This bit is quite challenging to describe in an unambiguous way.

Reviewer comment (Experiments):
What if we pretrain ‘CLOSE’ states first instead? Do we see better forward transfer?

Author response:
We thank the Reviewer for this interesting ablation study. To test it, we run additional experiments in the robotic domain, see Appendix D, paragraph “Other sequences”. In short, we do not observe forgetting and there is no significant gap between vanilla fine-tuning and knowledge retention methods. This is further evidence that the forgetting caused by the state visitation gap is at the heart of the problem of transfer.

Reviewer comment (Experiments):
Why is EWC missing in some of the subsequent experiments?

Author response:
Since our main goal is to show and analyze the forgetting problem in fine-tuning RL models, we decided to use relatively simple methods to avoid multiple overlapping effects. We discovered behavioral cloning to be better and more robust than L2 or EWC, so we decided to use it. For the same reason, we do not check more sophisticated continual learning methods on our benchmark. We believe that incorporating more sophisticated CL methods to the fine-tuning problem is an important next step, which we now mention in the Limitations section.

Reviewer comment (Experiments):
As an experimental paper studying forgetting, it lacks evaluation on many related methods. The paper only evaluates two kinds of methods: parameter regularization and behavior cloning. But there exists many other methods addressing the forgetting issue in the literature of continual RL and finetuning with RL, like using offline RL over previous data [1], adding KL-divergence loss to the pre-trained policy on the online data [2], and a lot of methods in sharing representations and structures [3].

Author response:
We thank the reviewer for this suggestion. Since our main goal is to show and analyze the forgetting problem in fine-tuning RL models, as we discuss in General response, we decided to use relatively simple tools to mitigate it. We discovered behavioral cloning to be a better performing and a more robust approach than L2 or EWC, so we decided to use it. For the same reason, we do not check more sophisticated continual learning methods on our benchmark, even though we believe they might also lead to amazing results. However, we agree with the Reviewer that testing other approaches would be a great idea for further studies.

We thank the reviewer for providing valuable references. Authors of [1] use batch RL in order to maintain performance on the previous tasks. We take advantage of the fact that SAC is an off-policy method that has been a basis for multiple offline RL algorithms [4, 5] and we are currently running experiments where we keep the data from the previous tasks in the buffer. We will update the experiments once they are done. At the same time, we would like to point out that this approach would not be viable in the case of NetHack, where we pre-train the model using behavioral cloning and as such we do not have a critic network to clone.

VPT [2] includes a different type of distillation, but they do not specify what data exactly they are using, e.g. are the trajectories samples from the student or the teacher policy. In order to better understand this problem, we perform further studies of two variations of distillation: (a) where the trajectories are obtained from the student (i.e. known as kickstarting in the literature [6]) (b) and where the targets are obtained from the expert dataset that the pre-trained policy aimed to clone rather than from the pre-trained policy itself. Surprisingly, we find that both of these approaches perform worse than our standard strategy. The full results and analysis are shown in Appendix F, paragraph “Different approaches to distillation”.

Finally, although we believe methods relying on sharing representations and structures [3] could bring substantial improvements to the studied problem, to the best of our knowledge most approaches in this family require clear task separation. In the RL fine-tuning setting it might not be trivial to partition the state space into FAR and CLOSE states, and as such applying these methods would be difficult. However, we believe finding efficient ways to apply such approaches to be an important future research direction. We agree with the reviewer that incorporating a wider variety of CL methods to the fine-tuning problem would be useful, and we add this point to the Limitations section in the revised manuscript.

Following your suggestion, we implemented and evaluated an additional method on Continual World; Episodic Memory (EM), in which we populate SAC’s replay buffer with trajectories from the pre-training tasks. We observe that EM works slightly better than EWC, but worse than behavioral cloning. 

Additionally, we run experiments with EWC on Montezuma’s Revenge. The results show that EWC provides improvements, but is not as good as Behavioral Cloning. Details of the experiments are presented in Figure 5.

We believe that overall new experiments and analyses considerably strengthen our paper. We, thus, gently ask you to consider raising the score or pointing out remaining deficiencies.

Reviewer comment (Evaluation):
The experimental results do not provide different insights into the two problems. All the results demonstrate that Finetuning+BC outperforms other methods and the vanilla Finetuning method suffers from forgetting on the FAR states. But beyond that, the results lack further analysis of the two problems and do not reflect the significance of dividing the forgetting problem into the two types.

Author response:
We use the two problems to highlight natural circumstances in which forgetting of pre-trained capabilities occurs. That is, when the state spaces of the upstream and downstream task overlap but are not identical [7, 8, 9], and when the model is pretrained on offline data and finetuned online [10, 11, 12]. 

There are subtle differences between these two settings. In particular, with the imperfect cloning gap, it is far more difficult to assess in which states we should minimize forgetting. As such, one should be careful about where to apply distillation techniques, as we show through new experiments in Appendix F.

Reviewer comment (Novelty):
The paper has no novel contributions in methods and techniques. It also cannot provide insights in how to better address forgetting in the future work.

Author response:
Based on the experimental results, are there any insights in improving the existing methods or further addressing the forgetting problem?

An important conclusion (see Appendix F) is that it is probably beneficial to use regularization techniques which ‘strength’ depends on the state and the phase of training. Yes, in the future work we provide multiple open research directions.

Reviewer comment (Novelty):
The main takeaway, which is essentially that co-training on the old tasks prevents forgetting when learning a new task, is pretty unsurprising and has been demonstrated before in previous works in continual learning both for the supervised case and the RL case. It's not clear what the contribution of this work adds.

Author response:
We respectfully disagree with the Reviewer regarding the value of the main takeaway. As detailed in the general answer, we consider the conceptualization of 'forgetting of pre-trained capabilities' to be non-obvious and important. We study a new perspective. More precisely, we are not aware of any work comprehensively showing that forgetting is detrimental to transfer in RL fine-tuning scenarios. Please let us know if you know any; we would be happy to include and discuss them.

Reviewer comment (Evaluation):
An obvious downside of co-training on previous tasks is that the memory requirement increases linearly with the number of tasks and the computation increases quadratically - this is not adequately discussed.

Author response:
We thank the Reviewer for raising this important issue. To address this we:
- include additional discussion in the Limitations section,
- provide additional experiments in Appendix D, which indicate that as little as 100 samples is sufficient to observe a significant improvement over fine-tuning. We typically use 10000 transitions sampled out of 5M steps for training, which amounts to roughly 3.5MB of memory. We argue that, compared to the current hardware capabilities, this is negligible.

Reviewer comment (Experiments):
It would have been nice to include result for the L2 and EWC on Montezuma and NetHack.

Author response:
In the robotic experiments, we discovered that behavioral cloning performs better and is more robust than L2 or EWC, and we anticipated the same for the other environments. However, we consider designing and benchmarking other methods to be a very important direction of future work, which is indicated in the Limitations section.  

Following your suggestion, we ran additional experiments with EWC on Montezuma’s Revenge. Our findings indicate that while EWC provides improvements, it still fails short compared to performance achieved by Behavioral Cloning. Detailed results and analyses of these experiments are presented in Figure 5 the revised manuscript.

Reviewer comment (Writing):
In the intro, it would be helpful to give a bit more details on the "knowledge retention techniques" used to mitigate the forgetting problems. Currently, the ready does not have much idea on the methodological aspects going into the paper.

Author response:
We thank the Reviewer for these suggestions. In the revised version of the manuscript we added a short introduction to knowledge retention techniques and we clarified the notation in the MDP examples. Please let us know if you have any further comments.

Reviewer comment (Writing):
Example in 2-state MDP: the notation here is confusing. Both $theta$ and $f_\theta$ are used before being defined. Please add the definitions in the main text.

Author response:
We thank the Reviewer for these suggestions. In the revised version of the manuscript we added a short introduction to knowledge retention techniques and we clarified the notation in the MDP examples. Please let us know if you have any further comments.

Reviewer comment (Writing):
The clarity of certain sections could be improved with more details. For example, for the initial toy example, it would help if the main text explained the motivation behind the MDPs design a little more: why that particular choice of transitions and rewards was made. Also, $f_0$ is not described in the main text and it looks like subfigures b) and c) are interchanged for this example.

Author response:
Thank you for bringing up the issue of clarity of the MDP examples. We have made improvements in the revised version. We also enhanced the clarity of the paper in other sections. All changes are highlighted in blue. Please let us know if you have any more suggestions.

Reviewer comment (Novelty):
The main proposed solution, behaviour cloning from the pretrained policy, seems to be somewhat limited. Behaviour cloning is inherently limited by the quality of the pretrained policy. The experiments show that it's possible to retain the pretrained policy's performance but not exceed it.

Author response:
We agree this is a relevant concern. However, in our experimental results, in many cases fine-tuning with knowledge retention outperforms the level of the pre-trained model. 

For example, on NetHack, the average return of a pre-trained model is 1000, which is outperformed by the BC-regularized fine-tuning agent by a factor of 3.5x, with the same pattern holding for deeper levels (aka FAR states); see Figure 6. On Montezuma’s Revenge, we observe that in terms of completing rooms present in pre-training, PPO-BC manages to outperform the pre-trained model in some cases, see Appendix D, Figure 19. In the robotic manipulation tasks, fine-tuning does not outperform the pre-trained policies since the latter are close to optimal.

We speculate that outperforming the pre-trained policy is due to the fact that regularization is a ‘soft bias’, which can be controlled with hyperparameters (e.g., the KL coefficient), expert data, or the policy network size. As such, this bias can be overwritten in the fine-tuning optimization. Although we found it rather easy to tune, there are also rare negative cases (we detail one of them in Appendix F, paragraph “Different approaches to distillation”). How to efficiently control this 'soft bias' throughout training is an interesting research question, and we included it in the Limitations section.

Reviewer comment (Evaluation):
Are the benefits of pretraining purely from learning a good policy? Are there benefits due to the representations learned in the pretraining phase?

Author response:
We thank the reviewer for these questions – the issue of disentangling the impact of the policy from the impact of the representation is an interesting one. Although there are several studies on the impact of representation learning on transfer in supervised learning [1, 2], the same question in RL remains relatively understudied. We run new experiments to better understand this problem. Please see Appendix D, paragraph “Impact of representation vs policy on transfer”, where we describe experiments with resetting the last layer of the networks before starting the training. As such, the policy at the beginning is random even on the tasks known from pre-training, but has features relevant to solving the downstream tasks. The results suggest that in fine-tuning the benefits come from transferring both policy and representations, and for EWC in particular the representation transfer is more prominent.

Reviewer comment (Experiments):
Have you experimented with the agent only learning a decent, but not great, policy on the FAR tasks? Could we expect to surpass the performance of the pretrained policy? Using behaviour cloning would seem to be limited by the pretrained policy.

Author response:
We agree this is a relevant concern. However, in our experimental results, in many cases fine-tuning with knowledge retention outperforms the level of the pre-trained model. 

For example, on NetHack, the average return of a pre-trained model is 1000, which is outperformed by the BC-regularized fine-tuning agent by a factor of 3.5x, with the same pattern holding for deeper levels (aka FAR states); see Figure 6. On Montezuma’s Revenge, we observe that in terms of completing rooms present in pre-training, PPO-BC manages to outperform the pre-trained model in some cases, see Appendix D, Figure 19. In the robotic manipulation tasks, fine-tuning does not outperform the pre-trained policies since the latter are close to optimal.

We speculate that outperforming the pre-trained policy is due to the fact that regularization is a ‘soft bias’, which can be controlled with hyperparameters (e.g., the KL coefficient), expert data, or the policy network size. As such, this bias can be overwritten in the fine-tuning optimization. Although we found it rather easy to tune, there are also rare negative cases (we detail one of them in Appendix F, paragraph “Different approaches to distillation”). How to efficiently control this 'soft bias' throughout training is an interesting research question, and we included it in the Limitations section.

Reviewer comment (Experiments):
Have you considered off-policy methods? It seems like there may be an advantage since these methods could simply keep around samples (or trajectories) from previous tasks in the replay buffer to learn from without having to necessarily imitate the previous behaviours.

Author response:
We agree with the Reviewer that off-policy methods seem well-suited to deal with the forgetting problem. However, on-policy approaches are often better adjusted to settings that require processing a high number of trajectories to facilitate learning: Montezuma's Revenge and NetHack are examples of such environments. As such, we focus on knowledge retention methods that can be used in a wider variety of problems. At the same time, we agree that this is an interesting research direction and we are currently conducting such experiments on the RoboticSequence benchmark with Soft Actor-Critic. We will update the manuscript once they are ready.

Reviewer comment (Experiments):
In Montezuma's revenge, how far is the agent able to reach without pretraining? Does it get past room 7 consistently? It would be nice to see the overall learning curves of the agent that has been pretrained vs. the one that has not.

Author response:
The agent is able to reach FAR states without pre-training, however only in the later stages of training. We include the visualization of time spent in different rooms across training time in Appendix E. It takes the agent learned from scratch more than 3x longer to start entering further rooms when compared with the fine-tuned one.

Reviewer comment (Experiments):
Nethack levels are generated procedurally. How are the sequence of levels chosen for these experiments? When the agent is reinitialized, is it to a fixed level with the same seed?

Author response:
In terms of training, we run the experiments on the standard NLE environment without any modifications: the levels in each episode are generated randomly. The FAR states in this scenario are not any specific levels but rather related to challenges in the later parts of the game. In this vein, the dungeon level N in Figure 7 denotes these states, which are encountered after passing the first N-1 level. These are stochastic, in practice, we collect a sample of 200 states. We thank the reviewer for this comment and we modify the manuscript to include a further explanation of these issues in Section 4 as well as Appendix B.2.

Reviewer comment (Evaluation):
For the Nethack experiment, Fig.6, how come the learning curve for finetuning only matches that of the original agent after 2 billion steps? It looks like, if training continued further, the pretrained agent would even do worse.

Author response:
Fine-tuning (solid green line) starts with the performance of the pre-trained agent (dashed blue horizontal line). After 2B steps an agent trained from scratch (solid orange line) matches the performance of the fine-tuned agent around the reward value of 2500. Our preliminary results suggest that the performance of these two plateaus at that level. This is consistent with our main takeaway: knowledge transfer from a pre-trained model is greatly hindered due to forgetting of capabilities on the ‘FAR’ states.

Reviewer comment (Evaluation):
The Sokoban results (fig.7) seem to be fairly poor for both agents since they can only fill less than 1.5 pits on average---not close to a solution. Is this to be expected?

Author response:
Yes, this is expected. The original Sokoban (i.e., moving boulders) itself is a complex game, save the fact that the Nethack agent still needs to fight monsters. In Appendix F, we added a more detailed explanation of this matter.

Reviewer comment (Presentation):
In Fig. 4, I would consider changing the text "pre-trained optimal policy" to "pre-trained expert policy" since we don't necessarily have the optimal policy in those environments.

Author response:
Thank you for the suggestion, we corrected it in the revised version of the manuscript.