Paper: TAIL: Task-specific Adapters for Imitation Learning with Large Pretrained Models (ICLR-2024)

Reviewer comment (Experiments):
Using a separate set of weights ("adapters") per task feels bit of cheating when comparing the model against the baselines. This alone ensures that you keep high performance in the previous tasks. While the argument for this is valid (with TAIL, you only need a handful of parameters vs. full network), I'd prefer a more apples-to-apples comparison.

Author response:
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.

Reviewer comment (Experiments):
Limited to only one environment. While I understand data for right setup is hard to come by, it is hard to tell if the results generalize outside this tested environment and data. Showing results in an another type of environment would solidify these insights.

Author response:
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.

Reviewer comment (Experiments):
Page 7, "Training, Adaptation and Evaluation" paragraph: Manuscript says "This limited data setup...". How do you consider this a "limited data" setup? Can you give context/examples what would not be "data limited" setup?

Author response:
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.

Reviewer comment (Experiments):
Page 7, last paragraph. The description of what parts were finetuned/update/adapted and when is somewhat unclear and involved. How did you come up with this setup? Could you clarify this paragraph or provide a figure to help with understand what was updated in which stage?

Author response:
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.

Reviewer comment (Writing):
Define acornyms only once and consistently use the acornym ever since (e.g., TAIL is defined multiple times).

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

Reviewer comment (Experiments):
Drawing a comparison between TAIL and the fine-tuning baseline methods cited in the paper is unfair. This is because TAIL is designed to learn dedicated adapters for each distinct task. Thus, the claims regarding its efficacy become less compelling.

Author response:
> **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.

Reviewer comment (Experiments):
Could you show the results when setting the rank for LORA to 32, 256, 512, and 768 (the dimensionality of the embeddings and hidden states of GPT2)? The ablation studies regarding the number of parameters should be added.

Author response:
> **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.

Reviewer comment (Experiments):
Could you provide some insights on whether combining the Lora and P-tuning can improve the performance?

Author response:
> **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.

Reviewer comment (Experiments):
One of the main weaknesses of this paper is the lack of comprehensive empirical analysis over multiple different benchmarks, especially those used by previous baseline methods. Particularly, the conclusions and empirical insights were drawn only from a single dataset, with no results from any other public datasets. As the main contribution, this paper shows that the LoRA-type integration of the adapter in continual imitation learning performs the best. The paper thus concludes that they “are contrary to many results from the vision and language model literature which show that full fine-tuning works better”. However, these results are drawn based on LIBERO robotic manipulation continual learning benchmark ONLY. They have not been verified by any other datasets. It is thus unclear if and how the results and conclusions can be generalized to other robot control tasks, e.g., Franka-Kitchen.

Author response:
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).

Reviewer comment (Experiments):
It is unclear if the baseline methods have been well tuned for the LIBERO benchmark. For example, the RoboAdapter method may need to tune the adapter locations and different pre-trained representations for its optimal performance. The paper, however, seems to take the default configs, even though those configs were originally proposed for other robot tasks. To have a fair comparison with RoboAdapter, the paper should consider some of datasets used by RoboAdapter paper, i.e., Metaworld, Franka-Kitchen, and RGB-Stacking task suites, replicate the reported performance of RoboAdapter and then demonstrate on at least one of those datasets that the TAIL performs better than RoboAdapter.

Author response:
Regarding the baseline setups, we have a detailed description of each baseline approach in Appendix B.1, including their method implementations, corresponding analysis, and insights.  We would like to remind the reviewer to check the updated Appendix C (highlighted in red) for the additional experiments.

Reviewer comment (Novelty):
The originality of this paper can be limited and incremental. In terms of technique, the adapter idea has been well studied by the RoboAdapter paper, which introduces the adapter layer (Sequential Integration in this paper) to imitation learning and considers its application in robot controls. Though the paper proposes two other adapter mechanisms, LoRA & Prefix prompt tuning, both of them are taken directly from other papers, without any substantial modifications or creative combinations to the continual imitation learning or robot control tasks.

Author response:
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.

Reviewer comment (Evaluation):
In terms of the empirical results, they may not offer any in-depth understanding of using adapters in continual imitation learning and the significance can be limited. The paper mentions that TAIL with LoRA “avoiding catastrophic forgetting and preserving adaptation plasticity”. Avoiding catastrophic forgetting can be evident since the adapter weights are task-specific and not shared across different tasks. But the empirical results shed little light on the adaptation plasticity. It is unclear if and how the adaptation plasticity is preserved given that the empirical results focus only on the success rate, as in Figure 4.

Author response:
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!

Reviewer comment (Experiments):
In appendix, “all methods share similar amount of parameters”. I’m not quite sure how to interpret this. Why all methods need to have a similar number of parameters? Even though the adapter integration mechanisms are different? Could this constraint on the number of parameters be biased towards to LoRA adapter? (since “filtering often requires a larger bottleneck size compared to that of LoRA, leading to more parameters”)

Author response:
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.

Reviewer comment (Evaluation):
In the experiment Table 1, “The BWT for TAIL methods are all 0 (no catastrophic forgetting)”. Is it self-evident? Since the TAIL methods produce an ensemble of models (fixed pretrained weights + task-specific adapter weights), each of which aims to solve a dedicated task and is updated independently to each other. Then why do we need to use BWT as a performance indicator.

Author response:
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).

Reviewer comment (Presentation):
The colours for Prefix Token & Frozen cannot be easily distinguished in Figure 2.

Author response:
Thank you, we have updated the colors.

Reviewer comment (Experiments):
It appears that the number of tasks can easily explode over time since each task instruction (and initial-state distribution) corresponds to a new task. A natural extension will be to mitigate the amount of adapters based on the similarity of task instructions, as hinted in appendix.

Author response:
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.

Reviewer comment (Experiments):
Since the problem setting indicates that a task definition is not w.r.t. the state-action space, it will be more convincing if there are experiments conducted on cross-embodiment (e.g. different arms or same arm with different inertial properties.)

Author response:
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.

Reviewer comment (Evaluation):
The evaluation metric is unclear — in particular under BWT, the equation seems to be different from what is described — should the *best FWT model* be the best \( F_i \) for task \( i \) after seeing \( k \) tasks?

Author response:
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.

Reviewer comment (Evaluation):
For the 10 validation episodes, I believe it has been shown multiple times that validation error does not necessarily correspond to success rates [1, 2, 3] — a model may achieve high validation error while still achieving high success rate.

Author response:
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.

Reviewer comment (Writing):
Under subsection **Training, Adaptation, and Evaluation**, what exactly is a validation scene?

Author response:
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.

Reviewer comment (Reproducibility):
Does ER retrain parameters from scratch, or does it continue training from the current set of parameters? Clearly if it is the latter the model will experience plasticity loss similar to the reinforcement learning setting [1].

Author response:
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.

Reviewer comment (Evaluation):
It appears that some tasks are repeated based on appendix. I am wondering if there is any result regarding the per-task success rate. Do we understand whether the parameters are well-adapted to specific tasks, or are they similar in performance?

Author response:
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.

Reviewer comment (Experiments):
In Figure 4, the authors primarily compare various design choices within the TAIL framework when presenting forward adaptation results. It would be valuable if the authors could extend this analysis to include a comparison of TAIL's best performance with the state of the art in lifelong imitation learning. Furthermore, in Table 1, the comparison of TAIL's forward transfer performance with previous continual learning methods may not be fair enough. This is because EWC and ER primarily address catastrophic forgetting and may not contain task-specific model parameters.

Author response:
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.

Reviewer comment (Experiments):
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. I would like to suggest the authors include a comparison with multi-task learning methods for a more comprehensive evaluation.

Author response:
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.

Reviewer comment (Experiments):
Is a separate adapter $\omega_k$ trained for each task suite? My understanding is that when re-evaluating previous task suites $j$, the previously trained adapter $\omega_j$ is reloaded instead of testing the previous task on the adapter $\omega_k$ trained for the current task $k$.

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

Reviewer comment (Experiments):
In Figure 4, it would be insightful to determine whether combining all integration techniques, including LoRA, Bottleneck Adapter, and Prefix Token, would yield further performance improvements.

Author response:
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.