Mastering Robot Manipulation with Multimodal Prompts through Pretraining and Multi-task Fine-tuning

We are grateful for the positive and detailed feedback from the reviewer. Please find our responses below. We hope that the following response will further clarify any concerns and assist in advocating for the acceptance of our paper.

It would be more solid if authors cloud show similar improvement on other robot learning benchmarks such as RLBench (James et al., 2020).

As mentioned by the reviewer, RLBench does not support multimodal prompts. Thus, we choose to extend the existing VIMA-BENCH by adding more tasks to showcase the in-context learning ability of our models.

I would encourage authors to extend existing VIMA-BENCH by adding more representative tasks to show the in-context learning ability of models trained with the proposed method.

We thank the reviewers for the constructive comments! We designed 4 new tasks with in-context learning examples in the task prompt. We evaluate our policy that is trained on the full VIMA-BENCH data on these 4 tasks and compare it with the VIMA policy. The table below shows that our method demonstrates superior performance over all baseline methods. On the other hand, the VIMA policy struggles with these tasks, revealing its inability to learn from the in-context demonstration. More experiment details, including task definitions, can be found in Appendix D.

Missing citations. Authors are encouraged to discuss the following recent related work

Thank you for bringing up these two works! We have updated our manuscript to cite and discuss these works in the related work section.

Although this paper is not designed to address real-robot manipulation, showing proof-of-concept demos would justify the feasibility of applying this method on real hardware.

We appreciate the reviewer's suggestion. We currently do not have access to a real robot and thus cannot conduct a demonstration. Nevertheless, we assume the same action space as in Transporter [1] and CLIPort [2] and share similar settings except for tackling multimodal prompts. These alignments suggest strong potential for real-world deployment of our method. We believe our contributions should still be valid in the absence of a real-robot demonstration.

With this shortcut, is there any performance difference between LMs with varying depth?

We evaluate our methods by replacing the T5-base with T5-small. The experiment result is shown below. We can see that the performance difference is minimal.

[1] Zeng et al., Transporter Networks: Rearranging the Visual World for Robotic Manipulation, CoRL 2020.

[2] Shridhar et al., CLIPORT: What and Where Pathways for Robotic Manipulation, CoRL 2021.

Dear Reviewer yFyT,

Thank you again for your time and effort. Your feedback has been valuable in helping us clarify, improve, and refine our work. We have carefully addressed your comments in our authors' responses to improve the quality of our paper. We thus kindly request that you take a moment to revisit our paper and consider the changes we have made. We hope our clarifications can help you better advocate for our paper's acceptance.

Best regards,

The authors

We appreciate the detailed feedback from the reviewer. Please find our responses below.

The pretraining method is not general enough: it only concern about instruction of "follow motion for ..." for a particular motion trajectory, and therefore it mainly tackles the tasks with prompts given a certain motion of a certain trajectory.

We would like to clarify an important misunderstanding: The goal of our pretraining is not to achieve zero-shot generalization for any unseen task. Instead, we aim to endow robots with an intuitive, human-like ability to learn and adapt to new tasks through in-context demonstrations. This focus is the cornerstone of our inverse dynamic pretraining design, facilitating the robot's understanding of the underlying transition dynamics suggested by the multimodal prompts.

Nevertheless, our pretraining can still improve our model's performance on tasks without trajectory instructions in the prompt. Evidence of this can be found in Table 1, which shows that our pretraining can improve the performance of Task 5 (Rearrange then restore) consistently on L1, L2, and L3 levels, despite Task 5 only containing subgoal images in its prompt, demonstrating the versatility of our pretraining approach.

This means it assumes the task at hand is always similar to follow motion.... it can do well for task T10, but for task T13... it cannot generalize.

We contend that the challenges in solving Task 13 (Sweep without Touching) are not indicative of shortcomings in our algorithm design. And thus, tackling Task 13 is outside the scope of our paper.

Task 13 is an L4 task in the VIMA-BENCH, which is out of the training distribution. Our analysis of the VIMA-BENCH training tasks and trajectories reveals two primary factors hindering generalization to Task 13:

1. End Effector Utilization: T13 employs a "spatula" end effector, a tool only used in Task T12 (Sweep without Exceeding) among all training tasks. 1. The majority (12 out of 13) of training tasks utilize a "suction cup" end effector. This distinction is critical as the suction cup is associated with "pick and place" motor skills, while the spatula is linked to "wipe" actions, as detailed in Section 2.
2. Layout and Prompt Specificity: The workspace layouts of Task 12 and Task 13 are nearly identical, yet they diverge significantly from the other training tasks. Task 13's prompt only differs from T12's by the term touching. However, the training data lacks the necessary signal to effectively conceptualize and respond to touching, impacting the model's performance on Task 13.

Consequently, policies trained via our methods and VIMA tend to replicate Task 12's action sequence in Task 13, leading to failure.

However, we kindly ask the reviewer to pay attention to our method's success in significantly improving performance on Task 5, 9, and 17. This improvement is not limited to motion-following tasks. Furthermore, while Task 10 (Follow Motion) shares similarities with our pretraining prompts, the demonstration image sequence in Task 10 includes distractors absent from the robot's workspace (see Figure 2). Conversely, our pretraining tasks incorporate the robot's observational sequence directly into the task prompt, as illustrated in Figure 4.

For the pretraining method to work, this method also assumes that the prompts contains the motion trajectory keypoint, which is a very narrow assumption and might not always hold.

First, We would like to clarify the misunderstanding that our pretraining method does not depend on the motion trajectory key point. Our method only requires a sequence of state-action pairs of the robot trajectory, which is always satisfied in the multi-task imitation settings. Second, the key points are a natural consequence of our action space $\mathcal{A} = (\mathcal{T}\_{intial}, \mathcal{T}\_{target})$, as each observation is gathered at the start/end of each action. Moreover, we emphasize that this versatile action space has been widely adopted in various robotics system, including VIMA, Transporter Network[1], and CLIPort [2].

The end users would not be expected to provide the entire trajectories all the time. Therefore the pretraining on motion following is a bit overfitting to the tasks that VIMA designed.

We are considering multi-task imitation learning where demonstration trajectories for each training task are provided. We argue that multi-task imitation learning is a general setting adopted in various works [1, 2]. We can always conduct our pretraining for any robot trajectory with a sequence state-action of pair [3]. Therefore, our pretraining strategies are NOT overfitting to the VIMA designed.

Dear Reviewer Y3nY,

Thank you again for your time and effort. Your feedback has been valuable in helping us clarify, improve, and refine our work. We have carefully addressed your comments in our authors' responses to improve the quality of our paper. We thus kindly request that you take a moment to revisit our paper and consider the changes we have made. We hope our clarifications warrant a more positive evaluation of our work.

Best regards,

The authors

We would like to clarify several major misunderstandings on the contributions of our work.

using inverse dynamic modeling loss and multi-task supervision loss are all intuitive and an easy follow-up step after the VIMA paper

We would like to emphasize that our core novelty does not lie in the specific training losses employed. It lies in equipping a multi-task policy with the capacity for in-context learning, which is achieved by applying a two-stage training pipeline. Our goal is to endow a multi-task robot with an intuitive, human-like ability to learn and adapt to new tasks through in-context demonstrations rather than striving for zero-shot generalization to entirely unseen tasks. Table 1, 2, and Appendix D shows that our method establishes SOTA performances across all four evaluation protocols while exhibiting a superior in-context learning ability. To the best of our knowledge, simultaneously equipping a robot with multi-task and in-context learning abilities has not been extensively explored in robotics, making our contribution a significant step forward in this domain.

(2) residual connections in the visual layers

We would like to clarify that our contribution is an effective multimodal prompt encoder that can capture visual and textual details. Specifically, it is constructed by adding a residual connection (RC) from the input visual tokens to the encoded embeddings, assisting the encoding process to retain more detailed visual information. Notably, this simple yet effective design is not limited to the task of robotics control. It can be incorporated into any multimodal LLM (MLLM), e.g., LLaVA [1] and AnyMAL [2], to facilitate general multimodal understanding.

(3) a small performance gain (10%) compared to the VIMA paper

Our method achieves average success rates of 97.8% (+10.6% over VIMA) on L1, 97.9% (+10.9% over VIMA) on L2, 93.4% (+9.4% over VIMA) on L3, and 59.1% (+9.5% over VIMA) on L4 on the standard VIMA-BENCH. Notably, we highlight our exceptional gains on Task 5 (Rearrange the restore, +31.8%), Task 9 (Twist, +86.3%), Task 10 (Follow Motion, + 41.0%) and Task17 (Pick then restore, +19.3%). These results clearly indicate a major advancement over the VIMA model, not just a small performance gain.

The reviewer found the contributions are not sufficient to be published as a long paper in ICLR.

We kindly remind that the reviewer has overlooked our policy's superior in-context learning ability, as empirically supported in Section 4.3 and Appendix D. Our model can learn from in-context examples for novel tasks, a feature where the baseline VIMA model falls short. We highlight the in-context learning ability is crucial for a generalist robot.

We assert the novelty of our research. Prior works have primarily focused on developing either a multi-task robot or one capable of learning from demonstrations or in-context examples. However, integrating in-context learning into a multi-task policy remains under-explored in robotics.

Moreover, our work sheds light on the emergence of the in-context learning ability in our model, a topic of significant interest in NLP. While one might initially credit inverse dynamic pretraining for this ability, Table 2 reveals that policies developed through our two-stage pipeline (both Our Method and Our Method w/o Modified FT) outperform those from inverse dynamic pretraining alone (Our Method w/ Pretrain Only). This phenomenon echoes the findings in FLAN [3], which indicate that a finetuned language model excels in in-context learning.

We thank the reviewer's invaluable feedback, which helped us clarify our contributions. We kindly ask the reviewer to revisit our paper in light of our response and consider whether the clarifications we have provided might warrant a reconsideration of the rating.

The authors add the appendix pages in the main paper, which exceeds the page limits. Please remove the appendix in revision.

We follow the official instructions of ICLR https://iclr.cc/Conferences/2024/CallForPapers. We make sure our main text is within the 9 pages limit. Paper length

There will be a strict upper limit of 9 pages for the main text of the submission, with unlimited additional pages for citations. This page limit applies to both the initial and final camera ready version.

• Authors may use as many pages of appendices (after the bibliography) as they wish, but reviewers are not required to read the appendix.

[1] Liu et al., Visual Instruction Tuning. NeurIPS 2023

[2] Moon et al., AnyMAL: An Efficient and Scalable Any-Modality Augmented Language Model. arXiv:2309.16058

[3] Wei et al. "Finetuned language models are zero-shot learners." ICLR 2022

Dear Reviewer YMJx,

Thank you again for your time and effort. Your feedback has been valuable in helping us clarify, improve, and refine our work. We have carefully addressed your comments in our authors' responses to improve the quality of our paper. We thus kindly request that you take a moment to revisit our paper and consider the changes we have made. We hope our clarifications warrant a more positive evaluation of our work.

Best regards,

The authors

We are grateful for the reviewer's positive feedback on our work. We hope the following response will further clarify any concerns and assist in advocating for the acceptance of our paper.

The prompts are quite controlled during pretraining versus the more complex prompts at test time. It is unclear if the pretraining fully transfers to the downstream tasks.

Our pretraining mainly focuses on enabling the agent to reason the inverse dynamics from the task prompt, and thus does not incorporate a wide range of language. Our approach can still benefit the downstream tasks without demonstration trajectory in the prompt. For example, it can improve the performance of Task 5 (Rearrange then restore) consistently across L1, L2, and L3, despite Task 5 only containing subgoal images in the prompt.

We acknowledge the potential benefits of a more varied language in our pretraining prompts and consider its inclusion as a valuable future extension of our research.

For the inverse dynamics pretraining, were other self-supervised objectives explored besides simply reconstructing the actions?

In this project, we only tried the inverse dynamic prediction as the pretraining tasks given limited resources. However, it might also be possible to reconstruct the observation images during pretraining, as shown by Radosavovic et al. [1]

What stopped the baseline VIMA model from reaching the same performance with just more compute/data?

1. The VIMA models each action dimension independently, which can lead to task failure when tackling tasks requiring coordination between the initial and target pose of the action.
2. The training data provided in VIMA-BENCH is imbalanced. We reveal this problem by finetuning a VIMA policy trained on the full multi-task data on only the trajectories of Task 5 (Twist). Its success rate increased from ~13% to ~92% on L1, L2, and L3 evaluations.
3. Performing multi-task imitation alone is insufficient to enable the VIMA policy to reason inverse dynamics given novel tasks with in-context examples in the prompt, which leads to the VIMA policy's inability to perform, e.g., Task 9 (Twist) and Task 10 (Follow Motion).

Is there other complementary information like force sensors that could augment the visual observations?

The current VIMA-BENCH only provides visual observations. However, in general, we can always include complementary information to augment the observation space, and this state-based information, like force sensors, can often improve the sample efficiency during training. For example, Radosavovic et al. incorporate proprioceptive robot states to assist in action prediction.

However, in-context examples/demonstrations are more accessible to gather as a sequence of images/video when tackling novel unseen tasks. Therefore, we consider pure image observations in this paper.

[1] Radosavovic et al., Robot Learning with Sensorimotor Pre-training, arXiv 2023.

Dear Reviewer t7wu,

Thank you again for your time and effort. Your feedback has been valuable in helping us clarify, improve, and refine our work. We have carefully addressed your comments in our authors' responses to improve the quality of our paper. We thus kindly request that you take a moment to revisit our paper and consider the changes we have made. We hope our clarifications warrant a more positive evaluation of our work.

Best regards,

The authors