VIMA: General Robot Manipulation with Multimodal Prompts

Contributions

• The work contributes a multimodal prompting formulation, a scalable benchmark, and an agent capable of multi-task and zero-shot generalization.Source: §1 · p. 1

Motivation

• Multimodal prompts interleave language and images or video frames to express diverse robot manipulation tasks.Examples include visual goal reaching and few-shot video imitation.Source: §1 · p. 1
• Different manipulation tasks previously relied on distinct policy architectures, objectives, data pipelines, and training procedures, producing siloed systems.Source: §1 · pp. 1–2

Contributions

• VIMA-BENCH provides 17 multimodal tasks, procedurally instantiated into thousands of instances, with four progressively harder generalization levels.Source: §1 · p. 2
• VIMA encodes interleaved textual and visual tokens and autoregressively decodes robot actions using a transformer architecture.Its controller uses cross-attention to condition actions on prompts.Source: §1 · p. 2

Prompt formulation

• A multimodal prompt is an ordered sequence of arbitrarily interleaved textual and visual elements.Each element is text or an image.Source: §2 · p. 3
• This formulation unifies goal conditioning, video demonstration, and natural language instruction within one sequence modeling problem.Source: §2 · p. 3

Task suite

• The task suite includes simple object manipulation, visual goal reaching, novel concept grounding, and one-shot video imitation.Source: §2 · p. 3
• Additional categories cover visual constraint satisfaction and visual reasoning, including appearance matching and visual memory.Source: §2 · p. 3
• Task categories are not mutually exclusive, so a task may combine novel concept grounding with video demonstration or visual reasoning.Source: §2 · p. 3

Benchmark

• VIMA-BENCH provides 17 multimodal task templates that generate thousands of instances from varied objects and textures.Scripted oracle agents generate large quantities of imitation-learning data.Source: §3 · p. 3

Simulation environment

• The simulator supplies frontal and top-down RGB observations, object annotations, and primitive motor skills such as pick-and-place and wipe.Source: §3 · p. 3

Dataset

• The dataset contains 50K trajectories per task and 650K successful trajectories in total, with held-out objects, textures, and tasks for evaluation.Source: §3 · p. 4

Evaluation

• Evaluation reports average binary task success across multiple simulator episodes without partial rewards.Source: §3 · p. 4
• The four evaluation levels progress from randomized placement to new object combinations, novel objects, and novel task templates.Each level deviates further from training and is strictly harder than the previous one.Source: §3 · p. 4

Architecture

• VIMA learns a robot policy π(a_t|P, H) from a multimodal prompt P and past interaction history H.Source: §4 · p. 4
• A frozen pre-trained language model encodes prompt tokens, while the decoder predicts waypoint commands through cross-attention.Source: §4 · pp. 4–5

Object-centric representation

• VIMA represents visual inputs with object tokens derived from bounding boxes and cropped RGB patches rather than raw images.Source: §4 · pp. 4–5

Robot controller

• The controller alternates prompt-to-history cross-attention with causal self-attention before mapping action tokens to discretized robot-arm poses.Source: §4 · p. 5

Training

• Training uses behavioral cloning on offline expert trajectories by minimizing the negative log-likelihood of predicted actions.Object augmentation injects false-positive detections to improve robustness to detection errors.Source: §4 · p. 5

Evaluation Results

• VIMA’s object tokens combined with cross-attention conditioning are the most effective recipe among the evaluated model designs.The comparison covers four generalization levels and different model and training-dataset sizes.Source: §5.2 · p. 6

Model Scaling

• VIMA consistently outperforms alternative architectures across model capacities ranging from 2M to 200M parameters.The encoder size is held constant and excluded from the parameter count.Source: §5.2 · p. 6

Data Scaling

• With 1% of the training data, VIMA matches baselines using 10× more data on L1 and L2 and surpasses full-dataset variants on L4.With 10% of the data, VIMA outperforms other architectures trained on the full dataset across all generalization levels.Source: §5.2 · pp. 6–7

Progressive Generalization

• On progressively harder zero-shot evaluations, VIMA has smaller performance regression than baselines, whose degradation reaches 20%.On L4 novel tasks, VIMA’s performance drop is only half as severe as the other baselines.Source: Figure 5 · p. 6, §5.2 · p. 7

Prompt Conditioning and Encoding

• Cross-attention is especially helpful in the low-parameter regime and on harder generalization tasks, while T5 encoder size produces no significant performance difference.The prompt-conditioning comparison contrasts cross-attention with a vanilla decoder, and the encoder study tests 30M, 111M, and 368M T5 capacities.Source: Figure 7 · p. 8, §5.3 · p. 8

Visual Tokenization

• Object-centric tokenization consistently outperforms raw-pixel tokenizers and fixed-length Object Perceiver representations.The Mask R-CNN pipeline incurs minimal loss relative to oracle bounding boxes, while directly passing the variable-length object sequence improves decision making over downsampling.Source: Figure 6 · p. 7, §5.3 · p. 8

• Transformers unify multi-task learning across language, vision, and embodied-agent domains through shared sequence-modeling or backbone approaches.Source: §6 · pp. 8–9
• Multimodal models such as Perceiver, Flamingo, Frozen, and Gato process interleaved modalities or support broad task coverage, but VIMA focuses on robot prompting.Source: §6 · p. 9
• Robot-manipulation research covers instruction following, one-shot imitation, rearrangement, constraint satisfaction, and reasoning across multiple simulation benchmarks.Source: §6 · p. 9
• VIMA-BENCH is presented as the first robot-learning benchmark supporting multimodal-prompted tasks and standardized evaluation of generalization capabilities.Source: §6 · p. 9

• VIMA converts diverse robot-manipulation tasks into a uniform sequence-modeling problem and uses one transformer agent across multiple task types.Source: §7 · p. 9
• VIMA-BENCH is built on Ravens with procedurally generated tabletop instances formed from extensible 3D objects and textures.Source: A. Simulator Details · p. 20, B. Task Suite · p. 22

Simple object manipulation

• Simple manipulation tasks specify objects, containers, textures, and rotations through multimodal prompts with corresponding success criteria and oracle trajectories.Source: §B.1 · pp. 22–23

Visual goal reaching

• Visual-goal tasks require rearranging objects to match prompted scene configurations, sometimes moving distractors and restoring the initial setup.Source: §B.2 · p. 24

Novel concept grounding

• Novel concept grounding includes learning the meaning of a dummy adjective from demonstrations and inferring the angle associated with the novel verb “twist”.Source: §B.3 · pp. 25, 27

• Visual reasoning tasks require decisions based on information conveyed through multimodal prompts, including reasoning and memorization.Source: §B.6 · p. 30

Appearance matching

• Tasks 14 and 15 require grouping objects by the container’s texture or top-down shape profile, respectively.Source: §B.6 · p. 30

Success criteria

• The benchmark defines success using containment, correct neighbor selection, ordered placement, and final restoration criteria.Source: §B.6 · pp. 30–32

Spatial memory

• Task 16 requires placing a target object and then selecting its former neighbor using a specified cardinal direction.Source: §B.6 · p. 31

Ordered manipulation

• Task 17 requires moving a target through multiple containers in order and restoring it to its original container.Source: §B.6 · p. 31

Baseline variants

• Direct modeling concatenates prompt and trajectory tokens, whereas VIMA-Flamingo uses a Perceiver Resampler and VIMA-GPT uses decoder-only causal modeling.Source: §C.1 · p. 33, §C.3.1 · p. 36, §C.3.2 · pp. 36–37

Prompt tokenization

• Multimodal prompts contain text, full-scene images, and single-object images that are converted into token sequences.Source: §C.2.1 · p. 33
• Scene objects are represented using normalized bounding-box features and cropped-image features, then mapped into object tokens.Source: §C.2.1 · p. 33

Observation and action encoding

• Observation tokens combine flattened object tokens from frontal and top-down RGB views with the end-effector state, while past actions become action tokens.Source: §C.2.2 · pp. 33–34, §C.2.3 · p. 34

VIMA

• VIMA uses cross-attention from trajectory-history tokens to prompt tokens, alternating with causal self-attention before predicting the next action token.Source: §C.2.4 · pp. 34–35

Action decoding

• Predicted action tokens are decoded through independent action heads into discrete pose components and then mapped to continuous actions.Source: §C.2.5 · p. 35

• Training uses AdamW, learning-rate warm-up, cosine annealing, and GEGLU activations across Transformer methods.Source: D. VIMA Training Details · p. 38, Table 7 · p. 38
• Object augmentation injects randomly sampled false-positive detections, bounding boxes, and cropped images into observation object tokens.The number of augmented objects is sampled independently from a categorical distribution at each time step.Source: D. VIMA Training Details · p. 38
• Seven model configurations vary embedding dimension and layer count across 2M–200M parameters.Cross-attention methods and causal-self-attention methods use separate configuration tables.Source: §D.1 · p. 38, Table 8 · p. 38, Table 9 · p. 38
• Experiments run on nodes with 8 NVIDIA V100 GPUs, using distributed data parallelism; the largest experiment takes approximately one day.Source: §E.1 · p. 39

• 2.9× improvement is obtained on hardest L4 generalization with a 20M model relative to VIMA-Gato.The comparison uses the relative performance gap against VIMA-Gato, retaining the first decimal digit.Source: §E.2.1 · p. 39, Table 10 · p. 39
• Baseline variants that learn from raw pixels are initialized with MVP-pretrained, MAE-fine-tuned ViTs using the same in-domain detector data.This setup is intended to equalize pretraining information across methods.Source: §E.3.1 · p. 39
• 2.7× improvement comes from VIMA trained with 10% data achieving 46% average success rate versus VIMA-Gato’s 12.2% with full data on L4.Source: §E.3.2 · p. 39, Table 11 · p. 40
• With reasonable in-domain data—1%, 10%, or 100%—the recommended recipe consistently outperforms other variants.The same pattern is reported when baseline ViTs are trained from scratch.Source: §E.3.3 · p. 40, Figure A.21 · p. 40, Table 12 · p. 41
• MVP-pretrained ViTs help most at 10% data, while their advantage diminishes or becomes detrimental at full, 1%, and 0.1% data.Source: §E.3.3 · p. 40
• Prompt-encoder size produces no significant performance difference across T5-small, T5-base, and T5-large, so T5-base is used by default.The decision-making component is fixed at 200M parameters.Source: §E.4 · pp. 40–41, Table 13 · p. 41

• Adding one distractor object causes minimal performance degradation for the 200M-parameter VIMA model.The result indicates robustness to irrelevant objects beyond the training distribution.Source: §E.5 · p. 41, Table 14 · p. 41
• Masked prompts produce minimal performance decrease, while corrupted prompts cause a minor decrease.The authors attribute this robustness to the pretrained T5 language backbone.Source: §E.5 · p. 42, Table 15 · p. 42

• VIMA inherits standalone object-detector errors, especially under occlusion or out-of-distribution object forms.The detector can be replaced with stronger or open-vocabulary models such as OWL-ViT.Source: G. Limitations and Further Discussions · p. 44
• VIMA-BENCH prioritizes multitask ability, generalization, and multimodal understanding over physical realism and task complexity.Future work may combine the formulation with more physically realistic simulators.Source: G. Limitations and Further Discussions · p. 44
• VIMA currently uses only high-level pick-and-place and wipe actions, although its algorithm requires minimal modification for lower-level action spaces.Source: G. Limitations and Further Discussions · p. 44