VIMA: General Robot Manipulation with Multimodal Prompts

In plain words

Robots can be told what to do in many ways: a sentence, a picture of the goal arrangement, a short video of someone doing the task, or a rule such as "do not touch that area". Usually each way needs its own system, so the systems do not combine well.

VIMA's idea is to put all of these into one prompt that mixes words and pictures of objects, like a sentence with pictures inside it. A single transformer reads the prompt together with the robot's recent history and outputs the next small movement, one step at a time. It treats each object as its own item, a picture plus a box, instead of looking at the whole table image. The results below test whether this design is worth it.

The model is trained by copying demonstrations made by a scripted expert in a simulator. It is then tested on tasks it has not seen, including new kinds of tasks, and the paper's main claim is that it holds up better than other designs as the tests move further from the training data.

TL;DR

VIMA is a robot policy for tabletop manipulation whose task is given as an interleaved text-and-image prompt. It uses object-centric tokens, a pretrained T5 prompt encoder and a decoder that alternates cross-attention to the prompt with causal self-attention. It is trained by behaviour cloning in the authors' simulator benchmark and reports the best zero-shot results among the designs it compares, at every model size (author-reported). Code is MIT licensed LICENSE; checkpoint licensing was not checked.

Problem

Robot task specification is fragmented: goal images, video demonstrations, language instructions and visual constraints have each been handled by separate policies with different architectures, objectives and data pipelines §1. The paper asks whether one sequence-modelling interface can express all of them, and whether one agent trained on that interface generalises to unseen objects, textures and tasks §1, §2.

No prior method works directly with multimodal prompts, so the authors compare against adapted transformer baselines §5.1.

Method

The policy is \pi(a_t \mid \mathcal{P}, \mathcal{H}), where \mathcal{P} is the prompt and \mathcal{H} = [o_1, a_1, \ldots, o_t] is the interaction history §4. A prompt is an ordered sequence of text and image elements §2.

Step 1: tokenize the prompt §4 Tokenization app. C.2.1.

• Text: T5 tokenizer and word embeddings, dimension 768 tab. 2.
• Full-scene images: a fine-tuned Mask R-CNN extracts each object as a bounding box plus a cropped image app. C.4. Each object becomes one token.
• Single-object images: the same, with a dummy box app. C.2.1.
• Object token: the box (centre x, centre y, height, width) is normalised and passed through an MLP; the crop is passed through a ViT trained from scratch; the two are concatenated and mapped to the embedding dimension app. C.2.1. Hyperparameters in tab. 2.
• Non-text tokens go through an adapter MLP before the T5 encoder §4 app. C.2.1.

Step 2: encode the prompt §4 Tokenization app. C.2.1. A pretrained T5-base encoder encodes the interleaved sequence. Only its last two layers are fine-tuned, with layer-wise learning-rate decay; the other layers are frozen app. C.2.1. The T5 size ablation is in app. E.4.

Step 3: tokenize the observation app. C.2.2 tab. 3. Each observation gives object tokens from two views, ordered frontal then top-down. The end-effector state is one-hot encoded, embedded to dimension 2, concatenated with the object tokens and projected to observation tokens of dimension 768. Positional embeddings are learned and absolute.

Step 4: tokenize past actions app. C.2.3. Each past action is embedded by a two-layer MLP with hidden size 256, then mapped to the token dimension.

Step 5: the robot controller §4 Robot Controller app. C.2.4. The decoder is a causal stack of L blocks, each alternating cross-attention and causal self-attention. Queries come from the trajectory tokens and keys and values come from the prompt tokens. The cross-attention step is written as

\mathcal{H}' = \mathrm{softmax}\!\left(\frac{Q_\mathcal{H} K_\mathcal{P}^\top}{\sqrt{d}}\right) V_\mathcal{P}

with residual connections, where d is the embedding dimension §4. The paper's appendix pseudocode adds a feed-forward block after each attention step app. C.2.4, Pseudocode 1. The number of blocks per model size is in tab. 8 and tab. 9; the paper does not give L in the text, so reading it from "Num Blocks" is an inference.

# Robot controller: cross-attention, then causal self-attention
# [§4] [app. C.2.4]
# Shapes: prompt (L_p, d); trajectory (2T-1, d), obs and action tokens

def controller(prompt, trajectory):
    # Add learned positional embeddings to both sequences  [app. C.2.4]
    x = trajectory + trajectory_position_embedding
    prompt = prompt + prompt_position_embedding
    # Repeat for L blocks  [app. C.2.4]
    for block in range(num_blocks):
        # Trajectory tokens query the prompt  [§4]
        x = x + cross_attention(query=x, key_value=prompt)
        x = x + feed_forward(x)  # [app. C.2.4, Pseudocode 1]
        # Causal self-attention over the trajectory  [§4]
        x = x + causal_self_attention(x)
        x = x + feed_forward(x)
    # The last token is the predicted action token  [app. C.2.4]
    return x[-1]

Step 6: decode the action app. C.2.5 tab. 4. An action is two SE(2) poses, each with a position and a quaternion rotation. Each pose uses six independent categorical heads: two for position (50 bins for x, 100 bins for y) and four for the quaternion (50 bins each). The heads are integrated and mapped to continuous values by an affine transform. The paper reports that independent heads are as good as autoregressive decoding in early ablations app. C.2.5.

Training §4 Training app. D tab. 7. Behaviour cloning minimises the negative log-likelihood of the expert actions over the T steps of each trajectory §4. Object augmentation randomly injects false-positive detections into every observation to make the model robust to detector errors §4:

# One behaviour-cloning update  [§4 Training] [app. D]
# Each observation receives extra false-positive objects  [app. D]

def training_step(prompt, history, expert_actions):
    # Sample how many fake objects to add: 0 with p=0.95, 1 with p=0.05
    extra = sample_categorical(values=[0, 1], probs=[0.95, 0.05])
    # Add that many random boxes and crops to the object tokens  [app. D]
    objects = add_random_objects(history.objects, count=extra)
    # Predict per-head action distributions  [app. C.2.5]
    dists = policy(prompt, objects, history.actions)
    # Negative log-likelihood summed over the trajectory  [§4 Training]
    loss = -sum(dists[t].log_prob(expert_actions[t]) for t in steps)
    return loss

Checkpoints are chosen by aggregated accuracy on a held-out validation set. The evaluation itself uses rollouts in the physics simulator §4 Training.

Data. Expert trajectories come from scripted oracles with privileged state in VIMA-Bench, built on the Ravens simulator §3. The paper says it has 50K trajectories per task and 650K successful trajectories in total §3. See Pitfalls for the arithmetic.

Model sizes. Seven robot controllers from 2M to 200M parameters; the T5 encoder (111M) is not counted §5.2 tab. 8 tab. 9.

Hyperparameters

Complete training hyperparameters tab. 7:

Prompt tokenization tab. 2:

Observation encoding tab. 3:

Action decoder tab. 4:

Other settings in the text: object augmentation K=2 with probabilities \{0: 0.95, 1: 0.05\} app. D; action-token MLP hidden size 256 app. C.2.3; Mask R-CNN fine-tuning uses learning rate 5\times10^{-4} and 10 epochs, with 100 trajectories per task giving 22,741 images and 61,822 annotations app. C.4; the model with high recall is favoured for the detector app. C.4.

Controller configurations for cross-attention models (VIMA and VIMA-Flamingo) tab. 8:

Controller configurations for causal self-attention models (VIMA-Gato and VIMA-GPT) tab. 9:

Baseline settings: the Gato-style ViT uses image size 64 × 128, patch size 32, width 768 and 4 layers tab. 5. The Flamingo-style Perceiver Resampler uses 4 latent queries, 4 blocks, 4 self-attention layers per block and 24 heads tab. 6.

Results

Author-reported. Benchmark: VIMA-Bench on the Ravens simulator, 17 task templates in six categories, evaluated on a four-level zero-shot protocol §3. Level 1 is placement generalisation, Level 2 combinatorial generalisation, Level 3 novel objects and textures, and Level 4 novel tasks §3 fig. 2. Each task gives a binary success signal §3. Baselines are VIMA-Gato, VIMA-Flamingo and VIMA-GPT, each re-implemented by the authors §5.1 app. C.3.

Headline comparison at 200M parameters, average success % tab. 10:

Model scaling, all sizes tab. 10 (columns are controller parameters):

Data scaling, fraction of the imitation dataset tab. 11 fig. 4. Baselines that learn from pixels use an MVP-pretrained ViT-B with MAE fine-tuning on the detector's in-domain data app. E.3:

Other reported results (all author-reported):

• Progressive generalisation: VIMA's drop from L1 to harder levels is small, especially L1 to L2 and L1 to L3; baselines degrade by as much as 20%; on L4 VIMA's drop is about half that of the other variants §5.2 fig. 5.
• T5 encoder size, 200M controller tab. 13: t5-small (30M) / t5-base (111M) / t5-large (368M): L1 78.8 / 81.5 / 80.8; L2 79.0 / 81.5 / 81.0; L3 80.3 / 78.7 / 81.0; L4 49.1 / 48.6 / 49.3. The paper finds no significant difference and uses t5-base §5.3 app. E.4.
• Distractors, 200M tab. 14: adding one distractor object per task lowers L1 81.5 → 78.5, L2 81.5 → 78.6, L3 78.7 → 72.9, L4 48.6 → 47.8 (relative drops 3.6%, 3.5%, 7.3%, 1.6%) app. E.5.
• Imperfect prompts, 200M tab. 15: incomplete prompts (words masked with <UNK>) L1 80.8, L2 81.1, L3 77.0, L4 48.0; corrupted prompts (words swapped) L1 78.2, L2 78.1, L3 73.8, L4 45.3 app. E.5.
• Ablations §5.3 fig. 6 fig. 7: object tokens beat raw-pixel tokenizers (image Perceiver, image patches, single image embedding) at every level; passing the variable-length object sequence beats compressing it with a Perceiver; cross-attention conditioning beats a GPT-style concatenated sequence, most clearly at small sizes and on L4. Using the oracle boxes instead of Mask R-CNN gives only a small gain, which the paper reads as the detector pipeline losing little. The per-task numbers are in tab. 16 to tab. 19.
• Compute: the largest experiment takes about one day on nodes with 8 V100 GPUs using DDP app. E.1.

Paper vs code

Code: official repository vimalabs/VIMA at commit 8449837 (cloned 2026-10-18), MIT licence LICENSE. The repository has the policy and model code and the pretrained model links; it does not contain the training script, the Mask R-CNN training code, or the object-augmentation code. The arXiv version used here is v2 (2023-05-28).

• Cross-attention block. §4 gives the bare formula \mathrm{softmax}(QK^\top/\sqrt{d})V with a residual §4. The code uses learned query, key and value projections, an output projection, a pre-LayerNorm, and a GEGLU feed-forward with residuals components.py. The attention logits are scaled by the per-head dimension, not d components.py. The appendix pseudocode matches the code, not the §4 formula app. C.2.4.
• Self-attention block normalisation. The causal GPT-style block is post-LayerNorm: n = ln_1(x + a), h = ln_2(n + m) components.py, while the cross-attention block is pre-LayerNorm components.py. The paper does not specify either.
• Position ids use the object mask. The paper says "learned absolute" positions app. C.2.2. The code takes position ids as a cumulative sum of the token mask, so masked or absent objects do not consume positions vima_policy.py xattn_gpt.py.
• Prompt length cap. The code asserts that the prompt has at most 256 tokens xattn_gpt.py (constructor argument at vima_policy.py). The paper gives no cap.
• Bounding-box normalisation. The paper says boxes are divided by an upper bound, without giving it app. C.2.1. The code divides by 256 for width and 128 for height, in the order it uses obj_encoder.py obj_encoder.py.
• Action embedding. The code embeds each pose component with a one-hidden-layer MLP of width 256 vima_policy.py, which is the paper's "two-layer MLP" app. C.2.3. The hidden_depth argument counts hidden layers utils.py. The four component embeddings are concatenated and mapped to the token dimension by a linear layer action_embd.py; the paper does not mention this.
• Action decoder heads. Six categorical heads per pose, 50/100 bins for x/y and 50 bins for each of four quaternion components, match the paper vima_policy.py tab. 4.
• Discretisation. Bin boundaries are linspace(0, 1, n) with bucketize, but de-discretisation divides the index by n, not n-1 vima_policy.py. See Pitfalls.
• T5 fine-tuning. The paper fine-tunes only the last two T5 layers with layer-wise learning-rate decay app. C.2.1 app. E.4. The released policy loads all of t5-base prompt_encoder.py and contains no freezing. A search for freeze and requires_grad found no freezing logic; the only requires_grad use is inside the T5 attention code prompt_encoder.py. The fine-tuning setup therefore cannot be confirmed from the code.
• Object augmentation. Described only in the paper §4 app. D. Not in the repository.
• Training data. The paper says 650K successful trajectories in total §3. The model card says 600K for training and 50K held out for validation model-card.md. The training data link is in README.md.
• Action bounds. The mapping from normalised values to simulator bounds is in the example script, not the policy example.py. The bounds are read from the environment metadata and are not in the repository.
• Dropout and GEGLU. Dropout 0.1 and GEGLU match the paper vima_policy.py tab. 7.

Pitfalls

• The headline ratios use different conventions, and one baseline is wrong. The "2.7×" in the abstract and introduction is computed as the gap over the baseline, divided by the baseline and truncated: (46.0 − 12.2) / 12.2 ≈ 2.77 app. E.3. The baseline used is Gato at full data (12.2), but GPT is higher at full data (14.3) tab. 11. Using GPT gives (46.0 − 14.3) / 14.3 ≈ 2.2. The "2.9×" in the introduction is not explained. At 200M on L4 it equals the plain ratio 48.6 / 16.8 tab. 10. Ratios at matched size reach about 3.8× (43M: 49.0 vs 12.8 for Gato). Quote success rates, not these ratios.
• Dataset arithmetic. "50K trajectories per task" over 17 tasks would give 850K, not the stated 650K §3. The number per task is therefore not reliable. The model card gives 600K training plus 50K validation model-card.md.
• Discretisation is not an exact inverse. The encoder places bin edges at k/(n-1) for n bins, but the decoder maps index k to k/n vima_policy.py vima_policy.py. A pose exactly at 1.0 decodes to 0.98 for x and rotation, and 0.99 for y. Errors are up to about one bin width. A reimplementation that decodes to bin centres, or to k/(n-1), will not reproduce the released action values exactly.
• Unstated training details. Batch size, total steps or epochs, AdamW betas, minimum learning rate of the cosine schedule, and number of evaluation episodes per task are not stated app. D §3. The released code does not include the training loop, so these cannot be filled from code.
• Layer count L. The text says "L alternating" blocks §4 but gives no value. tab. 8 and tab. 9 give "Num Blocks" per size; reading that as L is an inference.
• Attention masks. The causal mask uses fill value -10^4 rather than -10^9 components.py. Cross-attention padding masks use the dtype's minimum value components.py. Keep logits in float32 as the code does components.py.
• Object-token cost. Observation tokens scale with the number of detected objects, typically 3 to 8 per step tab. 1. Sequence length grows with the number of objects in each view. Direct sequence modelling of the prompt is quadratic in the number of observation tokens app. C.1; cross-attention avoids putting the prompt in the trajectory sequence §4. The paper gives no latency figures.
• Detector dependency. Object tokens depend on a Mask R-CNN fine-tuned on about 22,741 simulated images app. C.4. The paper reports minimal loss versus oracle boxes §5.3, but that is in simulation. Real-robot transfer is not tested; the authors say the checkpoints may not generalise to other simulators or the real world model-card.md.
• Baselines are the authors' reimplementations. The paper says it made its "best effort" to re-interpret Gato, Flamingo and GPT-style designs for multimodal prompts §5.1. Gains against these baselines measure the authors' versions, not the originals.
• Robustness of the T5 choice. The paper attributes robustness to the T5 backbone §5.3 app. E.5, but the T5 size ablation shows little difference between sizes tab. 13.

Lineage

• Builds on: T5 as prompt encoder §4, following the approach the paper cites as "Frozen" §4; object detection with Mask R-CNN app. C.4; the Gato, Flamingo and Decision Transformer designs for the baselines §5.1 app. C.2.3.
• Compared against: VIMA-Gato, VIMA-Flamingo, VIMA-GPT §5.1.
• Benchmark: VIMA-Bench, built on the Ravens simulator §3. The simulator is a separate repository, vimalabs/VimaBench model-card.md.
• Superseded by: not known from content in this wiki.
• Superseded by: not known from content in this wiki.