Learning to Modulate pre-trained Models in RL

Summary

Large-scale pretraining on a diverse dataset followed by finetuning on smaller datasets from downstream tasks has been wildly successful in domains such as computer vision and NLP. The closest analogue to this paradigm in the context of RL is arguably multi-task pretraining followed by finetuning on one or more unseen tasks. This paper investigates the efficacy of finetuning approaches popularized by supervised learning (CV/NLP) on RL problems cast as sequence modeling with Decision Transformers (DT). The authors construct a pretraining dataset that consists of 50 state-based tasks from Meta-World and DMControl (40 and 10 tasks, respectively), and evaluate finetuning methods on held-out tasks from each domain (10 and 6 tasks, respectively) in both single-task finetuning, multi-task finetuning, and continual learning settings. The authors find that their proposed finetuning method, L2M, which combines L2P and LoRA, consistently obtains good performance on unseen tasks after finetuning, while also retaining good performance on the pretraining tasks.

Strengths

The problem is interesting, paper is well written and easy to follow, the method is well motivated, and experiments appear sound. Sufficient discussion of related work. While many existing papers have considered multi-task pretraining and finetuning in RL, I appreciate that the authors take the time to thoroughly investigate trade-offs between different finetuning methods. Further, new finetuning strategies such as LoRA have become very popular in NLP, and this paper confirms that it (along with other modifications necessary to make it work for multi-task DTs as proposed by the authors) can also work well for DTs.

Weaknesses

- **Lack of clarity on experimental setup.** When going through the paper, I found it difficult to fully grasp what the experimental setup looks like and its potential assumptions / pitfalls / failure modes without repeatedly checking the appendix and/or reading between the lines. For example, it is not stated explicitly that finetuning is done strictly on offline replay data which is also collected by single-task SAC agents as in the pretraining dataset. I had to find this information in Appendix D. Likewise, it is not stated explicitly which tasks are included in the pretraining and finetuning datasets, I had to find this information in Appendix A. For the former, it is not really a problem that finetuning requires an offline dataset for the target task, but not making it clear is deceiving. For the latter, the authors do mention that the DMControl tasks include multiple embodiments but do not provide further details in the main paper. I find this problematic since different splits and/or sets of tasks would lead to very different finetuning performances (e.g. task difficulty and degree of overlap). - **Lack of discussion on limitations.** Continuing along the lines of the above, I also find that the paper generally lacks discussion of limitations. Given that the work is very data-driven and domain gaps generally are larger in RL than in NLP, it is important to clearly state assumptions / pitfalls / failure modes related to data collection and experimental setup. Ideally, these limitations (or properties, if you will) would be backed by data that shows, e.g., that finetuning is highly dependent on task similarity. The authors list this as future work but adding such an experiment would make the current submission more complete. For example, the authors could pretrain on Meta-World and finetuning on DMControl, and vice versa, and compare to the performance when including same-domain tasks in the pretraining dataset.

Questions

I would like the authors to address the comments I listed in "weaknesses". Additionally, two clarification questions: - It is stated in Appendix A that the task *reacher-hard* is included in both the pretraining and finetuning datasets. I assume that this is written in error, but would like the authors to please list the correct task splits. - The authors also state in Appendix A that the action spaces considered in DMControl range from 1 (cartpole) to 21 (humanoid) dimensions. However, I do not see any experiments on humanoid in the paper. Can the authors please clarify if they consider humanoid tasks or not?

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

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

3 good

Presentation

3 good

Contribution

3 good

Limitations

I would like to see more discussion on limitations. See my previous comments for constructive feedback.

Summary

This paper studies fine-tuning and continual learning of pre-trained decision transformers in RL. Extensive experiments are conducted to analyze naive fine-tuning, parameter-efficient fine-tuning, and prompt tuning methods on both Meta-world and DMC domains. This paper presents a new method L2M, which combines well-established LoRA and L2P methods and demonstrates the superiority of L2M on Continual-World and DMC benchmarks.

Strengths

1. The proposed method is well-motivated and carefully designed to enable a general agent on multiple domains and tasks. 2. Extensive experiments and ablation study on modulating pre-trained DT 3. Strong performance on the continual learning benchmark

Weaknesses

1. The most significant weakness is the relatively poor presentation of the methodology and experiments, mainly due to the absence of many details. See questions below. 2. Although a unified state space is manually designed for MDDT in this paper, if I understand correctly, it is hard to extend it for new domains with a distinct state space, such as RLBench.

Questions

On the methodology: 1. According to Eq. (2) and (3), it seems that each step of DT separately determines a distinct choice of the modulator. Does this mean that we have different weights in Eq. (1) for each token in the sequence? If so, can it break the advantage of training in parallel for transformers? 2. In Eq. (3), how do we update n(k)? What do we mean by selection count? Do we add n(k) by one, once we encounter a trajectory selecting k? 3. In Line 132, how the learnable keys are updated? I cannot find this additional term in the main text or appendix. 4. In Line 359, the authors state that future work includes combining and sharing modulators across tasks. However, if I understand correctly, we have already shared the modulation pool across tasks in the continual learning setting. On the experiments: 5. In Line 208, the authors claim that they also experimented with PrompDT and VIMA, but I cannot find any experimental details or results. 6. In Section 3.2, where are the details about L2M-oracle? What kind of information (e.g., textual task specification?) is provided to L2M, and how is this information provided in the model? I cannot find implementation details either. 7. In Figure 6, it seems there is a line of straightforward baselines, i.e., separately training a new LoRA or adapter for each new task. Why are these baselines not evaluated? 8. In Figure 7, why do PEFT methods hurt performance on the pretrained tasks? If I understand correctly, we freeze all the parameters of the pre-trained model and only fine-tune the modulators. Minor suggestions on presentation: 9. Since L2M combines L2P and LoRA, the authors should present L2P in the Background section, provide intuitions on how L2P mitigates forgetting, and highlight the difference between L2P and L2M (one uses prompt tuning and one uses LoRA). 10. In Figure 4, the captions of subfigures should be CW10 and DMC6, respectively. Overall, I appreciate the effort made by the authors to conduct extensive experiments and carefully design the method. I recommend the author continuously improve this paper and make it impactful. If the authors solve my questions properly and plan to revise their presentation, I will be happy to increase my rating. ------ Update: The authors have responded with detailed clarification in their rebuttal and most of my concerns are addressed. Thus I increase my rating from 4 to 5.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

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

2 fair

Contribution

4 excellent

Limitations

This work has discussed its limitations, future work, and social impacts. I recommend including weakness 2 mentioned above in the limitation part.

Summary

The authors study the problem of preventing catastrophic forgetting in DT finetuning. The proposed method leveraged a pool of LORA adaptors and only choose the relevant adaptor matrix during finetuning. The author achieve good results on continual world.

Strengths

The applicaiton of lora pools for finetuning DT is novel and the problem of continual learnning in DT is important. The experiment results show that the proposed method works and the released dataset should have a good impact to the community.

Weaknesses

The only weakness I think is the presentation of the method. It seems that the paper is largely inspired by Learing to prompt, and as a result, I believe many of the technical details are not explained in here. E.g., the exact loss function to optimize key. Section I should be polished to include more explanation of the method.

Questions

N/A

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

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

3 good

Presentation

2 fair

Contribution

3 good

Limitations

N/A

Summary

The authors propose an adaptation method for pretrained DT (decision transformer) that combines two finetuning techniques, learning-to-prompt and low rank adaptation (L2P + LoRA), which have been investigated in NLP and computer vision domains. This combined method aims at exploiting the benefits of finetuning and prompt-based learning so that adaptation can be achieved parameter-efficiently and without much catastrophic forgetting.

Strengths

The authors provide the evaluation and comparison on finetuning technique applications to DT-based RL policies, including full finetuning, finetuning with action head, adapters, LoRA, and prompt-tuning, and prefix-tuning that have been well investigated in NLP and computer vision domains. Based on the evaluation, the authors propose an adaptation method combining L2P and LoRA, by which the pretrained DT can be used to solve new tasks.

Weaknesses

The evaluation of different finetuning techniques as survey is meaningful and helpful for readers, but the proposed solution simply combines the two techniques, and little analysis has been conducted on it. There might be some other combination based on L2P, e.g., L2P with IA3, for adaptation. The authors do not explain clearly why multi-domain DT (MDDT) is considered, e.g., MDDT can be useful and effective for learning each domain through shared knowledge and representation. Does the proposed L2P+LoRA get benefited from MDDT? What if a single domain DT was tested? Minor errors: - In line 21, no Section 1 title. - In Figure 4(a), the caption should be CW10. - In Figure 4(b), DMC10 should be DMC6. - In Figure 6(b), Success rate should be Normalized score. - In line 326, a missing citation. - Some citation forms are incorrect.

Questions

Are there any other finetuning techniques that can be combined with L2P, similar to L2P with LoRA? HyperDT [1] handles parameter efficiency, where LoRA part is similar. Could the authors compare the proposed solution with HyperDT? [1] Xu, Mengdi, et al. "Hyper-decision transformer for efficient online policy adaptation." arXiv preprint arXiv:2304.08487 (2023).

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

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

3 good

Presentation

3 good

Contribution

2 fair

Limitations

No specific statement on limitations.

Summary

The paper considers the catastrophic forgetting problem in pre-training and fine-tuning RL setting. The paper proposes Learning-to-Modulate (L2M) to reduce the degradation of pretrained models by modulating the information flow of the frozen pre-trained model via a learnable modulation pool. L2M shows state-of-the-art performance on a continual learning benchmark, while retaining performance on the pre-training tasks.

Strengths

+ The paper is overall well written and easy to follow. + The paper proposes L2M which is a parameter-efficient fine-tuning and prompt-based tuning prompting method. The method is sound and achieves good results.

Weaknesses

- Enhanced baselines: How does L2M measure up against the latest advancements in training methods, such as https://arxiv.org/abs/2211.12740, https://arxiv.org/abs/2301.09816, https://arxiv.org/abs/2211.10869, and https://arxiv.org/abs/2305.16554? These techniques have demonstrated enhanced training of Transformers, resulting in improved generalization and scalability across numerous tasks. Considering the success of these methods in scaling up pretraining, does L2M still hold its relevance? - Absence of weight decay baseline: In the paper, the authors discuss the limitations of commonly used methods that encounter catastrophic forgetting during continual learning. However, it is worth investigating whether the authors explored the effectiveness of weight decay as a preventive measure against overfitting to new tasks. Weight decay, being a straightforward yet potent technique, has proven effective in mitigating overfitting during finetuning processes. Update: The authors have responded with detailed clarification in their rebuttal and most of my concerns are addressed. Thus I increase my rating to 6.

Questions

Is the embedding history state only, did the author consider state-action history? How does the embedding history size impact the results?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, 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

3 good

Presentation

2 fair

Contribution

2 fair

Limitations

Yes, the authors adequately addressed the limitations.

Rebuttal Acknowledged

I would like to thank the authors for their effort during the rebuttal. I appreciate the clarification on weight decay baselines and running extra experiments to ablate history size.

Thank you

Thank you for the clarifications and additional experiments. I believe that this paper will be useful to the NeurIPS community, and incorporating all of the feedback you have received during this rebuttal (from fellow reviewers and I) into a future revision will further strengthen it. I have raised my score accordingly.