Source-linked AI summary
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
TL;DR
VIMA addresses fragmented robot task specification by expressing diverse manipulation tasks through multimodal prompts. It introduces a benchmark and transformer agent, achieving stronger zero-shot generalization and data efficiency than alternative designs.
Problem
Robot manipulation tasks use varied specifications and specialized models, limiting a uniform interface for diverse capabilities.
Method
VIMA formulates diverse manipulation tasks as uniform sequence modeling with multimodal prompts, a procedurally generated benchmark, and a transformer-based agent.
Results
VIMA consistently outperforms alternative designs across four zero-shot generalization levels and model capacities, with reported gains up to 2.9× task success rate.
Takeaways & Limitations
Multimodal prompting provides a common interface for multi-task robot learning and supports model scalability, zero-shot generalization, and data efficiency.
Takeaways & Limitations
VIMA-BENCH uses a simulator with limited realism and task complexity, which are not its primary focus.
Abstract
from arXiv · showhide
Prompt-based learning has emerged as a successful paradigm in natural language processing, where a single general-purpose language model can be instructed to perform any task specified by input prompts. Yet task specification in robotics comes in various forms, such as imitating one-shot demonstrations, following language instructions, and reaching visual goals. They are often considered different tasks and tackled by specialized models. We show that a wide spectrum of robot manipulation tasks can be expressed with multimodal prompts, interleaving textual and visual tokens. Accordingly, we develop a new simulation benchmark that consists of thousands of procedurally-generated tabletop tasks with multimodal prompts, 600K+ expert trajectories for imitation learning, and a four-level evaluation protocol for systematic generalization. We design a transformer-based robot agent, VIMA, that processes these prompts and outputs motor actions autoregressively. VIMA features a recipe that achieves strong model scalability and data efficiency. It outperforms alternative designs in the hardest zero-shot generalization setting by up to $2.9\times$ task success rate given the same training data. With $10\times$ less training data, VIMA still performs $2.7\times$ better than the best competing variant. Code and video demos are available at https://vimalabs.github.io/
1. Introduction
The paper proposes multimodal prompts as a unified interface for diverse robot manipulation tasks and introduces VIMA and VIMA-BENCH to study scalable, generalist robot learning.
- Contributions: The work contributes a multimodal prompting formulation, a scalable benchmark, and an agent capable of multi-task and zero-shot generalization.
- 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.
- Motivation: Different manipulation tasks previously relied on distinct policy architectures, objectives, data pipelines, and training procedures, producing siloed systems.
- Contributions: VIMA-BENCH provides 17 multimodal tasks, procedurally instantiated into thousands of instances, with four progressively harder generalization levels.
- Contributions: 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.
2. Multimodal Prompts for Task Specification
The paper formulates task specification as an ordered sequence of interleaved text and image elements, enabling one interface for multiple manipulation paradigms.
- Prompt formulation: A multimodal prompt is an ordered sequence of arbitrarily interleaved textual and visual elements.Each element is text or an image.
- Prompt formulation: This formulation unifies goal conditioning, video demonstration, and natural language instruction within one sequence modeling problem.
- Task suite: The task suite includes simple object manipulation, visual goal reaching, novel concept grounding, and one-shot video imitation.
- Task suite: Additional categories cover visual constraint satisfaction and visual reasoning, including appearance matching and visual memory.
- Task suite: Task categories are not mutually exclusive, so a task may combine novel concept grounding with video demonstration or visual reasoning.
3. VIMA-BENCH: Benchmark for Multimodal Robot Learning
VIMA-BENCH is a procedurally generated simulator benchmark for multimodal robot learning, combining diverse tasks, expert trajectories, and increasingly difficult generalization tests.
- 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.
- Simulation environment: The simulator supplies frontal and top-down RGB observations, object annotations, and primitive motor skills such as pick-and-place and wipe.
- Dataset: The dataset contains 50K trajectories per task and 650K successful trajectories in total, with held-out objects, textures, and tasks for evaluation.
- Evaluation: Evaluation reports average binary task success across multiple simulator episodes without partial rewards.
- Evaluation: 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.
4. VIMA: Visuomotor Attention Agent
VIMA is an object-centric encoder-decoder transformer that conditions autoregressive robot control on multimodal prompts and interaction history.
- Architecture: VIMA learns a robot policy π(a_t|P, H) from a multimodal prompt P and past interaction history H.
- Architecture: A frozen pre-trained language model encodes prompt tokens, while the decoder predicts waypoint commands through cross-attention.
- Object-centric representation: VIMA represents visual inputs with object tokens derived from bounding boxes and cropped RGB patches rather than raw images.
- Robot controller: The controller alternates prompt-to-history cross-attention with causal self-attention before mapping action tokens to discretized robot-arm poses.
- 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.
5. Experiments
The experiments compare transformer-based agent designs across model sizes, dataset sizes, generalization levels, and component ablations. VIMA’s object-token representation with cross-attention conditioning is the most effective recipe among the considered designs, with strong scaling, data efficiency, and robustness.
- 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.
- 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.
- 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.
- 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.
- 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.
- 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.
6. Related Work
Prior work spans unified transformer and foundation-model approaches, multimodal interfaces, and diverse robot-manipulation benchmarks. VIMA-BENCH targets the missing combination of multimodal prompting and standardized generalization evaluation.
- Transformers unify multi-task learning across language, vision, and embodied-agent domains through shared sequence-modeling or backbone approaches.
- Multimodal models such as Perceiver, Flamingo, Frozen, and Gato process interleaved modalities or support broad task coverage, but VIMA focuses on robot prompting.
- Robot-manipulation research covers instruction following, one-shot imitation, rearrangement, constraint satisfaction, and reasoning across multiple simulation benchmarks.
- VIMA-BENCH is presented as the first robot-learning benchmark supporting multimodal-prompted tasks and standardized evaluation of generalization capabilities.
7. Conclusion
VIMA combines multimodal prompting, a procedurally generated simulator benchmark, and a transformer agent for diverse robot-manipulation tasks. The benchmark spans simple manipulation, visual goals, and novel concept grounding.
- VIMA converts diverse robot-manipulation tasks into a uniform sequence-modeling problem and uses one transformer agent across multiple task types.
- VIMA-BENCH is built on Ravens with procedurally generated tabletop instances formed from extensible 3D objects and textures.
- Simple object manipulation: Simple manipulation tasks specify objects, containers, textures, and rotations through multimodal prompts with corresponding success criteria and oracle trajectories.
- Visual goal reaching: Visual-goal tasks require rearranging objects to match prompted scene configurations, sometimes moving distractors and restoring the initial setup.
- 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”.
B.6. Visual Reasoning
VIMA-BENCH includes visual-reasoning tasks requiring agents to match appearances, use spatial memory, select neighboring objects, and execute ordered actions before restoration.
- Visual reasoning tasks require decisions based on information conveyed through multimodal prompts, including reasoning and memorization.
- Appearance matching: Tasks 14 and 15 require grouping objects by the container’s texture or top-down shape profile, respectively.
- Success criteria: The benchmark defines success using containment, correct neighbor selection, ordered placement, and final restoration criteria.
- Spatial memory: Task 16 requires placing a target object and then selecting its former neighbor using a specified cardinal direction.
- Ordered manipulation: Task 17 requires moving a target through multiple containers in order and restoring it to its original container.
C.1. Summary of Different Methods
VIMA conditions autoregressive action decoding on multimodal prompts and trajectory history, contrasting cross-attention with decoder-only and Perceiver-based alternatives. Its inputs are tokenized from text, scenes, objects, observations, and past actions.
- Baseline variants: Direct modeling concatenates prompt and trajectory tokens, whereas VIMA-Flamingo uses a Perceiver Resampler and VIMA-GPT uses decoder-only causal modeling.
- Prompt tokenization: Multimodal prompts contain text, full-scene images, and single-object images that are converted into token sequences.
- Prompt tokenization: Scene objects are represented using normalized bounding-box features and cropped-image features, then mapped into object tokens.
- 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.
- VIMA: VIMA uses cross-attention from trajectory-history tokens to prompt tokens, alternating with causal self-attention before predicting the next action token.
- Action decoding: Predicted action tokens are decoded through independent action heads into discrete pose components and then mapped to continuous actions.
D. VIMA Training Details
VIMA training uses standard Transformer optimization, object augmentation, and scalable model configurations across 2M–200M robot-controller parameters.
- Training uses AdamW, learning-rate warm-up, cosine annealing, and GEGLU activations across Transformer methods.
- 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.
- 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.
- Experiments run on nodes with 8 NVIDIA V100 GPUs, using distributed data parallelism; the largest experiment takes approximately one day.
E.2.1. NUMERICAL RESULTS
The numerical studies examine model and data scaling, pretrained-vision alternatives, prompt-encoder size, and training-from-scratch baselines. VIMA’s recommended design remains strongest in the reported practical scaling settings.
- 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.
- 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.
- 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.
- 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.
- MVP-pretrained ViTs help most at 10% data, while their advantage diminishes or becomes detrimental at full, 1%, and 0.1% data.
- 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.
E.5. Policy Robustness
VIMA’s robustness is evaluated against additional distractors and imperfect prompts, with results indicating small degradation under both types of perturbation.
- 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.
- Masked prompts produce minimal performance decrease, while corrupted prompts cause a minor decrease.The authors attribute this robustness to the pretrained T5 language backbone.
G. Limitations and Further Discussions
The paper identifies boundaries involving object detection, simulator realism and task complexity, and the high-level action space.
- 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.
- 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.
- VIMA currently uses only high-level pick-and-place and wipe actions, although its algorithm requires minimal modification for lower-level action spaces.