VIMA: G ENERAL R OBOT M ANIPULATION WITH M ULTIMODAL P ROMPTS
Anonymous authors
Paper under double-blind review
A BSTRACT
Prompt-based learning has emerged as a successful paradigm in natural languageprocessing, where a single general-purpose language model can be instructed toperform any task specified by input prompts.
Yet task specification in roboticscomes in various forms, such as imitating one-shot demonstrations, following language instructions, and reaching visual goals.
They are often considered differenttasks and tackled by specialized models.
This work shows that we can express awide spectrum of robot manipulation tasks with multimodal prompts , interleavingtextual and visual tokens.
We design a transformer-based robot agent, VIMA,that processes these prompts and outputs motor actions autoregressively.
To trainand evaluate VIMA, we develop a new simulation benchmark with thousands ofprocedurally-generated tabletop tasks with multimodal prompts, 600K+ expert trajectories for imitation learning, and four levels of evaluation protocol for systematicgeneralization.
VIMA achieves strong scalability in both model capacity and datasize.
It outperforms prior SOTA methods in the hardest zero-shot generalizationsetting by up to 2 .
9 × task success rate given the same training data.
With 10 × lesstraining data, VIMA still performs 2 .
7 × better than the top competing approach.
Video demos are available at https://iclr3081.github.io/ .
I NTRODUCTION
Transformers have given rise to remarkable multi-task consolidation across many AI domains.
Forexample, 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 naturallanguage understanding task to a general-purpose model.
We envision that a generalist robot agent should have a similarly intuitive and expressive interfacefor task specification.
What does such an interface for robot learning look like?
As a motivatingexample, consider a personal robot tasked with household activities.
We can ask the robot to bring usa cup of water by a simple natural language instruction.
If we require more specificity, we can insteadinstruct the robot to “bring me <image of the cup> ”.
For tasks requiring new skills, the robotshould be able to adapt preferably from a few video demonstrations (Duan et al., 2017).
Tasks thatneed interaction with unfamiliar objects can be easily explained via a few image examples for novelconcept grounding (Hermann et al., 2017).
Finally, to ensure safe deployment, we can further specifyvisual constraints like “do not enter <image> 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 manipulationtasks into one sequence modeling problem; 2) a new robot agent model capable of multi-taskand zero-shot generalization; and 3) a large-scale benchmark with diverse tasks to systematicallyevaluate the scalability and generalization of our agents.
We start with the observation that many robot manipulation tasks can be formulated by multimodalprompts 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 tomatch this { scene image } ”; Novel Concept Grounding looks like “This is a dax { new object } 1and this is a blicket { new object } 2 . Put two metal dax on the marble blicket.”
; Few-shot Imitation can embed video snippet in the prompt “Follow this motion trajectory for the wooden cube: [Figure]
for
This is a blicket
This is a wug
Put a wug into a blicket
Sweep all into without exceeding
One-shot Demonstration
Visual Goal: Rearrangement
[Figure]
{ frame 1 } , { frame 2 } , { frame 3 } , { frame 4 } ”; and expressing Visual Constraint is as simple asadding the clause “without touching { safety boundary } ”.
Multimodal prompts not only have more expressive power than individual modalities, but also enablea uniform sequence IO interface for training generalist robot agents.
Previously, different robotmanipulation tasks require distinct policy architectures, objective functions, data pipelines, andtraining 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, ourmultimodal prompt interface allows us to harness the latest advances in large transformer models (Linet al., 2021; Tay et al., 2020; Khan et al., 2021) for developing scalable multi-task robot learners.
To this end, we design a novel Vi suo M otor
A ttention model (VIMA).
The architecture follows theencoder-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-trainedlanguage model (Tsimpoukelli et al., 2021), and decodes robot control actions autoregressivelyfor each environment interaction step.
The transformer decoder is conditioned on the prompt viacross-attention layers that alternate with the usual causal self-attention.
Instead of operating on rawpixels, VIMA adopts an object-centric approach.
We parse all images in the prompt or observationinto objects by off-the-shelf detectors (He et al., 2017), and flatten them into sequences of objecttokens.
All these design choices combined deliver a conceptually simple architecture with strongmodel and data scaling properties.
To systematically evaluate our proposed algorithm, we introduce a new benchmark (VIMA-B ENCH )built on the Ravens simulator (Zeng et al., 2020; Shridhar et al., 2021).
We provide 17 representativemeta-tasks with multimodal prompt templates, which can be procedurally instantiated into thousands of individual tasks by various combinations of textures and tabletop objects.
VIMA-B ENCHestablishes a 4-level protocol to evaluate progressively stronger generalization capabilities, from randomized object placement to novel tasks altogether (Fig. 2).
To demonstrate the scalability of VIMA,we train a spectrum of 7 models ranging from 2M to 200M parameters.
Our approach outperformsstrong prior SOTA methods such as Gato (Reed et al., 2022), Decision Transformer (Chen et al.,2021), and Flamingo (Alayrac et al., 2022) across all 4 levels of zero-shot generalization and allmodel capacities, sometimes by a large margin (up to 2 .
9 × task success rate given the same amount oftraining data, and 2 .
7 × better even with 10 × less data).
We plan to open-source the simulation environment, training dataset, algorithm code, and pre-trained model checkpoints to ensure reproducibilityand facilitate future works from the community.
We attach supplementary materials as Appendix tothis PDF, and present video demos and anonymized code at https://iclr3081.github.io/ .
R ELATED W ORK
Multi-task Learning by Sequence Modeling.
Transformers have enabled task unification acrossmany AI domains (Raffel et al., 2020; Brown et al., 2020;
Chen et al., 2022a;b; Lu et al., 2022; Wanget al., 2022c; Alayrac et al., 2022).
For example, in NLP , T5 (Raffel et al., 2020) unifies all languageproblems into the same text-to-text format.
GPT-3 (Brown et al., 2020), PaLM (Chowdhery et al.,2022), and Megatron (Shoeybi et al., 2019) demonstrate emergent behaviours of intuitive task specifications by zero-shot prompting.
In computer vision , Florence (Yuan et al., 2021), BiT (Kolesnikovet al., 2020), and MuST (Ghiasi et al., 2021) pre-train a shared backbone model at scale for generalvisual representations and transfer it to downstream tasks.
Pix2Seq (Chen et al., 2022b) casts manyvision problems into a unified sequence format.
In multimodal learning , Flamingo (Alayrac et al.,2022) and Frozen (Tsimpoukelli et al., 2021) design a universal API that ingests an interleavingsequence of images and text and generates free-form text.
Gato (Reed et al., 2022) is a massively multitask model across NLP, vision, and embodied agents.
Our work is most similar in spirit to Gato, but wefocus primarily on enabling an intuitive, multimodal prompting interface for a generalist robot agent.
Foundation Models for Embodied Agents.
Foundation models (Bommasani et al., 2021; Brownet al., 2020; Raffel et al., 2020; Ramesh et al., 2022; Wei et al., 2022) have demonstrated strong emergent properties like zero-shot prompting and complex reasoning.
There are many ongoing efforts toreplicate this success for embodied agents, focusing on 3 aspects: 1) Transformer agent architecture : Decision Transformer (Chen et al., 2021; Janner et al., 2021; Zheng et al., 2022; Xu et al., 2022) and
Gato (Reed et al., 2022) leverage the powerful self-attention models for sequential decision making.
CLIPort (Shridhar et al., 2021) and Perceiver-Actor (Shridhar et al., 2022) apply large transformers torobot manipulation tasks.
Methods such as Dasari & Gupta (2020) and MOSAIC (Zhao et al., 2022)also leverage transformers to achieve superior performance in one-shot video imitation tasks.
Pretraining for better representations : Masked ViT (Gupta et al., 2022b), R3M (Nair et al., 2022), and Parisi et al.
(2022) pre-train general visual representations for robotic perception.
Li et al.
(2022); Reidet al.
(2022) finetune 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.
Large language models for robot learning : SayCan (Ahn et al., 2022) leverages the500B PaLM (Chowdhery et al., 2022) for zero-shot concept grounding.
Socratic Models (Zeng et al. ,2022) composes multiple vision and language foundation models (VLMs) for multimodal reasoningin videos.
Huang et al.
(2022a),
Inner Monologue (Huang et al., 2022b) and LM-Nav (Shah et al. ,2022) successfully apply LLMs to long-horizon robot planning.
VIMA differs from these works inour novel multimodal prompting formulation, which existing LLMs and VLMs do not easily support.
Robot Manipulation and Benchmarks.
There are a wide range of robot manipulation tasks thatrequire different skills and task specification formats, such as instruction following (Stepputtis et al. ,2020; Shridhar et al., 2021; Lynch & Sermanet, 2021), one-shot imitation (Finn et al., 2017; Dasari& Gupta, 2020; Duan et al., 2017), rearrangement (Batra et al., 2020;
Weihs et al., 2021; Szot et al.,2021), constraint satisfaction (Brunke et al., 2021a; Srinivasan et al., 2020; Thananjeyan et al., 2021),and reasoning (Shridhar et al., 2020; Gupta et al., 2019;
Ahmed et al., 2021; Toyer et al., 2020; Limet al., 2021).
Multiple physics simulation benchmarks are introduced to study the above tasks.
Forexample, iGibson (Shen et al., 2020; Li et al., 2021; Srivastava et al., 2021) simulates interactivehousehold 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. MOSAIC (Zhao et al., 2022)features a challenging benchmark built on top of Zhu et al. (2020) for one-shot imitation learning.
Our VIMA-B ENCH is the first robot learning benchmark to support multimodal-prompted tasks.
We alsostandardize the evaluation protocol to systematically measure an agent’s generalization capabilities.
A more extended literature review can be found in Appendix, Sec.
F.
M ULTIMODAL P ROMPTS FOR T ASK S PECIFICATION
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 wellacross tasks.
Our key insight is that various task specification paradigms (such as goal conditioning, [Figure] video demonstration, natural language instruction) can all be instantiated as multimodal prompts(Fig.
Concretely, a multimodal prompt P of length l is defined as an ordered sequence of arbitrarilyinterleaved texts and images P :=
[ x 0 , x 1 , . . .
, x l ] , where each element x i ∈ { text , image } .
Task Suite.
The flexibility afforded by multimodal prompts allows us to specify and build modelsfor a huge variety of task specification formats.
Here we consider the following six task categories.1.
Simple object manipulation : simple tasks like “put <object> into <container> ”, whereeach image in the prompt corresponds to a single object; 2.
Visual goal reaching : manipulatingobjects to reach a goal configuration, e.g. , Rearrangement (Batra et al., 2020); 3.
Novel conceptgrounding : the prompt contains unfamiliar words like “dax” and “blicket”, which are explained byin-prompt images and then immediately used in an instruction.
This tests the agent’s ability to rapidlyinternalize new concepts; 4.
One-shot video imitation : watching a video demonstration and learningto reproduce the same motion trajectory for a particular object; 5.
Visual constraint satisfaction :the robot must manipulate the objects carefully and avoid violating the (safety) constraints; 6.
Visualreasoning : tasks that require reasoning skills, such as appearance matching “move all objects withsame textures as <object> into a container”, and visual memory, “put <object> in containerand then restore to their original position”.
Note that these six categories are not mutually exclusive.
For example, a task may introduce apreviously unseen verb ( Novel Concept ) by showing a video demonstration, or combine goal reachingwith visual reasoning.
More details about the task suite are discussed in Appendix, Sec.
B.
V IMA -B ENCH : B ENCHMARK FOR MULTIMODAL ROBOT LEARNING
Simulation Environment.
Existing benchmarks are generally geared towards a particular taskspecification.
To our knowledge, there is no benchmark that provides a rich suite of multimodal tasksand a comprehensive testbed for targeted probing of agent capabilities.
To this end, we introducea new benchmark suite for multimodal robot learning that we call VIMA-B ENCH .
We built ourbenchmark by extending the Ravens robot simulator (Zeng et al., 2020).
VIMA-B ENCH supportsextensible collections of objects and textures to compose multimodal prompts and procedurallygenerate a large number of tasks.
Specifically, we provide 17 meta-tasks with multimodal prompttemplates, which can be instantiated into 1000s of individual tasks.
Each meta-task belongs to one ormore of 6 task categories mentioned above.
VIMA-B ENCH can generate large quantities of imitationlearning data via scripted oracle agents.
More details are elaborated in Appendix, Sec.
A.
Observation and Actions.
The observation space of our simulator includes RGB images renderedfrom both frontal view and top-down view.
Groundtruth object segmentations and bounding boxes arealso provided for training object-centric models (Sec. 5).
We inherit the high-level action space from
Zeng et al.
(2020), which consists of primitive motor skills like “pick and place” and “wipe”.
Theseare parameterized by poses of the end effector.
Our simulator also features scripted oracle programsthat can generate expert demonstrations by using privileged simulator state information, such as theprecise location of all objects, and the groundtruth interpretation of the multimodal instruction.
Training Dataset.
We leverage the pre-programmed oracles to generate a large offline dataset ofexpert trajectories for imitation learning.
Our dataset includes 50 K trajectories per meta-task, and650 K successful trajectories in total.
We hold out a subset of object [Figure]
Evaluating Zero-Shot Generalization.
Each task in VIMA-B ENCH has a binary success criterionand does not provide partial reward signals.
During test time, we execute the agent policies in thephysics simulator for multiple episodes to compute a success rate in percentage.
The average successrate over all evaluated meta-tasks will be the final reported metric.
We design a 4-level evaluation protocol (Fig. 2) to systematically probe the generalization capabilitiesof learned agents.
Each level deviates more from the training distribution, and is thus strictlyharder than the previous one — Level 1) placement generalization : all prompts are seen verbatimduring training, but only the placement of objects on the tabletop is randomized at testing; Level 2)combinatorial generalization : all materials (adjectives) and 3D objects (nouns) are seen duringtraining, but new combinations of them appear in testing; Level 3) novel object generalization :test prompts and the simulated workspace include novel adjectives and objects; Level 4) novel taskgeneralization : new meta-tasks with novel prompt templates at test time.
VIMA: V ISUOMOTOR A TTENTION M ODEL
Our goal is to build a robot agent capable of performing any task specified by multimodalprompts.
To learn an effective multi-task robot policy, we propose VIMA, a minimalistic multitask encoder-decoder architecture with object-centric design (Fig. 3).
Concretely, we learn a robotpolicy π ( a t |P , H ) , where H :=
[ o 1 , a 1 , o 2 , a 2 , . . .
, o t ] denotes the past interaction history, ando t ∈ O , a t ∈ A are observations and actions at each interaction steps.
We encode multimodal promptsvia a frozen pre-trained langauge model and decode robot motor commands conditioned on the encoded prompts via cross-attention layers.
Unlike prior works (Florence et al., 2019; Sieb et al., 2019),
VIMA adopts an object-centric token representation that computes features from bounding box coordinates and cropped RGB patches.
Tokenization.
There are 3 formats of raw input in the prompt — text, image of a single object, andimage of a full tabletop scene ( e.g. , for Rearrangement or imitation from video frames).
For textinputs , we use pre-trained T5 tokenizer and word embedding to obtain word tokens.
For imagesof full scenes , we first extract individual objects using off-the-shelf Mask R-CNN (He et al., 2017).
Each object is represented as a bounding box and a cropped image.
We then compute object tokens byencoding them with a bounding box encoder and a ViT, respectively.
Since Mask-RCNN is imperfect,
[Figure]
the bounding boxes can be noisy and the cropped image may have irrelevant pixels.
For images ofsingle 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.
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 understandingcapability and robustness properties. To accommodate tokens from new modalities, we insert MLPsbetween the non-textual tokens and T5. To prevent catastrophic forgetting, VIMA finetunes the lasttwo layers of the language encoder with layer-wise learning rate decay (He et al., 2021) but freezesall other layers.
Our positional embedding is learnable and absolute.
Robot Controller.
A challenging aspect of designing multi-task policy is to select a suitableconditioning mechanism.
In our schema (Fig. 3), the robot controller (decoder) is conditioned onthe prompt sequence P by a series of cross-attention layers between P and the trajectory historysequence H .
We compute key K P and value V P sequences from the prompt and query Q H fromthe trajectory history, following the encoder-decoder convention in T5 (Raffel et al., 2020).
Eachcross-attention layer then generates an output sequence H ′ =
softmax (cid:16) Q H K ⊺ P √ d (cid:17) V P , where d is theembedding dimension.
Residual connections (He et al., 2015) are added to connect higher layerswith 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, as demonstrated in VideoGPT (Yan et al., 2021) as well.
VIMAdecoder consists of L alternating cross-attention and self-attention layers.
Finally, we follow commonpractice (Baker et al., 2022) to map predicted action tokens to discretized coordinates of the robotarm. See Appendix, Sec. C.2 for more details.
Training.
We follow behavioral cloning to train our models by minimizing the negativelog-likelihood of predicted actions.
Concretely, for a trajectory with T steps, we minimizemin θ (cid:80) Tt =1 − log π θ ( a t |P , 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 augmentationby randomly injecting false-positive detection outputs.
After training, we select model checkpointsfor evaluation based on the aggregated accuracy on a held-out validation set.
The evaluation involvesinteracting with the physics simulator.
We follow the best practices to train Transformer modelsusing the AdamW optimizer (Loshchilov & Hutter, 2019), learning rate warm-up, cosine annealing(Loshchilov & Hutter, 2016),
See Appendix Sec.
D comprehensive training hyperparameters.
E XPERIMENTS
In this section, we aim to answer three main questions: (1) How does VIMA compare with prior
SOTA transformer-based agents on a diverse collection of multimodal-prompted tasks?
(2) What arethe scaling properties of our approach in model capacity and data size?
(3) How do different visualtokenizers, prompt conditioning, and prompt encoding affect decision making?
6. B ASELINES
Gato (Reed et al., 2022) introduces a decoder-only model that solves tasks from multiple domainswhere tasks are specified by prompting the model with the observation and action subsequence.
Forfair comparison, we provide the same conditioning as VIMA, i.e. , our multimodal embedded prompt.
Input images are divided into patches and encoded by a ViT (Dosovitskiy et al., 2020) model toproduce observation tokens.
Flamingo (Alayrac et al., 2022) is a vision-language model that learns to generate textual completionin response to multimodal prompts.
It embeds a variable number of prompt images into a fixednumber of tokens via a Perceiver Resampler (Jaegle et al., 2021b), and conditions the languagedecoder on the encoded prompt by cross-attention.
Flamingo does not work with embodied agentsout of the box.
We adapt it to support decision-masking by replacing the output layer with robotaction heads.
Multimodal GPT agent is a GPT-based behavior cloning agent conditioned on tokenized multimodalprompts. It autoregressively decodes next actions given instructions and interaction histories.
Similarto prior works of casting RL problems as sequence modeling (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 trajectorywith prompt tokens. This baseline does not involve cross-attention.
A more detailed comparison between these methods can be found in Appendix, Sec. C.1.
6. E VALUATION R ESULTS
We compare VIMA against other SOTA methods on the four levels of generalization provided in ourbenchmark for different model and training dataset sizes.
Model scaling.
We train all methods for a spectrum of model capacities from 2M to 200M parameters,evenly spaced on the log scale.
The encoder size is kept constant (pre-trained T5-Base) for all methodsand excluded from the parameter count.
Across all levels of zero-shot generalization, we find that
VIMA strongly outperforms prior work.
Although models like Gato and Flamingo show improvedperformance with bigger model sizes, VIMA consistently achieves superior performance over allmodel sizes.
We note that this can only be achieved with both cross-attention and object tokensequence representation without any downsampling — altering any component will degrade theperformance significantly, especially in the low model capacity regime (ablations in Sec. 6.3).
Data scaling.
Next we investigate how different methods scale with varying dataset sizes.
Wecompare model performance at 0 .
1% , 1% , 10% and full imitation learning dataset provided in
VIMA-B ENCH (Fig. 4).
VIMA is extremely sample efficient and with just 1% of the data canachieve performance similar to baseline methods trained with 10 × more data for L1 and L2 levels ofgeneralization.
In fact, for L4 we find that with just 1% of training data, VIMA already outperformsprior work trained with entire dataset.
Finally, across all levels with just 10% of the data, VIMAcan outperform prior work trained with the full dataset by a significant margin.
We hypothesize thatthe data efficiency can be attributed to VIMA’s object-centric representation, which is less prone tooverfitting than learning directly from pixels in the low-data regime.
This is consistent with findingsfrom Sax et al.
(2018), which demonstrates that embodied agents conditioned on mid-level visualrepresentations tend to be significantly more sample-efficient than end-to-end control from raw pixels.
Progressive Generalization.
Finally, we compare the relative performance degradation as we test themodels on progressively challenging zero-shot evaluation levels without further finetuning (Fig. 5).
Our method exhibits a minimal performance regression, especially between L 1 → L 2 and L 1 → L 3 .
In contrast, other methods can degrade as much as 20% , particularly in more difficult generalizationscenarios.
Although all methods degrade significantly when evaluated on L 4 ( Novel Tasks ), the drop [Figure] in performance for VIMA is only half as severe as all other baselines.
This results suggest that VIMAhas developed more generalizable policy and robust representations than the competing approaches.
6. A BLATION S TUDIES
[Figure] Through extensive experiments, we ablate different designchoices in VIMA and study their impact on robot decisionmaking.
We focus on the following 4 aspects: visual tokenization, prompt encoding, prompt conditioning variants,and robustness against distractors and imperfect prompts.
Visual tokenization.
As explained in Sec. 5, VIMA processes the prompt and observation images into a variablenumber of object tokens with an off-the-shelf Mask RCNN implementation.
How important is this particularchoice of visual tokenizer?
We study 5 different variantsand empirically evaluate their 4 levels of generalizationperformance on VIMA-B ENCH .
(1) Ours (Oracle) : instead of using Mask R-CNN, we directly read out thegroundtruth bounding box from the simulator.
In otherwords, we use a perfect object detector to estimate theupper bound on the performance of this study; (2) Object
Perceiver : we apply a Perceiver module (Jaegle et al.
,2021b;a) 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 Flamingo, which converts an imageto a small, fixed number of tokens; (4) Image patches : following Gato, we divide an RGB frame intosquare patches, and extract ViT embedding tokens.
The number of patches is more than the output of Image Perceiver; (5) Single image : Decision Transformer’s tokenizer, which encodes one image intoa single token.
Fig.
6 shows the ablation results.
We highlight a few findings.
First, we note that our Mask R-CNNdetection pipeline (Appendix, Sec. A.20) incurs a minimal performance loss compared to the oraclebounding boxes, thanks to the object augmentation (Sec. 5) that boosts robustness during training.
Second, tokenizing from raw pixels (Image Perceiver, patches, or single embedding) consistentlyunderperforms our object-centric format.
We hypothesize that these tokenizers have to allocate extra [Figure] internal capacity to parse the objects from low-level pixels, which likely impedes learning.
Saxet al.
(2018) echoes our finding that using mid-level vision can greatly improve agent generalizationcompared to an end-to-end pipeline.
Third, even though Ours and Object Perceiver both use the sameobject bounding box inputs, the latter is significantly worse in decision making.
We conclude thatit is important to pass the variable sequence of objects directly to the robot controller rather thandownsampling to a fixed number of tokens.
Prompt Conditioning.
VIMA conditions the robot controller (decoder) on the encoded prompt bycross-attention.
A simple alternative is to concatenate the prompt P and interaction history H into onebig sequence, and then apply a decoder-only transformer like GPT (Radford et al., 2018) to predictactions.
In this ablation, we keep the object tokenizer constant, and only switch the conditioningmechanism to causal sequence modeling.
Note that this variant is conceptually “Gato with objecttokens”.
Fig.
7 shows the comparison of VIMA ( xattn ) and the gpt-decoder variant across 4generalization levels.
While the variant achieves comparable performance in larger models, crossattention 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 onthe prompt instruction at each interaction step.
This bears resemblance to the empirical results in Sanh et al.
; Wang et al.
(2022b), which show that well-tuned encoder-decoder architecturescan outperform GPT-3 in zero-shot generalization.
Prompt Encoding.
We vary the size of the pre-trained T5 encoder to study the effect of promptencoding.
We experiment with three T5 capacities: small (30M), base (111M), to large (368M).
For all T5 variants, we fine-tune the last two layers and freeze all other layers.
We find no significantdifference among the variants (Appendix, Sec.
E.2), thus we set base as default for all our models.
Policy Robustness.
We study the policy robustness against increased amounts of distractors andimperfect task specifications.
See Appendix, Sec. E.3 for exact setup and results.
VIMA exhibitsminimal performance degradation with increased distractors and corrupted prompts.
We attribute thisrobustness to the high-quality, pre-trained T5 language backbones.
C ONCLUSION
Similar to GPT-3, a generalist robot agent should have an intuitive and expressive interface for humanusers to convey their intent.
In this work, we introduce a novel multimodal prompting formulation thatconverts diverse robot manipulation tasks into a uniform sequence modeling problem.
We propose
VIMA, a conceptually simple transformer-based agent capable of solving tasks like visual goal,one-shot video imitation, and novel concept grounding with a single model.
VIMA exhibits superiormodel and data scaling properties, and provides a strong starting point for future work.
The current VIMA experiments are not without limitations.
We identify the following weaknesses:(1) limited action primitives (only pick-and-place and wipe for now); (2) limited simulator realism;(3) reliance on domain-finetuned Mask R-CNN to provide object tokens. However, VIMA’s algorithmdesign is general-purpose and does not make assumptions about the particular observation and actionformats. This opens the door to future works that may address many of these weaknesses withmore sophisticated environments (e.g. BEHAVIOR (Srivastava et al., 2021)), stronger vision pipeline(large-scale open-vocabulary models like ViLD (Gu et al., 2021)), and temporally-extended robotcontrollers (such as MAPLE (Nasiriany et al., 2021)). With these stronger modules, VIMA couldpotentially scale to more challenging problems.
We open-source all code to facilitate future research.
R EPRODUCIBILITY S TATEMENT
We provide comprehensive details to reproduce our work in the Appendix.
Concretely, the specifications of each meta-task in the benchmarking suite are explained in Sec.
B. Model architecturesare elaborated in Sec.
C. Hyperparameter configurations are listed in Sec.
D. Furthermore, we hostanonymized code at https://iclr3081.github.io/ for review.
A S IMULATOR D ETAILS
We build our VIMA-B ENCH 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 × 1 m. Our benchmark containsextensible sets of object geometries and textures.
Instantiated from an object-texture combination,all object instances can be rendered as RGB images appeared in multimodal prompts.
Figure A.1displays all object geometries.
Figure A.2 displays all textures.
The observation space of VIMA-B ENCH includes RGB images from both frontal and topdown views.
It also includes a one-hot vector ∈ { 0 , 1 } 2 to indicate type of the end-effector∈ { 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-B ENCH 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 ∈ SE (2) specifyingtarget poses of the end effector.
For the “pick and place” primitive, they represent the pick pose andthe place pose. Fir 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-B ENCH provides scriptedoracles to generate successful demonstrations for all tasks.
We leverage them to construct an offlineimitation dataset for behavioral cloning.
Given a prompt, these pre-programmed bots can accessprivileged information such as the correct object to pick and target location to place.
B T ASK S UITE
We develop 17 meta tasks that belong to 6 diverse categories. Thousands of individual tasks and theircorresponding multimodal prompts can be procedually generated from these meta-task templates.
Weuse 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 applythe NVISII (Morrical et al., 2020) raytracing renderer to enhance the visual quality.
We elaborateeach meta task in the following subsections.
B. S IMPLE O BJECT M ANIPULATION
This task category asks agents to follow basic instructions specified by multimodal prompts.
Task 01: Pick the specified object(s) and place it into the specified object
• 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 ob- jects 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 twodifferent 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 containerobject(s), with specified shapes and textures provided in the prompt.• Oracle Trajectory : Shown in Fig.
A.3 with its multimodal prompt.
Task 02:
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 } object.
[Figure]
[Table]
(cid:5)(cid:0)(cid:2)(cid:2)(cid:16)(cid:24)((cid:29)(cid:24)(cid:13)(cid:8)(cid:2) (cid:11)(cid:8)(cid:0)(cid:0)(cid:4)((cid:29)(cid:24)(cid:13)(cid:8)(cid:2) (cid:3)(cid:2)(cid:1)(cid:0)((cid:29)(cid:24)(cid:13)(cid:8)(cid:2) (cid:20)(cid:1)(cid:8)(cid:20)(cid:2)(cid:0)((cid:29)(cid:24)(cid:13)(cid:8)(cid:2) (cid:8)(cid:0)(cid:21)((cid:20)(cid:6)(cid:13)(cid:29)(cid:2)(cid:0)(cid:5) (cid:5)(cid:0)(cid:2)(cid:2)(cid:16)(cid:24)((cid:20)(cid:6)(cid:13)(cid:29)(cid:2)(cid:0)(cid:5) (cid:11)(cid:8)(cid:0)(cid:0)(cid:4)((cid:20)(cid:6)(cid:13)(cid:29)(cid:2)(cid:0)(cid:5) (cid:3)(cid:2)(cid:1)(cid:0)((cid:20)(cid:6)(cid:13)(cid:29)(cid:2)(cid:0)(cid:5) (cid:20)(cid:1)(cid:8)(cid:20)(cid:2)(cid:0)((cid:20)(cid:6)(cid:13)(cid:29)(cid:2)(cid:0)(cid:5)
[Table]
[Figure]
Put the into the .
• 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% 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.
Put the green and blue stripe object in into the yellow paisley object.
Task 03: Rotate objects clockwise by certain degrees along z -axis.
Only rotationally asymmetricobjects 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 the dragged object needs to be rotated.
A target angle is sampled from 30 ◦ , 60 ◦ , 90 ◦ , 120 ◦ , and150 ◦ .
• 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. V ISUAL G OAL R EACHING
This task category requires agents to manipulate objects in the workspace to reach goal statesrepresented as images shown in prompts.
Rotate the 120 degrees.
Task 04: 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.
• Prompt : Rearrange to this { scene } .
• 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 defined distractor conflict rate, theposition 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. .
[Figure]
Rearrange to this . Task 05: 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 maxsteps restore all target objects to their initial configurations.
• Oracle Trajectory : Shown in Fig.
A.7 with its multimodal prompt.
[Figure]
B. NOVEL CONCEPT GROUNDING
[Figure] 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 withinterleaved texts, images, and even video frames; b) quickly internalize new concepts that are differentacross task instances, which even tests the ability to meta learn; and c) do complicated reasoning suchas comparing between “taller” vs “less taller” vs “shorter” and then ground this reasoning into therobot action space.
Prompts consist of two parts: a definition part followed by an instruction part.
In the definitionpart, novel conceptions are defined by multimodal illustrations with multiple support examples. Inthe instruction part, agents are asked to achieve the goal by properly applying concepts from thedefinition part.
The assignment of unique nonsense words is varied and independent for each taskinstance such that tasks can only be solved if the agent applies the reasoning correctly.
This ability isalso referred to as fast-mapping (Heibeck & Markman, 1987).
Task 06: Ground comparative adjectives by comparing the size or the textural saturation of objectsand 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 placeholderfor 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 thereal 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 totask 01 , where the agent is required to pick the specified dragged object(s) with the noveladjective attribute and then place it into the specified container object.
To avoid revealing thecorrect object to manipulate, we use a neutral texture for objects appeared in the instructionpart.
• Success Criteria : All target objects with the specified adjective attribute are within thebounds of the specified container object.• Oracle Trajectory : Shown in Fig.
A.8 with its multimodal prompt.
Task 07: Orthogonal to task 06 by requiring to learn mappings of novel nouns.
is kobar than . is kobar than .
is kobar than .
Put the kobar into the .
• Prompt : This is a { novel name } 1 { object } 1 .
This is a { novel name { object } 2 .
Put { novel name } 1 into a { novel name } 2 .
[Figure] • Description : Novel noun words are defined with the text placeholders { novel name } and { novel name } 2 , following their image placeholders { object } 1 and { object } 2 , for the target object and container object, respectively.
Novel nouns are sampled from { dax, [Figure] blicket, wug, zup } .
In the instruction part, objects are expressed as novel nouns [Figure] defined in the previous definition part.
Distractors are defined the same as task
[Figure] • Success Criteria : All target object(s) are within the bounds of the container object(s).
[Figure]
• Oracle Trajectory : Shown in Fig.
A.8 with its multimodal prompt.
[Figure]
Task 08:
Combination of tasks 06 and 07 . [Figure]
• Prompt : This is a { novel name } 1 { object } 1 .
This is a { novel name }
[Figure] { object
{ demo object } 1 is { adj } than { demo object } 2 . Put the [Figure]
{ adv } { novel adj } { novel name } 1 into the { novel name } 2 .
[Figure] • 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.
Task 09:
A novel verb ”twist” is defined as rotating a specific angle conveyed by several examples.
This task is similar to task 03 , but it requires the agent to infer what is the exact angle to rotate fromthe prompt and to ground novel verbs that are semantically
similar but different in exact definitions. [Table] • Prompt : "Twist" is defined as rotating object a specific angle. [Table]
For examples: From { before twist } i to { after twist } i . Now twist [Table] all { texture } objects.
[Table] • Description : Both { before twist } i and { after twist } i are scene placeholders where { before twist } i shows a randomly sampled object before ”twist” and [Table] { after twist } i shows the same object pose after ”twist”.
All examples illustrate the
[Table]
This is a wug .
This is a zup .
is blicker than . is blicker than .
is blicker than .
Put the blicker zup into the wug.
same sampled angle of the rotation.
In the workspace, the target objects have the texturespecified by { texture } and randomly sampled shapes.
• Success Criteria :
Same as the task
• Oracle Trajectory : Shown in Fig.
A with its
multimodal prompt.(cid:2)(cid:23)(cid:24)(cid:9)(cid:19)(cid:20)(cid:2)(cid:18)(cid:9)(cid:19)(cid:18)(cid:11)(cid:15)(cid:13)(cid:9)(cid:12)(cid:15)(cid:11)(cid:18)(cid:8)(cid:19)(cid:18)(cid:7)(cid:22)(cid:20)(cid:8)(cid:20)(cid:9)(cid:12)(cid:4)(cid:18)(cid:22)(cid:25)(cid:14)(cid:15)(cid:21)(cid:20)(cid:18)(cid:8)(cid:18)(cid:19)(cid:5)(cid:15)(cid:21)(cid:9)(cid:13)(cid:9)(cid:21)(cid:18)(cid:8)(cid:12)(cid:4)(cid:3)(cid:15)(cid:10)(cid:18)(cid:17)(cid:22)(cid:7)(cid:18)(cid:15)(cid:1)(cid:8)(cid:6)(cid:5)(cid:3)(cid:15)(cid:19)(cid:0)(cid:16) (cid:17)(cid:7)(cid:22)(cid:6)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:20)(cid:22)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:18)(cid:10) [Table]
[Figure] 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 theproblem by giving one video demonstration (represented as consecutive frames in prompts), thentest the learned imitator’s ability to produce target trajectories. This setup is challenging becausea) only one demonstration is available to the agent; b) the model needs to understand video framesinterleaved with textual instructions; and c) missing correspondences between demonstrations andtarget trajectories since demonstrations only show partial key frames.
Task
Follow motions for specific objects.
• Prompt : Follow this motion for { object } : { frame } 1 ... { frame } i ...{ frame } n .•
Description : Image placeholder { object } is the target object to be manipulated and {{ frame } i } is set of workspace-like scene placeholders to represent a video trajectory,where n 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 .
[Figure]
• Success Criteria : In each step, the pose of the target object matches the pose in thecorresponding video frame.
Incorrect manipulation sequences are considered as failures. [Figure]
• Oracle Trajectory : Shown in Fig.
A.12 with its multimodal prompt.
[Figure]
Task 11: Stack objects with the order illustrated in the prompt video.
[Figure]
• Prompt : Stack objects in this order { frame } 1 ... { frame } i ... { frame } n .
[Figure] • Description : There are multiple objects with the same shape but different textures spawnedin the workspace without any stacking initially.
Distractor objects with different shapes arespawned in the workspace but not in the prompt video.
At each step of the prompt video,one of the top objects is stacked over another object or put at an empty position.• Success Criteria : Similar as task 10.
[Figure]
• Oracle Trajectory : Shown in Fig.
A.13 with its multimodal prompt.
[Figure]
B. V ISUAL C ONSTRAINT S ATISFACTION
[Figure] This task category requires agents to wipe a specific number of objects in the workspace to a goal [Figure] region while also satisfy the given visual constraint
[Figure]
Task 12:
Sweep the designated number of objects into a specified region without exceeding theboundary.
• Prompt : Sweep { quantifier } { object } into { bounds } withoutexceeding { constraint } .
• Description : { object } is the image placeholder of the target object to be swept spawnedwith a random amount in the workspace.
Distractors have the same amount, same shape, butdifferent 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 , andall .
• Success Criteria :
The exact number of target objects to be swept are all inside the specifiedregion.
Failure reasons include 1) any distractor being wiped into the region, 2) target objectexceeding the constraint, or 3) incorrect number of target objects being swept into the goal [Figure] region. [Figure]
• Oracle Trajectory : Shown in Fig.
A.14 with its multimodal prompt
[Figure]
Task 13:
Sweep the designated number of objects into a specified region without touching the constraint. • Prompt : Sweep { quantifier } { object } into { bounds } withouttouching { constraint } .
• Description : Similar as task 12 but requiring different way to satisfy the constraint.
Theagent has to learn to avoid ”touching” the constraint line in this case.
• Success Criteria : Similar as task 12 except that the constraint is to not touch the red line.
[Figure]
• Oracle Trajectory : Shown in Fig.
A.15 with its multimodal prompt.
[Figure]
B. V ISUAL R EASONING
This task category requires agents to make decisions by reasoning over or memorizing informationconveyed through multimodal prompts.
Task 14:
By reasoning the ”same texture”, the agent is required to pick all objects in the workspacewith 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 } intoit.
• Description : { object } is the sampled goal container object.
In the workspace, there areobjects with the same texture as the container but potentially different shapes.
Distractorswith different textures are spawned.
• Success Criteria : All objects with the same texture as the goal container are within thebounds of the container.
• Oracle Trajectory : Shown in Fig.
A.16 with its multimodal prompt.
[Figure] as into it. Task 15: By reasoning the ”same shape”, the agent is required to pick all objects in the workspacewith the same top-down shape as the goal container specified in the prompt and place them into it.
For example, blocks and boxes have the same rectangular shape.
• 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 are with 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.
[Figure]
Task
: 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 and { object } 2 are the targetobject 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 } .
[Table]
• Success Criteria :
The target object and the correct neighboring object are inside thecontainer.
[Table]
• Oracle Trajectory : Shown in Fig.
A.18 with its multimodal prompt.
[Table]
[Figure]
Task 17: Pick and place the target object specified in the prompt into different containers in order [Figure] then restore to the initial container. [Figure]
• Prompt : Put { object } 1 into { object }
Finally restore it into [Figure] its original container.
[Figure] • Description : The object in the image placeholder { object } 1 is the target ob- [Figure] ject to be manipulated across the task.
There are more than one target containers (e.g. [Figure]
Put { object } 1 into { object } 2 then { object } 3 .Finally
restore it [Figure] into its original container.
for two target base objects to be placed in order).
[Figure]
The rest of spawned containers naturally becomes distractors. [Figure]
• Success Criteria :
The target object are first put into multiple containers following the
[Figure] specific order.
Finally it should be restored into its original container. [Figure]
• Oracle Trajectory : Shown in Fig.
A.19 with its multimodal prompt.
[Figure]
C M ODEL A RCHITECTURE
In this section, we provide comprehensive details about VIMA model architecture as well as otheradapted baseline methods.
We implement all models in PyTorch (Paszke et al., 2019) and adapt
Transformer-related implementation from Wolf et al.
C. S UMMARY OF D IFFERENT M ETHODS
We summarizes differences between VIMA and other baseline methods in Table 1.
In the column“Prompt Conditioning”, an alternative of cross-attention is to first concatenate prompt and interactioninto a big sequence, then repetitively apply transformer decoders to predict actions.
It is referredto as “direct modeling”.
The relative computation cost is quadratically proportional to number ofobservation tokens.
[Table] [Table] C. VIMA A RCHITECTURE C.2. M ULTIMODAL P ROMPT T OKENIZATIONS
As introduced in Section
[Table] there are 3 types of input formats in multimodal prompts, namely (1) textinputs , (2) images of full scenes [Table] and (3) images of single objects .
For text inputs , we follow the standard pipeline in NLP to first tokenize raw languages to discreteindices through pre-trained t5-base tokenizer.
We then obtain corresponding word tokens from theembedding look-up of the pre-trained t5-base model.
For images of full scenes , we first parse thescene through a fine-tuned mask R-CNN detection model (He et al., 2017; Wu et al., 2019) to extractindividual objects.
Each object representation contains a bounding box and a cropped image.
Thebounding box is in the format of [ x center , y center , height , width ] .
We normalize it to be within [0 , 1] bydividing each dimension with corresponding upper-bound value.
We then pass it through a boundingbox encoder MLP and obtain a feature vector.
To process the cropped image, we first pad non-squareimage to a square by padding along the shorter dimension.
We then resize it to a pre-configured sizeand pass it through a ViT (trained from scratch) to obtain the image feature.
Finally, an object tokenis obtained by concatenating the bounding box feature and the image feature and mapping to theembedding dimension.
For images of single objects , we obtain tokens in the same way except with adummy bounding box.
Detailed model hyperparameters about tokenizations are listed in Table 2.
After obtaining a sequence of prompt tokens, we follow Tsimpoukelli et al.
(2021) to pass it througha pre-trained t5-base encoder to obtain encoded prompt.
Note that we add adapter MLP between object tokens and the T5 encoder.
We adopt learned absolute positional embedding.
Modelhyperparameters are listed in Table 2 as well. [Table] [Table]
C.2. O BSERVATION E NCODING
Since all RGB observations are images of full scenes, we follow the same procedure discussedabove to obtain flattened object tokens.
Because we provide RGBs from two views (frontal andtop-down), we order object tokens by following the order of [ frontal, top-down ] .
We one-hot encodethe state of the end effector.
We then concatenate object tokens with the end-effector state andtransform to observation tokens.
We adopt learned absolute positional embedding.
Detailed modelhyperparameters about observation encoding is provided in Table 3.
C.2. A CTION E NCODING
Since our model is conditioned on observation-action interleaved history, we also tokenize pastactions.
We follow common practice in Chen et al.
; Zheng et al.
(2022) to encode past actionswith a two-layer MLP.
It has a hidden dimension of 256.
We then map outputs to token dimensionand obtain action tokens.
C.2. S EQUENCE M ODELING
The robot controller in VIMA is a causal decoder that autoregressively predicts actions.
To conditionthe 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 thekey-value sequence into cross-attention blocks.
The output prompt-aware trajectory tokens then gothrough causal self-attention blocks.
We alternate cross-attention and self-attention L times.
Thisprocedure is technically described in Pseudocode 1. [Table] [Table]
[Table] , # the [ [Table] [Table]
action tokens traj [Table] pos [Table]
embd, # learned positional embedding for trajectory
[Table]
, # learned positional embedding for prompt ): # interleave obs and action tokens traj
[Table]
tokens = interleave(obs [Table] , act_tokens) # [2T-1, d] # add positional embedding to trajectory tokens x = traj [Table]
+ [Table] # add positional embedding to prompt tokens [Table]
= [Table] tokens + prompt [Table]
pos_embd # apply xattn and causal self-attn for i in range(num_layers): # [Table] attention [Table]
[Table]
, kv [Table] [Table] kv
[Table]
[Table]
[Table]
last
[Table]
is
[Table]
predicted
[Table]
[Table] return [Table] Pseudocode 1: Cross-attention operation that conditions the trajectory history on prompt.
We repetitively alternate cross-attention and self-attention to model the trajectory given a specific task.
C.2. A CTION D ECODING
After obtaining the predicted action token, we map it to the action space A and obtain the predictedaction.
This is achieved though a group of action heads.
Since the action space consists of two
SE (2) poses, for each pose we use six independent heads to decode discrete actions (two for xycoordinate and four for rotation represented in quaternion).
These discrete actions are then dediscretized and mapped to continuous actions through affine transformation.
The two poses aremodeled independently.
Early ablations show that this independent modeling is equally good asalternatives techniques like autoregressive decoding (Vinyals et al., 2019; OpenAI et al., 2019).
Detailed model hyperparameters are listed in Table 4.
C. B ASELINES A RCHITECTURES
In this section, we elaborate model architectures for baseline methods.
Some components such as theaction decoder are same across all baseline methods and ours.
Therefore, we only discuss uniquemodel components.
C.3. G ATO
Gato (Reed et al., 2022) introduces a decoder-only model that solves tasks from multiple domainsincluding 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 tasksinvolving 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. Similar to ourmethod, Gato also predicts actions in an autoregressive manner. Gato and our method share the sametraining philosophy to only optimize the causal behavior cloning objective.
However, unlike ourmethod that adopts an object-centric representation to treat individual objects as observation tokens,
Gato divides input images into patches and encodes them by a ViT (Dosovitskiy et al., 2020) modelto produce observation tokens.
Furthermore, Gato relies on causal self-attention to model entiretrajectory sequences starting with prompt tokens.
Hyperparameters of Gato’s ViT is listed in Table 5.
The transformer-decoder style sequence modeling is technically illustrated in Pseudocode 2.
[Equation] #the[L, d] prompt tokens (L=prompt length) sep_token, # the [1, d] learned token to separate prompt and trajectory history obs_tokens, # the [T, d] obs tokens (T=time step) act_tokens, # the [T-1, d] action tokens pos_embd, # learned positional embedding ): # interleave obs and action tokens
traj_tokens = interleave(obs_tokens, act_tokens) #
[2T-1, d] # assemble input tokens x = [Equation] [Equation] _tokens, sep_token, traj_tokens]) x
= x + pos_embd # apply GPT layers with causal mask for i in range(num_layers): [Equation] # the last token is the predicted action token predicted_ [Equation]
[Equation] return predicted_act_token Pseudocode 2: Plain sequence modeling that temporally concatenates prompt and trajectory history and repetitively perform causal self-attention operation. C.
F LAMINGO Flamingo (Alayrac et al., 2022) is a vision-language model that learns to generate textual completionin response to multimodal prompts.
It embeds a variable number of prompt images into a fixed numberof tokens via the Perceiver Resampler module (Jaegle et al., 2021b), and conditions the languagedecoder on encoded prompts by cross-attention.
Flamingo does not work with embodied agents outof the box.
We adapt it by replacing the output layer with robot action heads (hyperparameters listedin Table 4) and using tokenized rollout histories as inputs.
We train it end-to-end with causal behaviorcloning loss. The modified Flamingo agent differs from ours since it processes image observationsinto a fixed number of visual tokens through a learned Perceiver Resampler.
Model hyperparametersfor our reimplementation of the Perceiver Resampler is listed in Table 6.
C.3. M ULTIMODAL GPT A GENT
Multimodal GPT agent (Brown et al., 2020) is a behavior cloning agent conditioned on tokenized multimodal prompts with the GPT architecture. It autoregressively decodes next actionsgiven multimodal prompts and interaction histories.
We optimize this method end-to-end withcausal 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 sequencesprepended 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 withhyperparameters listed in Table 5.
C. M ASK R-CNN D ETECTION M ODEL
Finally, we elaborate the mask R-CNN model (He et al., 2017) for scene parsing and object extraction.
We fine-tuned 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.
A visualization of its output isprovided in Figure A.20.
We do not use the predicted object names in our models.
D VIMA T RAINING D ETAILS
[Equation] We follow the best practices 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.
[Table]
To make trained models robust to detection inaccuracies and failures, we apply object augmentationby randomly injecting false-positive detection outputs.
Concretely, for observation at each time step,we sample number of augmented objects i.i.d.
n augmented objects ∼ Cat ( K, p ) , where Cat ( · ) denotes amulti-categorical distribution with K supports parameterized by p .
For each augmented object, wethen randomly sample a bounding box and corresponding cropped image to add to object tokens.
Inour experiments, we set p = { 0 : 0 .
95 , 1 : 0 . 05 } with K = 2 .
D. V ARY M ODEL C APACITY
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 Flamingo)in Table 8, and configurations for methods only with causal self-attention (i.e., Gato and DT) in Table 9. [Table] [Table]
E M ORE E XPERIMENT R ESULTSE. B REAKDOWN R ESULTS
We show breakdown results for Figure 4 in Tables 10, 11, 12, and 13, respectively.
E. V ARY
[Equation]
E NCODER S IZES
We vary the size of the pre-trained T5 encoder (Raffel et al., 2020) to study the effect of promptencoding.
We experiment with three T5 model capacities: t5-small (30M), t5-base (111M), tot5-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 14, we findno significant difference among the variants.
Thus we set the standard t5-base as default for allour models.
E. P OLICY R OBUSTNESS
Increased Amounts of Distractors.
We study the policy robustness against increased amounts ofdistractors in scenes. For all tasks being evaluated, we add one more distractor object. We ran ourlargest VIMA model with 200M parameters. The result is presented in Table 15.
It turns out that the performance of VIMA degrades minimally with more distractors than the trainingdistribution.
This indicates that our agent has learned a reasonably robust policy against objects thatare irrelevant to the task.
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
Imperfect Prompts .
We then study the policy robustness against imperfect prompts
[Table]
including incomplete prompts (randomly masking out words with <UNK> token) and corrupted prompts (randomlyswapping words, which could have changed the task meaning altogether).
We ran our largest VIMAmodel with 200M parameters, results are shown in Table
[Table]
[Table]
[Table]
Our well-trained model exhibits minimal performance decrease when evaluated on masked promptsand minor decrease on corrupted prompts.
We attribute this robustness to the high-quality, pre-trained
T5 language backbones.
[Table]
F E XTENDED R ELATED W ORK
In this section
[Table]
we provide an extended review of related work as complementary to Section 2.
Multi-task Learning by Sequence Modeling.
In NLP domain, the Natural Language Decathlon (McCann et al., 2018) adopts a consistent question-answering format for a suite of {+
+} NLP tasks.
In computer vision, Mask R-CNN (He et al., 2017), UberNet (Kokkinos, 2016), and 12in-1 (Lu et al., 2020) leverage a single backbone model with multiple independent heads for differenttasks.
UVim (Kolesnikov et al., 2022) is another unified approach for vision that uses a languagemodel to generate the guiding code for a second model to predict raw vision outputs.
In multimodallearning, 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, though most of which are not equippedwith decision-making facilities.
Perceivers (Jaegle et al., 2021b;a) propose an efficient architectureto handle general-purpose inputs and outputs.
BEiT-3 (Wang et al., 2022c) performs masked datamodeling 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.
Foundation Models for Embodied Agents.
Embodied agent research (Duan et al., 2022; Batraet al., 2020; Ravichandar et al., 2020; Collins et al., 2021) is adopting the large-scale pre-trainingparadigm, powered by a collection of learning environments (Abramson et al., 2020; Shridharet al., 2020; Savva et al., 2019; Puig et al., 2018; Team et al., 2021; Toyama et al., 2021; Shi et al.,
From the aspect of pre-training for better representations , 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.
From the perspective of leveraging transformeras 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-attentionmechanism with auxiliary losses such as inverse dynamics loss (Dasari & Gupta, 2020) andcontrastive loss (Zhao et al., 2022) to learn robot controllers. Our work differs from them mainlyin three aspects: a) our method employs a transformer backbone to autoregressively predictactions; b) we utilize pre-trained language models (Raffel et al., 2020) and best practices from
Tsimpoukelli et al.
(2021) to learn policies conditioned on prompts with interleaved texts, images,and even videos; and c) while these works mainly focus on solving the single task of one-shot videoimitation with highly customized objectives, conceptually simple but effective, our model is learnedin a multi-task way with only the behavior cloning objective to solve a strict superset of tasks.
Robot Manipulation and Benchmarks.
There are many prior works that are not mentioned in themain paper that study different robotic manipulation tasks, such as constraint satisfaction (Bharadhwajet al., 2021), one-shot imitation (Paine et al., 2018; Huang et al., 2019; Aceituno et al., 2021;
Zhao et al., 2022), and rearrangement (Liu et al., 2021; Ehsani et al., 2021; Gan et al., 2021;
Stengel-Eskin et al., 2022).
Multiple simulation benchmarks are introduced to study the abovetasks: 1) Indoor simulation environments : Habitat (Savva et al., 2019; Szot et al ., 2021) isequipped with a high-performance 3D simulator for fast rendering and proposes a suite of commontasks for assistive robots.
AI2-THOR (Ehsani et al., 2021; Deitke et al., 2022) is a frameworkthat supports visual object manipulation and procedural generation of environments.
2) Tabletopenvironments : Meta-World (Yu et al., 2019), RLBench (James et al., 2019), and SURREAL (Fanet al., 2018; 2019) are widely used simulator benchmarks studying robotics manipulation withtabletop settings.
CausalWorld
(Ahmed et al., 2021) is a benchmark for causal structure and transferlearning in manipulation, requiring long-horizon planning and precise low-level motor control.
MOSAIC (Zhao et al., 2022) features a challenging benchmark built on top of Zhu et al.
(2020)to evaluate one-shot imitation learning.
It proposes a three-step test setting to evaluate therepresentational and generalization capability.
Compared to it, ours supports a wide spectrum ofmanipulation tasks, including one-shot imitation learning.
All these aforementioned simulators andbenchmarks do not natively support task specification and prompting with multiple modalities.