TAIL: Task-specific Adapters for Imitation Learning with Large Pretrained Models

Thank you for your detailed and positive review! We are glad you find the paper well-written and the experiments and results comprehensive. We respond to individual concerns and questions below. We also included new experiments in the revised version of the paper to provide an even more comprehensive view of our paper.

Only one environment?

We have added new experiments in Franka-Kitchen, using the same pre-trained model and following the experiment setup presented in the RoboAdapter and R3M paper. The detailed results are shown in Appendix C.7, where TAIL-LoRA performs the best, achieving an 80% success rate while RoboAdapter achieves 65%. The results indicate that our findings and claims from the extensive experiments on the large-scale LIBERO datasets can be generalized well on other small-scale domains, such as Franka-Kitchen.

We would also like to emphasize that the LIBERO dataset is a large-scale, comprehensive benchmark for evaluating continual imitation learning. It comprises over 100 different tasks in various environments, including the Kitchen, Living Room, and Study Room. This dataset aims to evaluate the diverse aspects of continual learning algorithms, including spatial relationship reasoning, object recognition, and understanding of task goals. Therefore, the main results and findings of our work should be generalizable to similar robotic task setups.

Per-task weights for TAIL is unfair comparison?

We clarify that we compare against Frozen Pre-trained Features (FPF) which also learns a new set of per-task weights. TAIL greatly outperforms this baseline in all adaptation environments (Fig 5), indicating that TAIL’s adapter integration is important for forward transfer even against baselines that learn per-task suite new weights.

To further illustrate this, we also add a new experiment comparing against PackNet [1], which prunes weights before re-training the pruned params (as new params) for new task suites. Results demonstrate that TAIL outperforms this baseline by up to 3x in specific task suites. See the updated Fig 5 for further details.

[1]: PackNet: Adding Multiple Tasks to a Single Network by Iterative Pruning. Arun Mallya and Svetlana Lazebnik.

“How do you consider this a ‘limited data’ setup? Can you give context/examples what would not be ‘data limited’ setup?”

Thanks for flagging this. When we talk about "limited data," we mean that only a few demonstrations are available for each task. This is a realistic assumption, considering the expense involved in collecting demonstrations. We have reworded Section 5.2 to make this clearer.

Page 7, last paragraph: Clarify description of what parts were finetuned/update/adapted.

Thank you for the suggestion and we apologize for the confusion. Since we only have pretrained CLIP weights for the vision and text encoders, all other parameters (including the GPT-2 model) must be trained in the pretraining stage. After that, during the adaptation stage for the subsequent task suites, we add adapters for the CLIP vision and text encoders and the temporal GPT2 decoder while keeping their transformer backbone parameters frozen. Since the fusion module and policy head are lightweight, we also tune them fully for each adaptation stage and store their parameters together with corresponding adapters per task suite.

In short, the CLIP spatial and instruction encoders are frozen across all stages. The temporal GPT2 decoder weights are trained in the pretraining stage (Kitchen) but frozen in the adaptation stages. The fusion module and the policy head are fully tuned with adapters per task suite. We have revised the last paragraph of Section 5.2 to address the confusion.

Minor comments: change names, acronyms

Thank you for these suggestions. We will address this in the final revision.

Please let us know if you have any more questions!

Thank you for your insightful comments! We appreciate that you find the paper performs rigorous tests and our TAIL framework could be a valuable tool for resource-constrained settings. We have addressed all your comments below. We hope you consider increasing your score after seeing our detailed responses.

Novelty: largely tests existing PEFT methodologies

Our focus is not on proposing new adapter strategies but rather on proposing a general framework (TAIL) for efficient adaptation in a continual imitation learning setting and testing various adaptation strategies for this setting. No prior work considers this adaptation setting with large pre-trained decision-making models, and prior work like RoboAdapter [1] only considers a single adapter strategy.

Our study contrasts multiple adapter strategies in this setting, including the RoboAdapter strategy, and we find great performance differences among the techniques (Fig 4), with RoboAdapter performing on average the worst among all techniques in LIBERO. Our new experiments demonstrate that RoboAdapter performs worse even in Franka Kitchen, a domain from the paper.

Please see the dedicated “To summarize our contributions…” paragraph in the Response to all reviewers for further details about our contributions.

[1]: Lossless Adaptation of Pretrained Vision Models for Robotic Manipulation (RoboAdapter). Sharma et al, 2023.

Comparison between TAIL and fine-tuning baselines unfair b/c TAIL learns new adapters?

We respectfully disagree that the comparison is unfair. FPF (frozen pre-trained features), like TAIL, also learns dedicated additional parameters per task (in the output head) yet performs much worse. As another comparison, the Experience Replay (ER) method maintains a data buffer for all previous data, which takes much more storage space than TAIL’s adapter weights. Moreover, except FPF, all other baselines (FFT, EWC, ER) are standard techniques for continual learning that tune the full parameters for each new task. In contrast, TAIL only tunes a small portion, yielding better forward transfer performance than them. Hence, TAIL doesn’t have an unfair advantage in our continual learning setting.

To further illustrate this, we have also added a new baseline, PackNet [1], which prunes the network and then re-trains pruned params for the new task. Results demonstrate that TAIL is superior in adaptation performance for almost all task suites, up to a factor of 3x. Please see the updated Fig. 5 for more details.

[1]: PackNet: Adding Multiple Tasks to a Single Network by Iterative Pruning. Arun Mallya and Svetlana Lazebnik.

Can you show LoRA dimensionality ablations?

We agree that this is an important ablation. Thus, we have performed new experiments ablating LoRA ranks from 2 to 768, demonstrating success rates increase to a rank of 32 and then slightly decrease as rank increases to 768. This validates our hypothesis that fine-tuning more parameters on limited data can lead to overfitting. These experiments and discussions have been added to Appendix C.4 and Fig. 9.

Can you combine LoRA and P-tuning for better performance?

Thank you for this excellent suggestion! We’ve now included a new experiment in which we test all combinations of LoRA, bottleneck layers, and p-tuning. The results demonstrate that all combinations that include LoRA perform similarly, while all combinations without LoRA perform far worse, indicating the necessity of LoRA. See Appendix C.6 for more info.

Please let us know if you have any more questions!

Thank you for the comprehensive additional experimental results, which have effectively addressed my previous concerns about the ablations. Hence, I've decided to raise my scores. However, I still believe that the comparison is not entirely fair for a couple of reasons. 1) Even though ER uses all the previous data, the lack of a MoE setting may also lead to worse performance. 2) The different number of trainable parameters remains a concern. Additionally, the ablation studies related to LoRA's rank suggest that increased parameters do not necessarily equate to improved performance.

Thank you for the insightful review. We appreciate that you find the paper easy to follow, intuitive, and that TAIL performs well. We have addressed all your comments below. We hope that you consider increasing your score after seeing our detailed responses and new experiments.

Cross-embodiment tasks?

This is a great suggestion. Following your suggestion, we have added new experiments to use the model pre-trained on LIBERO data to adapt to the Franka kitchen tasks. The detailed results are shown in Appendix C.7, where TAIL-LoRA performs the best, achieving an 80% average success rate while the best baseline achieves 67%. The results indicate that our findings and claims from the extensive experiments on the large-scale LIBERO datasets can be generalized well on other small-scale domains, such as Franka-Kitchen.

“Do we expect the policy that trains purely using the task-specific data to perform better than the models in Figures 5 and 6? How would a random policy perform?”

Thank you for the suggestion regarding the task-specific agent and random agents. First, using a random policy without any fine-tuning always results in a 0 success rate, since our tasks are relatively long and generally require 200-400 steps. However, per request, we added new ablation studies for the task-specific agent in Appendix C.5 of the revised paper.

In summary, if we only train the model with corresponding task data, it does perform better than sequential full fine-tuning in the later stages. However, our TAIL-LoRA still outperforms this task-specific model. In addition, maintaining an entire copy of the model for each adaptation task consumes much more storage space than our approach, and as shown in our results in Table 2 and Appendix C.3, training a large model on a narrow domain with limited data could reduce the model’s plasticity in adapting to new tasks. Therefore, these results further strengthen the advantage of our approach in terms of both resource efficiency and performance.

Tasks explode over time b/c new adapter for each task?

We use a single adapter for each task suite, which is of at least 8 tasks, corresponding to just 5 adapters for the 5 task suites (totaling 40 different tasks) in LIBERO in Fig. 5. We have clarified this in Section 5.2 of the revised paper. While we agree with your suggestion about using adapters based on the similarity of task instructions as it would be an exciting direction, it is beyond the scope of this work and we leave it for future research direction.

Unclear evaluation metric

Thanks for flagging this; there is indeed a typo for the BWT metric calculation. The BWT metric for the k-th task is computed by: $B_k = \frac{1}{k-1} \sum_{i=1}^{k-1} (S_{i} - F_{i})$, where $S_{i}$ values are from the best FWT model at task k, and $F_{i}$ values are from the best FWT model in the corresponding i-th task. We have revised this in Section 5.2.

Validation error doesn’t correspond to success rate?

To clarify, we do not use the validation error to perform model selection. Success rates are simply computed either at the end of adaptation fine-tuning (Table 1) or displayed as a curve over training epochs (Fig 5) (see “training, adaptation, and evaluation” in Section 5.2). As mentioned by your referred paper Appendix G [1], validation loss is in fact a poor measure of policy performance. The purpose of displaying the validation errors is to show that full fine-tuning has a higher risk to overfit to the small data.

[1]: Mandlekar, A., Xu, D., Wong, J., Nasiriany, S., Wang, C., Kulkarni, R., ... & Martín-Martín, R. (2021). What matters in learning from offline human demonstrations for robot manipulation. arXiv preprint arXiv:2108.03298.

“Regarding success rates, what are the baseline performances? That is, what is the success rate of the expert demonstrations?”

The success rate of expert demonstrations is 100%; we have clarified this in Section 5.2.

Some tasks are repeated in evaluation?

Although specific tasks may share similar descriptions, they are not identical due to environment variations, including goals, object arrangements, or object layouts (see Fig 3 for concrete examples). Consequently, these tasks are distinct. We have updated Appendix Section D with this discussion.

What’s a validation scene?

Each task comes with 50 demonstrations with different scene setups. We use 40 for training and 10 for validation. The validation scenes are the scenes from which the validation trajectories were created. We have clarified this in Section 5.2 of the revised paper.

Does ER retrain from scratch or continue fine-tuning?

Experience replay (ER) continues fine-tuning—we do not retrain from the original pre-trained model as it would require storing the entire model per task suite and is extremely storage inefficient for the continual learning setting.

Please let us know if you have any more questions.

Significance: results don’t focus on adaptation plasticity?

We clarify that we do have adaptation plasticity experiments in Table 2 in the main paper and Appendix Table 6. Upon re-visiting previously seen task suites, full fine-tuning suffers large drops in performance despite re-training on the same data (Table 2)

Meanwhile, because TAIL allows us to simply store ~7.8MB of weights per task suite for later reuse, we can reuse the weights when returning to previously seen tasks. As a comparison, storing the entire Kitchen datasets (for the ER method) and the entire model per task suite takes *28GB and 660MB, respectively. Furthermore, adding TAIL w/ LoRA after full fine-tuning is worse than LoRA without full fine-tuning on all tasks (Table 6), indicating that full fine-tuning leads to adaptation plasticity loss.

We have clarified in Section 5.3 and Appendix C of the revised paper that these experiments are studying adaptation plasticity specifically. We hope we have addressed both of your main concerns about originality & significance!

Why do all methods share similar # params? Does this favor LoRA?

We kept the parameter count the same to study how effective adaptation techniques were under the same memory constraint. However, per your request, we have performed a new experiment ablating the bottleneck layer size for the sequential Bottleneck Adapter method, demonstrating that LoRA still outperforms in rank scaling — achieving nearly double the performance at identical rank in some instances. Notably, the sequential Bottleneck Adapters with full rank size, i.e., 768, still under-perform LoRA with rank size 8, showing advantages of the parallel integration style in the continual learning setup. We have added this to Appendix C.4 of the revised paper.

Why study Backwards Transfer (BWT) in TAIL methods? Isn’t it self-evident?

BWT is a relevant and important metric for all non-TAIL methods we compare; it helps analyze catastrophic forgetting among our baselines. In our setting, BWT transfer is poor even among methods made explicitly to resolve this issue (Experience Replay, Elastic Weight Consolidation), indicating the need for TAIL (see Table 1).

Fig 2: colours hard to distinguish

Thank you, we have updated the colors.

Please let us know if you have any more questions!

Thank you for the detailed review. We are happy you found the presentation very good and our motivation excellent. We have addressed all your comments below. We hope you consider increasing your score after seeing our detailed responses and new experiments.

Draw conclusions from more environments, e.g, Franka Kitchen from RoboAdapter?

Thank you for the suggestion; we have now added new experiments in Franka-Kitchen, using the same pre-trained model and following the experiment setup presented in the RoboAdapter and R3M [1] paper. The detailed results are shown in Appendix C.7, where our TAIL-LoRA performs the best, achieving an 80% success rate while RoboAdapter achieves 65%. The results indicate that our findings and claims from the extensive experiments on the large-scale LIBERO datasets can be generalized well on other small-scale domains, such as Franka-Kitchen.

We would also like to emphasize that the LIBERO dataset is a large-scale, comprehensive benchmark for evaluating continual imitation learning. It comprises over 100 different tasks in various environments, including the Kitchen, Living Room, and Study Room. This dataset aims to evaluate the diverse aspects of continual learning algorithms, including spatial relationship reasoning, object recognition, and understanding of task goals. Therefore, the main results and findings of our work should be generalizable to similar robotic task setups.

We hope this addresses your quality & contribution concerns.

[1] Nair, Suraj, et al. "R3m: A universal visual representation for robot manipulation." arXiv preprint arXiv:2203.12601 (2022).

Originality: adapters studied in other works?

Our focus is not on proposing new adapter strategies but on proposing a general framework (TAIL) for efficient adaptation in a continual imitation learning setting and testing various adaptation strategies for this setting. No prior work considers this adaptation setting with large pre-trained decision-making models, and prior work like RoboAdapter [1] only considers a single adapter strategy.

Our study contrasts multiple adapter strategies in this setting, including the RoboAdapter strategy, and we find great performance differences among the techniques (Fig 4), with RoboAdapter performing on average the worst among all techniques in LIBERO. Our new experiments demonstrate that RoboAdapter performs worse even in Franka Kitchen, a domain from the paper.

Please see the dedicated “To summarize our contributions…” paragraph in the Response to all reviewers for further details about our contributions.

[1]: Lossless Adaptation of Pretrained Vision Models for Robotic Manipulation (RoboAdapter). Sharma et al, 2023.

Thank you for your detailed review. We appreciate that you found our paper well-written and the experiments sound! We have addressed all your comments below. We hope you consider increasing your score after seeing our detailed responses and new experiments.

What is the technical contribution?

Please see our answer to this above in the Response to all reviewers that addressed this question: “To summarize our contributions…”

“The proposed TAIL framework, which incorporates lightweight adapters and task-specific heads, shares similarities with existing methods in continual learning and multi-task learning, as demonstrated in the paper from Rebuffi et al. (2017)”

Thank you for this pointer. While both papers have similar motivations, there are major differences. In particular, Rebuffi et al does not (1) study continual imitation learning, (2) investigate decision-making settings, (3) or investigate transformer-based models.

(1) and (2) are important for deploying decision-making agents in realistic settings (e.g., robots) and (3) is important given the impact that transformers have had on large-scale pre-training.

We’ve now discussed Rebuffi et al. in our updated related works section.

Compare TAIL against SOTA in continual imitation learning?

That is a good suggestion. We have now added a comparison against PackNet [1], a SOTA method as shown in the original LIBERO paper, which prunes parameters to be then re-learned for every new task. The detailed results are updated in Fig. 5 of the revised paper. TAIL with LoRA outperforms PackNet in forward imitation in every task suite, for example, 2 times in the Goal task suite and 3 times in the Living Room task suite. Furthermore, PackNet still suffers from severe catastrophic forgetting problems after long continual adaptation stages while TAIL does not.

[1]: PackNet: Adding Multiple Tasks to a Single Network by Iterative Pruning. Arun Mallya and Svetlana Lazebnik.

“In Figure 5, a comparison between TAIL and EWC/ER may not be entirely equitable, considering that these prior methods do not assume known task IDs or retain task-specific model parameters”

We clarify that all methods can access task ID: the language descriptions as shown in Fig 1. Thus, task-specific model parameters are also contained in EWC/ER in the language encoder and the entire base model is fine-tuned in EWC/ER so that they can fairly perform forward transfer. We have clarified this in Section 5.2 of the revised paper.

Try combining LoRA, bottleneck, and p-tuning?

Thank you for this excellent suggestion! We’ve now included new experiments in Appendix C.6  of where we test all combinations of LoRA, bottleneck layers, and p-tuning. All combinations including LoRA perform similarly (75% success rate) and significantly better than other combinations (41-65%), indicating that LoRA is necessary for good adaptation performance.

Separate adapter for each task suite?

Yes, we include a separate adapter for each task suite. We have clarified this in Section 5.2.

Please let us know if you have any more questions.