paper7

VIMA: General Robot Manipulation with Multimodal Prompts

Yunfan Jiang, Agrim Gupta, Zichen Zhang, Guanzhi Wang, Yongqiang Dou, Yanjun Chen, Li Fei-Fei, Anima Anandkumar, Yuke Zhu, Linxi Fan

cs.ROcs.AIcs.LG

Introduction

Transformer models (Vaswani et al., 2017) have given rise to remarkable multi-task consolidation across many AI domains. For example, users can describe a task using natural language prompt to GPT-3 (Brown et al., 2020), allowing the same model to perform question answering, machine translation, text summarization, etc. Prompt-based learning provides an accessible and flexible interface to communicate a natural language understanding task to a general-purpose model.

We envision that a generalist robot should have a similarly intuitive and expressive interface for task specification. What does such an interface for robot learning look like? As a motivating example, consider a personal robot tasked with household activities. We can ask the robot to bring us a cup of water by a simple natural language instruction. If we require more specificity, we can instead instruct the robot to “bring me ”. For tasks requiring new skills, the robot should be able to adapt, preferably from a few video demonstrations (Duan et al., 2017). Tasks that need interaction with unfamiliar objects can be easily explained via a few image examples for novel concept grounding (Hermann et al., 2017). Finally, to ensure safe deployment, we can further specify visual constraints like “do not enter room”.

To enable a single agent with all these capabilities, we make three key contributions in this work: 1) a novel multimodal prompting formulation that converts a wide spectrum of robot manipulation tasks into one sequence modeling problem; 2) a large-scale benchmark with diverse tasks to systematically evaluate an agent’s scalability and generalization; and 3) a multimodal-prompted robot agent capable of multi-task and zero-shot generalization.

We start with the observation that many robot manipulation tasks can be formulated by multimodal prompts that interleave language and images or video frames (Fig. 1). For example, Rearrangement (Batra et al., 2020), a type of Visual Goal, can be formulated as “Please rearrange objects to match this {scene_image}”; Few-shot Imitation can embed video snippet in the prompt “Follow this motion trajectory for the wooden cube: {frame1},{frame2},{frame3},{frame4}\{\texttt{frame}_{1}\},\{\texttt{frame}_{2}\},\{\texttt{frame}_{3}\},\{\texttt{frame}_{4}\}”. Multimodal prompts not only have more expressive power than individual modalities but also enable a uniform sequence IO interface for training generalist robots. Previously, different robot manipulation tasks required distinct policy architectures, objective functions, data pipelines, and training procedures (Aceituno et al., 2021; Stengel-Eskin et al., 2022; Lynch & Sermanet, 2021), leading to siloed robot systems that cannot be easily combined for a rich set of use cases. Instead, our multimodal prompt interface allows us to harness the latest advances in large transformer models (Lin et al., 2021; Tay et al., 2020; Khan et al., 2021) for developing scalable multi-task robot learners.

To systematically evaluate agents with multimodal prompts, we develop a new benchmark, named VIMA-Bench, built on the Ravens simulator (Zeng et al., 2020; Shridhar et al., 2021). We provide 17 representative tasks with multimodal prompt templates. Each task can be procedurally instantiated into thousands of instances by various combinations of textures and tabletop objects. VIMA-Bench establishes a four-level protocol to evaluate progressively stronger generalization capabilities, from randomized object placement to novel tasks (Fig. 2).

To this end, we introduce the VisuoMotor Attention agent (VIMA) to learn robot manipulation from multimodal prompts. The model architecture follows the encoder-decoder transformer design proven to be effective and scalable in NLP (Raffel et al., 2020). VIMA encodes an input sequence of interleaving textual and visual prompt tokens with a pre-trained language model (Tsimpoukelli et al., 2021) and decodes robot control actions autoregressively for each environment interaction step. The transformer decoder is conditioned on the prompt via cross-attention layers that alternate with the usual causal self-attention. Instead of operating on raw images, VIMA adopts an object-centric approach. We parse all images in the prompt or observation into objects by off-the-shelf then domain fine-tuned detectors (He et al., 2017) and flatten them into sequences of object tokens. To demonstrate the scalability of VIMA, we train a spectrum of 7 models ranging from 2M to 200M parameters. Our approach outperforms other design alternatives, such as image patch tokens (Reed et al., 2022), image Perceiver (Jaegle et al., 2021b; Alayrac et al., 2022), and decoder-only conditioning (Radford et al., 2018). VIMA obtains consistent performance gains across all four levels of zero-shot generalization and all model capacities, in some cases by a large margin (up to 2.9×2.9\times task success rate given the same amount of training data, and 2.7×2.7\times better even with 10×10\times less data). We open-source the simulation environment, training dataset, algorithm code, and pre-trained model checkpoints to ensure reproducibility and facilitate future work from the community. These materials along with video demos are available at vimalabs.github.io.

Multimodal Prompts for Task Specification

A central and open problem in robot learning is task specification (Agrawal, 2022). In prior literature (Stepputtis et al., 2020; Dasari & Gupta, 2020; Brunke et al., 2021b), different tasks often require diverse and incompatible interfaces, resulting in siloed robot systems that do not generalize well across tasks. Our key insight is that various task specification paradigms (such as goal conditioning, video demonstration, natural language instruction) can all be instantiated as multimodal prompts (Fig. 1). Concretely, a multimodal prompt P\mathcal{P} of length ll is defined as an ordered sequence of arbitrarily interleaved texts and images P:=[x1,x2,…,xl]\mathcal{P}\vcentcolon=\begin{bmatrix}x_{1},x_{2},\ldots,x_{l}\end{bmatrix}, where each element xi∈{text,image}x_{i}\in\{\text{text},\text{image}\}.

The flexibility afforded by multimodal prompts allows us to specify and build models for a variety of task specification formats. Here we consider the following six categories.

Simple object manipulation. Simple tasks like “put into ”, where each image in the prompt corresponds to a single object;

Visual goal reaching. Manipulating objects to reach a goal configuration, e.g., Rearrangement (Batra et al., 2020);

Novel concept grounding. The prompt contains unfamiliar words like “dax” and “blicket”, which are explained by in-prompt images and then immediately used in an instruction. This tests the agent’s ability to rapidly internalize new concepts;

One-shot video imitation. Watching a video demonstration and learning to reproduce the same motion trajectory for a particular object;

Visual constraint satisfaction. The robot must manipulate the objects carefully and avoid violating the (safety) constraints;

Visual reasoning. Tasks that require reasoning skills, such as appearance matching “move all objects with same textures as into ”, and visual memory “put in and then restore to their original position”.

Note that these six categories are not mutually exclusive. For example, a task may introduce a previously unseen verb (Novel Concept) by showing a video demonstration, or combine goal reaching with visual reasoning. More details about the task suite are discussed in Appendix, Sec. B.

VIMA-Bench: Benchmark for Multimodal Robot Learning

Existing benchmarks are generally geared towards a particular task specification. To our knowledge, there is no benchmark that provides a rich suite of multimodal tasks and a comprehensive testbed for targeted probing of agent capabilities. To this end, we introduce a new benchmark suite for multimodal robot learning called VIMA-Bench. We build our benchmark by extending the Ravens robot simulator (Zeng et al., 2020). VIMA-Bench supports extensible collections of objects and textures to compose multimodal prompts and to procedurally generate a large number of tasks. Specifically, we provide 17 tasks with multimodal prompt templates, which can be instantiated into thousands of task instances. Each task belongs to one or more of the 6 task categories mentioned above. VIMA-Bench can generate large quantities of imitation learning data via scripted oracle agents. More details are elaborated in Appendix, Sec. A.

The observation space of our simulator includes RGB images rendered from both frontal view and top-down view. Ground-truth object segmentation and bounding boxes are also provided for training object-centric models (Sec. 4). We inherit the high-level action space from Zeng et al. (2020), which consists of primitive motor skills like “pick and place” and “wipe”. These are parameterized by poses of the end effector. Our simulator also features scripted oracle programs that can generate expert demonstrations by using privileged simulator state information, such as the precise location of all objects, and the ground-truth interpretation of the multimodal instruction.

We leverage oracles to generate a large offline dataset of expert trajectories for imitation learning. Our dataset includes 50K trajectories per task, and 650K successful trajectories in total. We hold out a subset of objects and textures for evaluation and designate 4 out of 17 tasks as a testbed for zero-shot generalization.

Each task in VIMA-Bench has a binary success criterion and does not provide partial reward. During test time, we execute agent policies in the simulator for multiple episodes to compute a percentage success rate. The average success rate over all evaluated tasks will be the final reported metric.

We design a four-level evaluation protocol (Fig. 2) to systematically probe the generalization capabilities of learned agents. Each level deviates more from the training distribution, and is thus strictly harder than the previous one.

Placement generalization. All prompts are seen verbatim during training, but only the placement of objects on the tabletop is randomized at testing;

Combinatorial generalization. All textures and objects are seen during training, but new combinations of them appear in testing;

Novel object generalization. Test prompts and the simulated workspace include novel textures and objects;

Novel task generalization. New tasks with novel prompt templates at test time.

VIMA: Visuomotor Attention Agent

Our goal is to build a robot agent capable of performing any task specified by multimodal prompts. There is no prior method that works out of the box with multimodal prompts. To learn an effective multi-task robot policy, we propose VIMA, a robot agent with a multi-task encoder-decoder architecture and object-centric design (Fig. 3). Concretely, we learn a robot policy π(at∣P,H)\pi(a_{t}|\mathcal{P},\mathcal{H}), where H:=[o1,a1,o2,a2,…,ot]\mathcal{H}\vcentcolon=\begin{bmatrix}o_{1},a_{1},o_{2},a_{2},\ldots,o_{t}\end{bmatrix} denotes the past interaction history, and ot∈O,at∈Ao_{t}\in\mathcal{O},a_{t}\in\mathcal{A} are observations and actions at each interaction steps. We encode multimodal prompts via a frozen pre-trained language model and decode robot waypoint commands conditioned on the encoded prompts via cross-attention layers. Unlike prior work (Florence et al., 2019; Sieb et al., 2019; Zhu et al., 2022), VIMA adopts an object-centric representation that computes tokens from bounding box coordinates and cropped RGB patches.

There are 33 formats of raw input in the prompt — text, image of a single object, and image of a full tabletop scene (e.g., for Rearrangement or imitation from video frames). For text inputs, we use pre-trained T5 tokenizer and word embedding to obtain word tokens. For images of full scenes, we first extract individual objects using domain fine-tuned Mask R-CNN (He et al., 2017) (Appendix, Sec. C.4). Each object is represented as a bounding box and a cropped image. We then compute object tokens by encoding them with a bounding box encoder and a ViT (Dosovitskiy et al., 2020), respectively. Since Mask R-CNN is imperfect, the bounding boxes can be noisy and the cropped images may have irrelevant pixels. For images of single objects, we obtain tokens in the same way except with a dummy bounding box. Prompt tokenization produces a sequence of interleaved textual and visual tokens. We then follow the practice in Tsimpoukelli et al. (2021) and encode the prompt via a pre-trained T5 encoder (Raffel et al., 2020). Since T5 has been pre-trained on large text corpora, VIMA inherits the semantic understanding capability and robustness properties. To accommodate tokens from new modalities, we insert MLPs between non-textual tokens and T5.

A challenging aspect of designing a multi-task policy is to select a suitable conditioning mechanism. In our schema (Fig. 3), the robot controller (decoder) is conditioned on the prompt sequence P\mathcal{P} by a series of cross-attention layers between P\mathcal{P} and the trajectory history sequence H\mathcal{H}. We compute key KPK_{\mathcal{P}} and value VPV_{\mathcal{P}} sequences from the prompt and query QHQ_{\mathcal{H}} from the trajectory history, following the encoder-decoder convention in Raffel et al. (2020). Each cross-attention layer then generates an output sequence H′=softmax(QHKP⊺d)VP\mathcal{H}^{\prime}=\text{softmax}\left(\frac{Q_{\mathcal{H}}K^{\intercal}_{\mathcal{P}}}{\sqrt{d}}\right)V_{\mathcal{P}}, where dd is the embedding dimension. Residual connections are added to connect higher layers with the input rollout trajectory sequence. The cross-attention design enjoys three advantages: 1) strengthened connection to prompt; 2) intact and deep flow of the original prompt tokens; and 3) better computational efficiency. VIMA decoder consists of LL alternating cross-attention and self-attention layers. Finally, we follow common practice (Baker et al., 2022) to map predicted action tokens to discretized poses of the robot arm. See Appendix, Sec. C.2 for more details.

We follow behavioral cloning to train our models by minimizing the negative log-likelihood of predicted actions. Concretely, for a trajectory with TT steps, we optimize min⁡θ∑t=1T−log⁡πθ(at∣P,H)\min_{\theta}\sum_{t=1}^{T}-\log\pi_{\theta}(a_{t}|\mathcal{P},\mathcal{H}). The entire training is conducted on an offline dataset with no simulator access. To make VIMA robust to detection inaccuracies and failures, we apply object augmentation by randomly injecting false-positive detection outputs. After training, we select model checkpoints for evaluation based on the aggregated accuracy on a held-out validation set. The evaluation involves interacting with the physics simulator. We follow the best practices to train Transformer models. See Appendix, Sec. D for comprehensive training hyperparameters.

Experiments

In this section, we aim to answer three main questions:

What is the best recipe for building multi-task transformer-based robot agents with multimodal prompts?

What are the scaling properties of our approach in model capacity and data size?

How do different components, such as visual tokenizers, prompt conditioning, and prompt encoding, affect robot performance?

Because there is no prior method that works out of the box with our multimodal prompting setup, we make our best effort to select a number of representative transformer-based agent architectures as baselines, and re-interpret them to be compatible with VIMA-Bench:

Gato (Reed et al., 2022) introduces a decoder-only model that solves tasks from multiple domains where tasks are specified by prompting the model with the observation and action subsequence. For a fair comparison, we provide the same conditioning as VIMA, i.e., our multimodal encoded prompts. Input images are divided into patches and encoded by a ViT model to produce observation tokens. This variant is referred to as “VIMA-Gato”.

Flamingo (Alayrac et al., 2022) is a vision-language model that learns to generate textual completion in response to multimodal prompts. It embeds a variable number of prompt images into a fixed number of tokens via Perceiver (Jaegle et al., 2021b), and conditions the language decoder on the encoded prompt by cross-attention. Flamingo does not work with embodied agents out of the box. We adapt it to support decision-making by replacing the output layer with robot action heads. We denote the method as “VIMA-Flamingo”.

VIMA-GPT is a decoder-only architecture conditioned on tokenized multimodal prompts. It autoregressively decodes the next actions given instructions and interaction histories. Similar to prior work (Chen et al., 2021; Janner et al., 2021), it encodes an image into a single state token by a ViT encoder and prepends the rollout trajectory with prompt tokens. This baseline does not use cross-attention.

A more detailed comparison between these variants can be found in Appendix, Sec. C.1.

2 Evaluation Results

We compare VIMA against the baseline variants on four levels of generalization provided in our benchmark for different model and training dataset sizes. Our empirical results demonstrate that VIMA’s choice of object tokens combined with cross-attention conditioning is the most effective recipe among the model designs we consider.

We train all methods for a spectrum of model capacities from 2M to 200M parameters, evenly spaced on the log scale (Fig. 4). The encoder size is kept constant (T5-Base, 111M) for all methods and excluded from the parameter count. Across all levels of zero-shot generalization, we find that VIMA strongly outperforms other alternatives. Although models like VIMA-Gato and VIMA-Flamingo show improved performance with bigger model sizes, VIMA consistently achieves superior performance over all model sizes. We note that this can only be achieved with both cross-attention and object token sequence representations — altering any component will significantly degrade the performance, especially in the low model capacity regime (ablations in Sec. 5.3).

Next we investigate how different methods scale with varying dataset sizes. We compare model performance at 0.1%0.1\%, 1%1\%, 10%10\% and full imitation learning dataset provided in VIMA-Bench (Fig. 4). Note that to ensure all methods are fairly pre-trained on the same amount of data, we initialize baseline variants that directly learn from raw pixels with MVP pre-trained ViT (Xiao et al., 2022; Radosavovic et al., 2022). It is further MAE fine-tuned (He et al., 2021), using the same in-domain data as for the Mask R-CNN object detector. See Appendix, Sec. E.3 for detailed setup. VIMA is extremely sample efficient and, with just 1%1\% of the data, can achieve performance similar to baseline methods trained with 10×10\times more data on L1 and L2 levels of generalization. In fact, for L4 we find that with just 1%1\% of training data, VIMA already surpasses other variants trained with entire dataset. Finally, across all levels with just 10%10\% of the data, VIMA can outperform other architectures trained with the full dataset by a significant margin. We hypothesize that the data efficiency can be attributed to the object-centric representation employed in the VIMA recipe, which is less prone to overfitting than learning directly from pixels in the low-data regime. This is consistent with findings from Sax et al. (2018), which demonstrates that embodied agents conditioned on mid-level visual representations tend to be significantly more sample-efficient than end-to-end control from raw pixels.

Finally, we compare the relative performance degradation as we test the models on progressively challenging zero-shot evaluation levels without further fine-tuning (Fig. 5). Our method exhibits a minimal performance regression, especially between L1→L2L1\rightarrow L2 and L1→L3L1\rightarrow L3. In contrast, the baselines can degrade as much as 20%20\%, particularly in more difficult generalization scenarios. Although all methods degrade significantly when evaluated on L4L4 (Novel Tasks), the performance drop for VIMA is only half as severe as all other baselines. These results suggest that VIMA has developed a more generalizable policy and robust representations than the alternative approaches.

3 Ablation Studies

Through extensive experiments, we ablate different design choices in VIMA and study their impact on robot decision making. We focus on four aspects: visual tokenization, prompt conditioning, prompt-encoding language models, and policy robustness against distractions and corruptions.

As explained in Sec. 4, VIMA processes the prompt and observation images into a variable number of object tokens with a domain fine-tuned Mask R-CNN implementation. How important is this particular choice of visual tokenizer? We study 5 different variants and empirically evaluate their 4 levels of generalization performance on VIMA-Bench. 1) Ours (Oracle): instead of using Mask R-CNN, we directly read out the ground-truth bounding box from the simulator. In other words, we use a perfect object detector to estimate the upper bound on the performance of this study; 2) Object Perceiver: we apply a Perceiver module to convert the variable number of objects detected in each frame to a fixed number of tokens. Perceiver is more computationally efficient because it reduces the average sequence length; 3) Image Perceiver: the same architecture as the Perceiver Resampler in VIMA-Flamingo, which converts an image to a small, fixed number of tokens; 4) Image patches: following VIMA-Gato, we divide an RGB frame into square patches, and extract ViT embedding tokens. The number of patches is greater than the output of Image Perceiver; 5) Single image: VIMA-GPT’s tokenizer, which encodes one image into a single token.

Fig. 6 shows the ablation results. We highlight a few findings. First, we note that our Mask R-CNN detection pipeline (Appendix, Sec. C.4) incurs a minimal performance loss compared to the oracle bounding boxes, thanks to the object augmentation (Sec. 4) that boosts robustness during training. Second, tokenizing from raw pixels (Image Perceiver, patches, or single embedding) consistently underperforms our object-centric format. We hypothesize that these tokenizers have to allocate extra internal capacity to parse the objects from low-level pixels, which likely impedes learning. Sax et al. (2018) echoes our finding that using mid-level vision can greatly improve agent generalization compared to an end-to-end pipeline. Third, even though Ours and Object Perceiver both use the same object bounding box inputs, the latter is significantly worse in decision making. We conclude that it is important to directly pass the variable-length object sequence to the robot controller rather than downsampling to a fixed number of tokens.

VIMA conditions the robot controller (decoder) on the encoded prompt by cross-attention. A simple alternative is to concatenate the prompt P\mathcal{P} and interaction history H\mathcal{H} into one big sequence, and then apply a decoder-only transformer like GPT (Radford et al., 2018) to predict actions. In this ablation, we keep the object tokenizer constant and only switch the conditioning mechanism to causal sequence modeling. Note that this variant is conceptually “VIMA-Gato with object tokens”. Fig. 7 shows the comparison of VIMA (xattn) and the gpt-decoder variant across 4 generalization levels. While the variant achieves comparable performance in larger models, cross-attention still dominates in the small-capacity range and generalizes better in the most challenging L4 (Novel Task) setting. Our hypothesis is that cross-attention helps the controller stay better focused on the prompt instruction at each interaction step. This bears a resemblance to the empirical results in Sanh et al. (2021); Wang et al. (2022b), which show that well-tuned encoder-decoder architectures can outperform GPT-3 in zero-shot generalization.

We vary the size of the pre-trained T5 encoder to study the effect of prompt encoding. We experiment with three T5 capacities: small (30M), base (111M), and large (368M). We further fix the parameter count of the decision-making part to be 200M. For all T5 variants, we fine-tune the last two layers and freeze all other layers. We find no significant difference among the variants (Appendix, Sec. E.4), thus we set base as default for all our models.

We study the policy robustness against increasing number of distractors and corrupted task specifications, including incomplete prompts (randomly masking out words with token) and corrupted prompts (randomly swapping words, which could have changed the task meaning altogether). See Appendix, Sec. E.5 for exact setup and results. VIMA exhibits minimal performance degradation with increased distractors and minor decrease with corrupted prompts. We attribute this robustness to the high-quality pre-trained T5 backbone.

Related Work

Transformers (Vaswani et al., 2017) have enabled task unification across many AI domains (Brown et al., 2020; Chen et al., 2022a, b; Lu et al., 2022; Wang et al., 2022c). For example, in NLP, the Natural Language Decathlon (McCann et al., 2018) adopts a consistent question-answering format for a suite of 10 NLP tasks. T5 (Raffel et al., 2020) unifies all language problems into the same text-to-text format. GPT-3 (Brown et al., 2020) and Megatron (Shoeybi et al., 2019) demonstrate emergent behaviours of intuitive task specifications by zero-shot prompting. In computer vision, Pix2Seq (Chen et al., 2022b) casts many vision problems into a unified sequence format. Florence (Yuan et al., 2021), BiT (Kolesnikov et al., 2020), and MuST (Ghiasi et al., 2021) pre-train shared backbone models at scale for general visual representations and transfer them to downstream tasks. In multimodal learning, Perceiver (Jaegle et al., 2021b, a) proposes an efficient architecture to handle structured inputs and outputs. Flamingo (Alayrac et al., 2022) and Frozen (Tsimpoukelli et al., 2021) design a universal API that ingests interleaving sequences of images and text and generates free-form text. Gato (Reed et al., 2022) is a massively multi-task model across NLP, vision, and embodied agents. Our work is most similar in spirit to Gato, but we focus primarily on enabling an intuitive multimodal prompting interface for a generalist robot agent.

Foundation models (Bommasani et al., 2021) have demonstrated strong emergent properties. There are many ongoing efforts to replicate this success for embodied agents (Yang et al., 2023), focusing on 3 aspects. 1) Transformer agent architecture: Decision Transformer and Trajectory Transformer (Chen et al., 2021; Janner et al., 2021; Zheng et al., 2022; Xu et al., 2022, 2023) leverage the powerful self-attention models for sequential decision making. CLIPort (Shridhar et al., 2021), Perceiver-Actor (Shridhar et al., 2022), and RT-1 (Brohan et al., 2022) apply large transformers to robot manipulation tasks. BeT (Shafiullah et al., 2022) and C-BeT (Cui et al., 2022) design novel techniques to learn from demonstrations with multiple modes with transformers. 2) Pre-training for better representations: MaskViT (Gupta et al., 2022b), R3M (Nair et al., 2022), VIP (Ma et al., 2022), and VC-1 (Majumdar et al., 2023) pre-train general visual representations for robotic perception. Li et al. (2022b) fine-tunes from LLM checkpoints to accelerate policy learning. MineDojo (Fan et al., 2022) and Ego4D (Grauman et al., 2021) provide large-scale multimodal databases to facilitate scalable policy training. 3) LLMs for robot learning: SayCan (Ahn et al., 2022) leverages PaLM (Chowdhery et al., 2022) for zero-shot concept grounding. Huang et al. (2022a), Inner Monologue (Huang et al., 2022b) and LM-Nav (Shah et al., 2022) apply LLMs to long-horizon robot planning. PaLM-E (Driess et al., 2023) is instead a multimodal language model that can be repurposed for sequential robotic manipulation planning. Ours differs from these works in our novel multimodal prompting formulation, which existing LLMs do not easily support.

A wide range of robot manipulation tasks require different skills and task specification formats, such as instruction following (Stepputtis et al., 2020), one-shot imitation (Finn et al., 2017; Duan et al., 2017), rearrangement (Batra et al., 2020), constraint satisfaction (Brunke et al., 2021a), and reasoning (Shridhar et al., 2020). Multiple physics simulation benchmarks are introduced to study the above tasks. For example, iGibson (Shen et al., 2020; Li et al., 2021; Srivastava et al., 2021; Li et al., 2022a) simulates interactive household scenarios. Ravens (Zeng et al., 2020) and Robosuite (Zhu et al., 2020; Fan et al., 2021) design various tabletop manipulation tasks with realistic robot arms. CALVIN (Mees et al., 2021) develops long-horizon language-conditioned tasks. Meta-World (Yu et al., 2019) is a widely used simulator benchmark studying robotics manipulation with tabletop settings. CausalWorld (Ahmed et al., 2021) is a benchmark for causal structure and transfer learning in manipulation, requiring long-horizon planning and precise low-level motor control. AI2-THOR (Ehsani et al., 2021; Deitke et al., 2022) is a framework that supports visual object manipulation and procedural generation of environments. Our VIMA-Bench is the first robot learning benchmark to support multimodal-prompted tasks. We also standardize the evaluation protocol to systematically measure an agent’s generalization capabilities.

An extended review can be found in Appendix, Sec. F.

Conclusion

In this work, we introduce a novel multimodal prompting formulation that converts diverse robot manipulation tasks into a uniform sequence modeling problem. We instantiate this formulation in VIMA-Bench, a diverse benchmark with multimodal tasks and systematic evaluation protocols for generalization. We propose VIMA, a conceptually simple transformer-based agent capable of solving tasks such as visual goal reaching, one-shot video imitation, and novel concept grounding with a single model. Through comprehensive experiments, we show that VIMA exhibits strong model scalability and zero-shot generalization. Therefore, we recommend our agent design as a solid starting point for future work.

Acknowledgement

We are extremely grateful to Shyamal Buch, Jonathan Tremblay, Ajay Mandlekar, Chris Choy, De-An Huang, Silvio Savarese, Fei Xia, Josiah Wong, Abhishek Joshi, Soroush Nasiriany, and many other colleagues and friends for their helpful feedback and insightful discussions. We also thank the anonymous reviewers for offering us highly constructive advice and kind encouragement during the review period. NVIDIA provides the necessary computing resource and infrastructure for this project. This work is done during Yunfan Jiang and Guanzhi Wang’s internships at NVIDIA. Guanzhi Wang is supported by the Kortschak fellowship in Computing and Mathematical Sciences at Caltech.

References

Appendix A Simulator Details

We build our VIMA-Bench simulation suite upon the Ravens physics simulator (Zeng et al., 2020; Shridhar et al., 2021). Specifically, it is supported by PyBullet (Coumans & Bai, 2016–2021) with a Universal Robot UR5 arm. The size of the tabletop workspace is 0.5×10.5\times 1m. Our benchmark contains extensible sets of 3D objects and textures. Instantiated from an object-texture combination, all object instances can be rendered as RGB images appeared in multimodal prompts. Figure A.1 displays all 3D objects. Figure A.2 displays all textures.

The observation space of VIMA-Bench includes RGB images from both frontal and top-down views. It also includes a one-hot vector ∈{0,1}2\in\{0,1\}^{2} to indicate type of the end-effector ∈{suction cup, spatula}\in\{\text{suction cup, spatula}\}. While a suction cup is equipped in most manipulation tasks, a spatula is used in particular for visual constraint tasks, where an agent is asked to “wipe” objects. VIMA-Bench inherits the same action space from Zeng et al. (2020) and Shridhar et al. (2021), which consists of primitive actions of “pick and place” for tasks with a suction cup as the end effector, or “push” for tasks with a spatula. Both primitive actions contain two poses ∈\in SE(2)\mathbf{SE}(2) specifying target poses of the end effector. For the “pick and place” primitive, they represent the pick pose and the place pose. For the “push” primitive, they represent the push starting pose and push ending pose.

Similar to prior work (Zeng et al., 2020; Shridhar et al., 2021), VIMA-Bench provides scripted oracles to generate successful demonstrations for all tasks. We leverage them to construct an offline imitation dataset for behavioral cloning. Given a prompt, these programmed bots can access privileged information, such as the correct object to pick and target location to place.

Appendix B Task Suite

We develop 17 task templates that belong to 6 diverse categories. Thousands of individual task instances and their corresponding multimodal prompts can be procedurally generated from these task templates. We use PyBullet (Coumans & Bai, 2016–2021) as our backend and the default renderer to produce the RGB frames for training data and interactive test environments. For demonstration purpose, we apply the NVISII (Morrical et al., 2020) ray tracing to enhance the visual quality. We elaborate on each task in the following subsections.

This task category asks agents to follow basic instructions specified by multimodal prompts.

Pick the specified object(s) and place it (them) into the specified container.

Prompt: Put the {object}1 into the {object}2.

Description: The image placeholder {object}1 is the object to be picked and the {object}2 is the container object. The agent requires to recognize the objects with the correct color-shape combinations. To extend the difficulties, it supports more than one object to be picked or placed. For example, the prompt “Put the {object}1 and {object}2 into the {object}3” asks to pick two different objects and place into a target container. We uniformly sample different color-shape combos for objects to be picked and containers.

Success Criteria: All specified object(s) to pick are within the bounds of the container object(s), with specified shapes and textures provided in the prompt.

Oracle Trajectory: Shown in Fig. A.3 with its multimodal prompt.

In the workspace, put the objects with a specified texture shown in the scene image in the prompt into container object(s) with a specified color. This task requires the agent to find the correct object to manipulate by grounding the textural attributes from both natural language descriptions and the visual scene images.

Prompt: Put the {texture}1 object in {scene} into the {texture}2 object.

Description: The text placeholder {texture}1 and {texture}2 are sampled textures for objects to be picked and the container objects, respectively. The number of dragged objects with the same texture can be varied. {scene} is the workspace-like image placeholder. There is a designated number of distractors with different textures (and potentially different shapes) in the scene. For each distractor in the workspace, it has 50%50\% chance to be either dragged or container distractor object with different textures from those specified in the prompt.

Success Criteria: All objects in the workspace with {texture}1 are within the bounds of the container object with {texture}2.

Oracle Trajectory: Shown in Fig. A.4 with its multimodal prompt.

Rotate objects clockwise by certain degrees along zz-axis. Only rotationally asymmetric objects are considered in this task.

Prompt: Rotate the {object}1 {angles} degrees.

Description: The agent is required to rotate all objects in the workspace specified by the image placeholder {object}1. There are also objects with different color-shape combinations in the workspace as distractors. {angles} is the sampled degree that needs to be rotated. A target angle is sampled from 30∘30^{\circ}, 60∘60^{\circ}, 90∘90^{\circ}, 120∘120^{\circ}, and 150∘150^{\circ}.

Success Criteria: The position of the specified object matches its original position, and the orientation matches the orientation after rotating specific angles.

Oracle Trajectory: Shown in Fig. A.5 with its multimodal prompt.

B.2 Visual Goal Reaching

This task category requires agents to manipulate objects in the workspace to reach goal states represented as images shown in prompts.

Rearrange target objects in the workspace to match goal configuration shown in prompts. Note that to achieve the goal configuration, distractors may need to be moved away first.

Description: Objects in the scene placeholder {scene} are target objects to be manipulated and rearranged. In the workspace, the same target objects are spawned randomly, potentially with distractors randomly spawned as well. With a pre-defined distractor conflict rate, the position of each distractor has this probability to occupy the position of any target object such that the rearrangement can only succeed if moving away that distractor first.

Success Criteria: The configuration of target objects in the workspace matches that specified in the prompt.

Oracle Trajectory: Shown in Fig. A.6 with its multimodal prompt.

Extend the task 04 by requiring the agent to restore rearranged objects to the initial setup after the “rearranging” phase.

Prompt: Rearrange objects to this setup {scene} and then restore.

Description: Same as the task 04, except introducing the instruction “restore”.

Success Criteria: Meet the success criteria of the task 04, and then within the allowed max steps restore all target objects to their initial configurations.

Oracle Trajectory: Shown in Fig. A.7 with its multimodal prompt.

B.3 Novel Concept Grounding

This task category requires agents to ground new concepts of adjectives, nouns, or verbs via visual perception and language understanding. Similar task design can be found in prior work (Hill et al., 2021). Completing these tasks are challenging, because the model should a) first understand prompts with interleaved texts, images, and even video frames; b) quickly internalize new concepts that are different across task instances, which even tests the ability to meta-learn; and c) do complicated reasoning such as comparing between “taller” vs “less taller” vs “shorter” and then ground this reasoning into the robot action space.

Prompts consist of two parts: a definition part followed by an instruction part. In the definition part, novel concepts are defined by multimodal illustrations with multiple support examples. In the instruction part, agents are asked to achieve the goal by properly applying concepts from the definition part. The assignment of dummy object names is varied and independent for each task instance such that tasks can only be solved if the agent applies the reasoning correctly. This ability is also referred to as fast-mapping (Heibeck & Markman, 1987).

Ground comparative adjectives by comparing the size or the textural saturation of objects and manipulating the correct object(s) instructed in the prompt.

Prompt: {demo_object}1 is {novel_adj} than {demo_object}2. Put the {adv} {novel_adj} {object}1 into the {object}2.

Description: The sampled adjective {novel_adj} is a dummy adjective placeholder for agent to ground. By default, the novel adjective set is {daxer, blicker, modier, kobar}. The real meaning can be related to size (smaller/larger) or textural saturation (lighter/darker texture). The image placeholders {demo_object}1 and {demo_object}2 illustrate how the novel adjective is defined. For example, if the real comparison is ”taller”, then the sampled object in {demo_object}1 is taller than {demo_object}2. The choices of the novel adjective and the real meaning are independently sampled for different task instances. For the instruction part, this task is similar to task 01, where the agent is required to pick the specified object(s) with the novel adjective attribute and then place it into the specified container object. To avoid revealing the correct object to manipulate, we use a neutral texture for objects appeared in the instruction part.

Success Criteria: All target objects with the specified adjective attribute are within the bounds of the specified container object.

Oracle Trajectory: Shown in Fig. A.8 with its multimodal prompt.

Orthogonal to task 06 by requiring to learn mappings of novel nouns.

Prompt: This is a {novel_name}1 {object}1. This is a {novel_name}2\texttt{This is a \{novel\_name\}}_{1}\texttt{ \{object\}}_{1}\texttt{. This is a \{novel\_name\}}_{2}  {object}2. Put {novel_name}1 into a {novel_name}2.\texttt{ \{object\}}_{2}\texttt{. Put \{novel\_name\}}_{1}\texttt{ into a \{novel\_name\}}_{2}.

Description: Novel noun words are defined with the text placeholders {novel_name}1\texttt{\{novel\_name\}}_{1} and {novel_name}2\texttt{\{novel\_name\}}_{2}, following their image placeholders {object}1\texttt{\{object\}}_{1} and {object}2\texttt{\{object\}}_{2}, for the target object and container object, respectively. Novel nouns are sampled from {dax, blicket, wug, zup}. In the instruction part, objects are expressed as novel nouns defined in the previous definition part. Distractors are defined the same as task 01.

Success Criteria: All target object(s) are within the bounds of the container object(s).

Oracle Trajectory: Shown in Fig. A.9 with its multimodal prompt.

Prompt: This is a {novel_name}1 {object}1. This is a {novel_name}2\texttt{This is a \{novel\_name\}}_{1}\texttt{ \{object\}}_{1}\texttt{. This is a \{novel\_name\}}_{2}  {object}2. {demo_object}1 is {adj} than {demo_object}2. Put the \texttt{ \{object\}}_{2}\texttt{. \{demo\_object\}}_{1}\texttt{ is \{adj\} than \{demo\_object\}}_{2}.\texttt{ Put the } {adv} {novel_adj} {novel_name}1 into the {novel_name}2.\texttt{\{adv\} \{novel\_adj\} \{novel\_name\}}_{1}\texttt{ into the \{novel\_name\}}_{2}\texttt{.}

Description: See task description for task 06 and task 07.

Success Criteria: Similar as tasks 06 and 07.

Oracle Trajectory: Shown in Fig. A.10 with its multimodal prompt.

A novel verb “twist” is defined as rotating a specific angle illustrated by several examples. This task is similar to task 03, but it requires the agent to infer what is the exact angle to rotate from the prompt and to ground novel verbs that are semantically similar but different in exact definitions.

Prompt: "Twist" is defined as rotating object a specific angle. For examples: From {before_twist}i\texttt{For examples: From \{before\_twist\}}_{i} to {after_twist}i.\texttt{\{after\_twist\}}_{i}. Now twist all {texture} objects.

Description: Both {before_twist}i\texttt{\{before\_twist\}}_{i} and {after_twist}i\texttt{\{after\_twist\}}_{i} are scene placeholders where {before_twist}i\texttt{\{before\_twist\}}_{i} shows a randomly sampled object before “twisting” and {after_twist}i\texttt{\{after\_twist\}}_{i} shows the same object pose after “twisting”. All examples illustrate the same sampled angle to rotate. In the workspace, the target objects have the texture specified by {texture} and randomly sampled shapes.

Oracle Trajectory: Shown in Fig. A.11 with its multimodal prompt.

B.4 One-Shot Video Imitation

This task category requires agents to imitate motions demonstrated through videos shown in prompts. We follow prior works (Finn et al., 2017; Dasari & Gupta, 2020; Duan et al., 2017) to formulate the problem by giving one video demonstration (represented as key frames in prompts), then test the learned imitator’s ability to produce target trajectories. This setup is challenging because a) only one demonstration is available to the agent; b) the model needs to understand video frames interleaved with textual instructions; and c) missing correspondences between demonstrations and target trajectories since demonstrations only show partial key frames.

Prompt: Follow this motion for {object}: {frame}1...{frame}i...\texttt{Follow this motion for \{object\}}\texttt{: \{frame\}}_{1}\texttt{...\{frame\}}_{i}\texttt{...} {frame}n.\texttt{\{frame\}}_{n}\texttt{.}

Description: Image placeholder {object} is the target object to be manipulated and {{frame}i}\texttt{\{\{frame\}}_{i}\} is set of workspace-like scene placeholders to represent a video trajectory, where nn is the trajectory length. There is an object spawned at the center in both the workspace and the prompt video but with different textures as a distractor. The initial position of the target object matches that in {frame}1\texttt{\{frame\}}_{1}.

Success Criteria: In each step, the pose of the target object matches the pose in the corresponding video frame. Incorrect manipulation sequences are considered as failures.

Oracle Trajectory: Shown in Fig. A.12 with its multimodal prompt.

Stack objects with the order illustrated in the prompt video.

Prompt: Stack objects in this order {frame}1...{frame}i...{frame}n.\texttt{Stack objects in this order \{frame\}}_{1}\texttt{...\{frame\}}_{i}\texttt{...\{frame\}}_{n}\texttt{.}

Description: There are multiple objects with the same shape but different textures spawned in the workspace without any stacking initially. Distractor objects with different shapes are spawned in the workspace but not in the prompt video. At each step of the prompt video, one object is stacked over another or put at an empty position.

Oracle Trajectory: Shown in Fig. A.13 with its multimodal prompt.

B.5 Visual Constraint Satisfaction

This task category requires agents to wipe a specific number of objects in the workspace to a goal region while also satisfy the given visual constraint.

Sweep the designated number of objects into a specified region without exceeding the boundary.

Prompt: Sweep {quantifier} {object} into {bounds} without exceeding {constraint}.

Description: {object} is the image placeholder of the target object to be swept spawned with a random amount in the workspace. Distractors have the same amount, same shape, but different color from target objects. {quantifier} is the text placeholder to determine the target quantity of objects to be wiped, sampled from any, one, two, three, and all. {bounds} is the image placeholder for a three-sided rectangle as the goal region. {constraint} is the constraint line.

Success Criteria: The exact number of target objects to be swept are all inside the specified region. Potential failure cases include 1) any distractor being wiped into the region, 2) target object exceeding the constraint, or 3) incorrect number of target objects being swept into the goal region.

Oracle Trajectory: Shown in Fig. A.14 with its multimodal prompt.

Sweep the designated number of objects into a specified region without touching the constraint.

Prompt: Sweep {quantifier} {object} into {bounds} without touching {constraint}.

Description: Similar as task 12 but requiring a different way to satisfy the constraint. The agent has to learn to avoid contacting the constraint line in this case.

Success Criteria: Similar as task 12 except that the constraint is to not touch the red line.

Oracle Trajectory: Shown in Fig. A.15 with its multimodal prompt.

B.6 Visual Reasoning

This task category requires agents to make decisions by reasoning over or memorizing information conveyed through multimodal prompts.

By reasoning the “same texture”, the agent is required to pick all objects in the workspace with the same texture as the container objects specified in the prompt and place them into it.

Prompt: Put all objects with the same texture as {object} into it.

Description: {object} is the sampled goal container object. In the workspace, there are objects with the same texture as the container but potentially different shapes. Distractors with different textures are spawned.

Success Criteria: All objects with the same texture as the goal container are within the bounds of the container.

Oracle Trajectory: Shown in Fig. A.16 with its multimodal prompt.

By reasoning the “same shape”, the agent is required to pick all objects in the workspace with the same top-down profile as the goal container specified in the prompt and place them into it. For example, blocks and boxes have the same rectangular profile.

Prompt: Put all objects with the same profile as {object} into it.

Description: Similar to the task 14 except the objects to be picked and placed have the same shape. There are three different shapes: rectangular-like (e.g. block and pallet), circle-like (e.g. ring and bowl), and undetermined for the rest.

Success Criteria: All objects with the same shape as the container are within the container.

Oracle Trajectory: Shown in Fig. A.17 with its multimodal prompt.

Put the target object into the container, and then put one of its old neighbors into the same container.

Prompt: First put {object}1 into {object}2 then put the object that was previously at its {direction} into the same {object}2.

Description: Objects in image placeholders {object}1\texttt{\{object\}}_{1} and {object}2\texttt{\{object\}}_{2} are the target object to be picked and the container, respectively. We then ask the agent to put one of old neighbors of the previous target object into the same container. The old neighboring object is specified through cardinal directions {north, south, west, east}.

Success Criteria: The target object and the correct neighboring object are inside the container.

Oracle Trajectory: Shown in Fig. A.18 with its multimodal prompt.

Pick and place the target object specified in the prompt into different containers in order then restore to the initial container.

Prompt: Put {object}1 into {object}2 . Finally restore it into its original container.

Description: The object in the image placeholder {object}1 is the target object to be manipulated across the task. There are more than one target containers (e.g. “Put {object}1 into {object}2 then {object}3. Finally restore it into its original container” for two target containers to be placed in order). The rest of spawned containers naturally becomes distractors.

Success Criteria: The target object is first put into multiple containers following the specific order. Finally it should be restored into its original container.

Oracle Trajectory: Shown in Fig.A.19 with its multimodal prompt.

Appendix C Model Architecture

In this section, we provide comprehensive details about VIMA model architecture as well as other adapted baseline methods. We implement all models in PyTorch (Paszke et al., 2019) and adapt Transformer-related implementation from Wolf et al. (2019).

We summarize differences between VIMA and other baseline variants in Table 1. In the column “Prompt Conditioning”, an alternative to cross-attention is to first concatenate prompt and interaction into a big sequence, then repetitively apply transformer decoders to predict actions. It is referred to as “Direct modeling”. The relative computation cost is quadratically proportional to the number of observation tokens.

C.2 VIMA Architecture

As introduced in Section 4, there are 3 types of input formats in multimodal prompts, namely (1) text inputs, (2) images of full scenes, and (3) images of single objects.

For text inputs, we follow the standard pipeline in NLP to first tokenize raw language to discrete indices through pre-trained t5-base tokenizer. We then obtain corresponding word tokens from the embedding look-up of the pre-trained t5-base model. For images of full scenes, we first parse the scene through a fine-tuned Mask R-CNN detection model (He et al., 2017; Wu et al., 2019) to extract individual objects. Each object representation contains a bounding box and a cropped image. The bounding box is in the format of [xcenter,ycenter,height,width]\begin{bmatrix}x_{\text{center}},y_{\text{center}},\text{height},\text{width}\end{bmatrix}. We normalize it to be within $$ by dividing each dimension with corresponding upper-bound value. We then pass it through a bounding box encoder MLP and obtain a feature vector. To process the cropped image, we first pad non-square image to a square by padding along the shorter dimension. We then resize it to a pre-configured size and pass it through a ViT (trained from scratch) to obtain the image feature. Finally, an object token is obtained by concatenating the bounding box feature and the image feature and mapping to the embedding dimension. For images of single objects, we obtain tokens in the same way except with a dummy bounding box. Detailed model hyperparameters about tokenization are listed in Table 2.

After obtaining a sequence of prompt tokens, we follow Tsimpoukelli et al. (2021) to pass it through a pre-trained t5-base encoder to obtain encoded prompt. Note that we add adapter MLP between object tokens and the T5 encoder. To prevent catastrophic forgetting, VIMA only fine-tunes the last two layers of the language encoder with layer-wise learning rate decay (He et al., 2021) but freezes all other layers. We adopt learned absolute positional embedding. Model hyperparameters are listed in Table 2 as well.

C.2.2 Observation Encoding

Since all RGB observations are images of full scenes, we follow the same procedure discussed above to obtain flattened object tokens. Because we provide RGBs from two views (frontal and top-down), we order object tokens by following the order of [frontal, top-down]\begin{bmatrix}\text{frontal, top-down}\end{bmatrix}. We one-hot encode the state of the end effector. We then concatenate object tokens with the end-effector state and transform to observation tokens. We adopt learned absolute positional embedding. Detailed model hyperparameters about observation encoding is provided in Table 3.

C.2.3 Action Encoding

Since our model is conditioned on observation-action interleaved history, we also tokenize past actions. We follow common practice in Chen et al. (2021); Zheng et al. (2022) to encode past actions with a two-layer MLP. It has a hidden dimension of 256. We then map outputs to token dimension and obtain action tokens.

C.2.4 Sequence Modeling

The robot controller in VIMA is a causal decoder that autoregressively predicts actions. To condition the decoder on prompt tokens, we perform cross-attention between history tokens and prompt tokens (Figure 3). Concretely, we pass history tokens as the query sequence and prompt tokens as the key-value sequence into cross-attention blocks. The output prompt-aware trajectory tokens then go through causal self-attention blocks. We alternate cross-attention and self-attention LL times. This procedure is technically described in Pseudocode LABEL:supp:code:xattn.

C.2.5 Action Decoding

After obtaining the predicted action token, we map it to the action space A\mathcal{A} and obtain the predicted action. This is achieved though a group of action heads. Since the action space consists of two SE(2)\mathbf{SE}(2) poses, for each pose we use six independent heads to decode discrete actions (two for xy coordinate and four for rotation represented in quaternion). These discrete actions are then integrated and mapped to continuous actions through affine transformation. The two poses are modeled independently. Early ablations show that this independent modeling is equally good as alternative techniques, such as autoregressive decoding (Vinyals et al., 2019; OpenAI et al., 2019). Detailed model hyperparameters are listed in Table 4.

C.3 Baselines Architectures

In this section, we elaborate model architectures for adapted baseline methods. Some components such as the action decoder are same across all models. Therefore, we only discuss unique model components.

Gato (Reed et al., 2022) introduces a decoder-only model that solves tasks from multiple domains including robotics, video game, image captioning, language modeling, etc. Different tasks are specified by supplying the model with an initial sequence of corresponding tokens. For example, in tasks involving decision making, these tokens include observation and action tokens. For fair comparison, we provide the same conditioning as VIMA, i.e., our multimodal tokenized prompts. This adapted baseline variant is referred to as “VIMA-Gato”. Similar to our method, VIMA-Gato also predicts actions in an autoregressive manner. VIMA-Gato and our method share the same training philosophy to only optimize the causal behavior cloning objective. However, unlike our method that adopts an object-centric representation to treat individual objects as observation tokens, VIMA-Gato divides input images into patches and encodes them by a ViT (Dosovitskiy et al., 2020) to produce observation tokens. Furthermore, VIMA-Gato relies on causal self-attention to model entire trajectory sequences starting with prompt tokens. Hyperparameters of VIMA-Gato’s ViT is listed in Table 5. The transformer-decoder style sequence modeling is technically illustrated in Pseudocode LABEL:supp:code:selfattn.

C.3.2 VIMA-Flamingo

Flamingo (Alayrac et al., 2022) is a vision-language model that learns to generate textual completion in response to multimodal prompts. It embeds a variable number of prompt images into a fixed number of tokens via the Perceiver Resampler module (Jaegle et al., 2021b), and conditions the language decoder on encoded prompts by cross-attention. Flamingo does not work with embodied agents out of the box. We adapt it by replacing the output layer with robot action heads (hyperparameters listed in Table 4) and using tokenized rollout histories as inputs. We thus call it “VIMA-Flamingo”. We train it end-to-end with causal behavior cloning loss. VIMA-Flamingo differs from ours since it processes image observations into a fixed number of visual tokens through a learned Perceiver Resampler. Model hyperparameters for our reimplementation of the Perceiver Resampler is listed in Table 6.

C.3.3 VIMA-GPT

VIMA-GPT is a GPT-based behavior cloning agent conditioned on tokenized multimodal prompts with the GPT architecture. It autoregressively decodes next actions given multimodal prompts and interaction histories. We optimize this method end-to-end with causal behavior cloning loss. Similar to prior works of casting RL problems as sequence modeling (Chen et al., 2021; Janner et al., 2021; Zheng et al., 2022), it encodes an image into a single “state” token through a learned ViT encoder. It also directly models entire trajectory sequences prepended with prompt tokens. Therefore, it differs from our method in the representation of observation tokens and prompt conditioning. For visual tokenizer, we employ a learned ViT with hyperparameters listed in Table 5.

C.4 Mask R-CNN Detection Model

Finally, we elaborate on the mask R-CNN model (He et al., 2017) for scene parsing and object extraction. We fine-tune a pre-trained lightweight mask R-CNN (mask_rcnn_R_50_FPN_3x) from Wu et al. (2019) to adapt to scenes and images in our tabletop environment. We fine-tune it on a subset of agent training dataset. It contains 100 trajectories for each task, resulting in 22,741 images and 61,822 annotations in total. We use learning rate 5×10−45\times 10^{-4} and train for 10 epochs. During model selection, we particularly favor models with high recall to reduce the number of missed objects. To compensate for resulting false-positives, we adopt object augmentation during agent training (Appendix, Sec. D).

A visualization of its output is provided in Figure A.20. We do not use the predicted object names in our models.

Appendix D VIMA Training Details

We follow the best practice to train Transformer models using the AdamW optimizer (Loshchilov & Hutter, 2019), learning rate warm-up, cosine annealing (Loshchilov & Hutter, 2017), etc. Training hyperparameters are provided in Table 7. We use GEGLU activation (Shazeer, 2020) inside Transformer models across all methods.

To make trained models robust to detection inaccuracies and failures, we apply object augmentation by randomly injecting false-positive detection outputs. Concretely, for observation at each time step, we sample number of augmented objects i.i.d. naugmented objects∼Cat(K,p)n_{\text{augmented objects}}\sim\text{Cat}(K,\mathbf{p}), where Cat(⋅)\text{Cat}(\cdot) denotes a categorical distribution with KK supports parameterized by p\mathbf{p}. For each augmented object, we then randomly sample a bounding box and corresponding cropped image to add to object tokens. In our experiments, we set p={0:0.95,1:0.05}\mathbf{p}=\{0:0.95,1:0.05\} with K=2K=2.

We train a spectrum of 7 models ranging from 2M to 200M parameters. To vary the model capacity, we follow prior work (Chowdhery et al., 2022) to change embedding dimension and number of layers. We list configurations for methods with cross-attention prompt conditioning (i.e., ours and VIMA-Flamingo) in Table 8, and configurations for methods only with causal self-attention (i.e., VIMA-Gato and VIMA-GPT) in Table 9.

Appendix E Extended Experiment Results

All experiments are conducted on cluster nodes, each with 8 NVIDIA V100 GPUs. The largest experiment takes approximately one day. We utilize DDP (distributed data parallel) to accelerate the training.

E.2 Model Scaling

We present numerical results that constitute Fig. 4 in Table 10. The claim of “up to 2.9×2.9\times improvement” made in Abstract and Sec. 1 is calculated as follows. The best competing variant is VIMA-Gato. On the hardest L4, our method shows the most significant relative improvement with a model size of 20M. We compute the performance gap, divide by VIMA-Gato’s performance, and only keep the first digit after decimal to obtain the result.

E.3 Data Scaling

To ensure all methods are fairly pre-trained on the same amount of data (i.e., they have roughly the same amount of built-in information, thus the x-axis in Fig. 4 faithfully corresponds to the extra bits of information seen during further training), we initialize variants that directly learn from raw pixels with MVP pre-trained ViT (Xiao et al., 2022; Radosavovic et al., 2022). It is further MAE fine-tuned (He et al., 2021), using the same in-domain data as for the Mask R-CNN object detector. Note that the MVP pre-trained then domain fine-tuned ViT also updates weights jointly with robot controllers later on. We use the ViT-B backbone from MVP. The in-domain data for fine-tuning include 100 trajectories for each task.

E.3.2 Numerical Results

We present numerical results that constitute Fig. 4 in Table 11. The claim of “2.7×2.7\times improvement” made in Abstract and Sec. 1 is calculated as follows. The best competing variant is VIMA-Gato that achieves 12.2%12.2\% average success rate trained with full data on L4. Our method trained with 10%10\% data achieves 46%46\% average success rate on the same level. We compute the performance gap, divide by VIMA-Gato’s performance, and only keep the first digit after decimal to obtain the result.

E.3.3 What If Baseline Variants’ ViT Is Trained from Scratch?

We further investigate what if baseline variants’ ViT is trained from scratch and end-to-end with the robot controllers. We visualize the results in Fig. A.21 and numerically present them in Table 12. We annotate with arrows to indicate performance increase (↑\uparrow) and decrease (↓\downarrow). We highlight two findings.

First, MVP pre-trained ViT is most beneficial in the setting with sufficient in-domain training data (i.e., the 10% data scenario). It boosts the performance for the most competing baseline variant VIMA-Gato. However, in other settings with abundant in-domain data (i.e., the full data scenario) or insufficient in-domain data (i.e., 1% and 0.1% scenarios), the advantage of MVP pre-trained ViT diminishes and it even becomes detrimental. This aligns with the finding in previous empirical studies (Hansen et al., 2022). Second, in settings with reasonable amounts of in-domain data (i.e., the 1%, 10%, and 100% scenarios), our recommended recipe always outperforms other variants. We notice that such a data demand generally can be satisfied by both simulated robotics data (Mandlekar et al., 2021) and real robotics data (Dasari et al., 2019; Brohan et al., 2022). Therefore, it demonstrates that our recommended recipe is highly sample-efficient compared to alternative designs, especially under practical settings.

E.4 Vary T5 Encoder Sizes

We vary the size of the pre-trained T5 encoder (Raffel et al., 2020) to study the effect of prompt encoding. We experiment with three T5 model capacities: t5-small (30M), t5-base (111M), and t5-large (368M). For all T5 variants, we fine-tune the last two layers and freeze all other layers. We fix the parameter count of the decision-making part to be 200M. As shown in Table 13, we find no significant difference among the variants. Thus we set the standard t5-base as default for all our models.

E.5 Policy Robustness

We study the policy robustness against increasing amounts of distractors in scenes. For all tasks being evaluated, we add one more distractor object. We run our largest VIMA model with 200M parameters. The result is presented in Table 14.

It turns out that the performance of VIMA degrades minimally with more distractors than the training distribution. This indicates that our agent has learned a reasonably robust policy against objects that are irrelevant to the task.

We then study the policy robustness against imperfect prompts, including incomplete prompts (randomly masking out words with token) and corrupted prompts (randomly swapping words, which could have changed the task meaning altogether). We run our largest VIMA model with 200M parameters, results are shown in Table 15.

Our well-trained model exhibits minimal performance decrease when evaluated on masked prompts and minor decrease on corrupted prompts. We attribute this robustness to the high-quality pre-trained T5 language backbone.

Appendix F Extended Related Work

In this section, we provide an extended review of related work as complementary to Section 6.

In computer vision, Mask R-CNN (He et al., 2017), UberNet (Kokkinos, 2016), and 12-in-1 (Lu et al., 2020) leverage a single backbone model with multiple independent heads for different tasks. UVim (Kolesnikov et al., 2022) is another unified approach for vision that uses a language model to generate the guiding code for a second model to predict raw vision outputs. In multimodal learning, numerous works (Lu et al., 2022; Wang et al., 2022a; Zellers et al., 2021, 2022; Buch et al., 2022; Fu et al., 2021; Yang et al., 2022) investigate the unification of image, video, audio, and/or language modalities to deliver multi-purpose foundation models, although most of which are not equipped with decision-making capabilities. BEiT-3 (Wang et al., 2022c) performs masked data modeling on images, texts and image-text pairs to pre-train a backbone for various downstream tasks. MetaMorph (Gupta et al., 2022a) learns a universal controller over a modular robot design space.

Embodied agent research (Duan et al., 2022; Batra et al., 2020; Ravichandar et al., 2020; Collins et al., 2021) is adopting the large-scale pre-training paradigm (Yang et al., 2023), powered by a collection of learning environments (Abramson et al., 2020; Shridhar et al., 2020; Savva et al., 2019; Puig et al., 2018; Team et al., 2021; Toyama et al., 2021; Shi et al., 2017). From the aspect of pre-training for better representations, Reid et al. (2022) fine-tunes from LLM checkpoints to accelerate policy learning. LaTTe (Bucker et al., 2022) and Embodied-CLIP (Khandelwal et al., 2021) leverage the frozen visual and textual representations of CLIP (Radford et al., 2021) for robotic manipulation. MaskDP (Liu et al., 2022a) pre-trains bidirectional transformers for various downstream embodied tasks. From the perspective of leveraging transformer as agent architecture, methods such as Dasari & Gupta (2020) and MOSAIC (Zhao et al., 2022) achieve superior performance in one-shot video imitation tasks. They both use the self-attention mechanism with auxiliary losses such as inverse dynamics loss (Dasari & Gupta, 2020) and contrastive loss (Zhao et al., 2022) to learn robot controllers. InstructRL (Liu et al., 2022b) leverages jointly pre-trained vision-language models as robot agents to perform manipulation tasks. From the perspective of large language models for robot learning, Socratic Models (Zeng et al., 2022) composes multiple vision and language foundation models for multimodal reasoning in videos. ROSIE (Yu et al., 2023) leverages text-to-image diffusion models to augment existing robotic dataset (Brohan et al., 2022) via inpainting. MOO (Minderer et al., 2022) adopts a similar object-centric representation as ours for open-world object manipulation. Furthermore, Voyager (Wang et al., 2023) develops a LLM-powered agent operating in an open-ended virtual world (Fan et al., 2022).

There are many prior works that are not mentioned in the main paper that study different robotic manipulation tasks, such as instruction following (Shridhar et al., 2021; Lynch & Sermanet, 2021), constraint satisfaction (Bharadhwaj et al., 2021; Srinivasan et al., 2020; Thananjeyan et al., 2021), one-shot imitation (Paine et al., 2018; Huang et al., 2019; Dasari & Gupta, 2020; Aceituno et al., 2021; Zhao et al., 2022), rearrangement (Weihs et al., 2021; Szot et al., 2021; Liu et al., 2021; Ehsani et al., 2021; Gan et al., 2021; Stengel-Eskin et al., 2022), and reasoning (Gupta et al., 2019; Ahmed et al., 2021; Toyer et al., 2020; Lim et al., 2021). Multiple simulation benchmarks are introduced to study the above tasks: 1) Indoor simulation environments: Habitat (Savva et al., 2019; Szot et al., 2021) is equipped with a high-performance 3D simulator for fast rendering and proposes a suite of common tasks for assistive robots. 2) Tabletop environments: RLBench (James et al., 2019) and SURREAL (Fan et al., 2018, 2019) are other widely used simulator benchmarks studying robotics manipulation with tabletop settings. STRETCH-P&P (Zhang & Weihs, 2023) studies generalization across goals for reset-free reinforcement learning. All these aforementioned simulators and benchmarks do not natively support task specification and prompting with multiple modalities.

Appendix G Limitations and Further Discussions

VIMA inherits the errors from the standalone object detector, which may struggle in the cases of occlusion or out-of-distribution object forms. However, using object detectors is not entirely without merits. First, it allows us to seamlessly switch to stronger detection models when they become available. For example, we can switch to object detectors that are more robust and open-vocabulary, such as OWL-ViT (Minderer et al., 2022). This would enable VIMA to transfer to real-world scenarios with minimal modifications. Second, by leveraging pre-trained vision pipelines, several concurrent works have demonstrated the superiority of object-centric representation in robot manipulation. For example, VIOLA (Zhu et al., 2022) achieves better performance with a pre-trained Region Proposal Network (Ren et al., 2015). MOO (Stone et al., 2023) also shows that a robot agent with OWL-ViT (Minderer et al., 2022) as the object detector significantly outperforms RT-1 (Brohan et al., 2022), which directly learns from raw pixels, on various real-world manipulation tasks. In fact, MOO (Stone et al., 2023) includes a baseline called “VIMA-like” that already demonstrates strong performance on real robots under real-world scenarios. As we witness image segmentation is becoming more robust and general-purpose (Kirillov et al., 2023), we envision such design choice will become more effective and further gain more popularity.

Our goal with VIMA-Bench is to explore the multi-task ability, generalization, and understanding of multi-modality. Therefore, these aspects are not the primary focus of this work. However, we envision future works can combine this formulation with more physically realistic simulators such as Zhu et al. (2020), Srivastava et al. (2021), and Mittal et al. (2023).

We inherit the same high-level action space from well-established prior works, such as Transporter (Zeng et al., 2020). While “pick-and-place” and “wipe” seem simple, they do cover a wide range of tabletop manipulation tasks and are crucial to industrial use cases like warehouse robots (Yoon et al., 2003; Berscheid et al., 2020; Devin et al., 2020; Song et al., 2019). While VIMA is currently using these two actions, the algorithm design is general-purpose and does not make assumptions about the particular action choices. For example, VIMA would require only minimal modifications to support more low-level action spaces like joint-torque control.

Appendix H Full Tables

This section contains more detailed tables that correspond to the results in Figure 4. Specifically, we show breakdown results on each task that constitute the model scaling results in Tables 16, 17, 18, and 19.