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

Seminal October 2022

Visual goal plus memory: rearrange objects to match a scene image, then restore the original layout.

TL;DR One transformer agent reads prompts that interleave text with images or video frames and outputs tabletop manipulation actions across 17 simulated tasks. On unseen tasks it scores up to 2.9x the success rate of the best alternative design, and with 10% of the data it still beats baselines trained on all of it.

What they do

Robot tasks are usually specified in one format per model: a language instruction, a goal image or a demonstration video. VIMA expresses all of them as one multimodal prompt, an interleaved sequence of text and images, and trains a single transformer to map prompt and interaction history to actions. They also build VIMA-Bench, a Ravens-based tabletop simulator with 17 task templates across six categories (simple manipulation, visual goals, novel concepts, one-shot video imitation, visual constraints, visual reasoning) and four generalization levels: new placements, new combinations, new objects, new tasks.

How it works

Object tokens. A fine-tuned Mask R-CNN parses every prompt image and observation into objects; each becomes a token from its bounding box and a ViT encoding of its crop. A pre-trained T5-Base encodes the interleaved prompt.
Cross-attention controller. A causal transformer decoder alternates self-attention over the history with cross-attention to the prompt and outputs discretized end-effector poses for pick-and-place and wipe primitives.
Training. Behavior cloning on 650K oracle trajectories (50K per task). Randomly injected false-positive detections make the policy robust to detector errors.
Baselines. Gato-, Flamingo- and GPT-style variants get the same prompt encoder; all are trained at 2M-200M parameters.

Results

At 200M parameters VIMA reaches 81.5 / 81.5 / 78.7 / 48.6% success on levels L1-L4; the strongest baseline, VIMA-Gato, gets 58.0 / 53.1 / 46.0 / 16.8%.
On novel tasks (L4) at 20M parameters VIMA scores 48.8% against 12.3% for VIMA-Gato, the 2.9x headline. Trained on 10% of the data it scores 46.0% on L4; VIMA-Gato trained on the full dataset gets 12.2%.
Ablations: ground-truth boxes barely beat the Mask R-CNN pipeline; pixel-based tokenizers, or compressing objects into a fixed number of tokens, do worse; cross-attention beats a decoder-only variant mainly at small model sizes and on L4. T5 encoder size (30M-368M) makes no significant difference.

Limitations

Simulation only, with high-level pick-and-place and wipe primitives rather than low-level control; physical realism was not a focus. VIMA inherits errors from the standalone object detector, which may struggle with occlusion or out-of-distribution object shapes. Success on novel tasks stays below 50% at every model size.

Videos

Visual constraint: sweep two objects into a region without touching a boundary line.

Visual reasoning: put an object into a container, then the object that was south of it into the same one.

Figures

From Figure 1: Many manipulation tasks can be written as multimodal prompts that interleave language with images or v... — From Figure 1: Many manipulation tasks can be written as multimodal prompts that interleave language with images or video frames (left); VIMA reads the prompt and controls a robot arm to solve the task (right).

From Figure 3: VIMA architecture. A pre-trained T5 encodes the multimodal prompt; a causal transformer decoder with a... — From Figure 3: VIMA architecture. A pre-trained T5 encodes the multimodal prompt; a causal transformer decoder with alternating self- and cross-attention layers predicts motor commands from the prompt and the interaction history.

From Figure 2: The four zero-shot evaluation levels in VIMA-Bench, from new object placements to new tasks; each leve... — From Figure 2: The four zero-shot evaluation levels in VIMA-Bench, from new object placements to new tasks; each level deviates further from the training distribution than the last.

From Figure 4: Top, success rate of VIMA and the baselines at model sizes from 2M to 200M parameters on each generali... — From Figure 4: Top, success rate of VIMA and the baselines at model sizes from 2M to 200M parameters on each generalization level. Bottom, success rate of 92M-parameter models trained on 0.1%, 1%, 10% and the full dataset.

Abstract

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/

Summaries are written automatically from each paper's text. Figures and videos come from the paper and its project page; all credit goes to the authors.

Matt Strong