VIMA: General Robot Manipulation with Multimodal Prompts

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. This work shows that we can express a wide spectrum of robot manipulation tasks with multimodal prompts, interleaving textual and visual tokens. We design a transformer-based generalist robot agent, VIMA, that processes these prompts and outputs motor actions autoregressively. To train and evaluate VIMA, we develop a new simulation benchmark with thousands of procedurally-generated tabletop tasks with multimodal prompts, 600K+ expert trajectories for imitation learning, and four levels of evaluation protocol for systematic generalization. VIMA achieves strong scalability in both model capacity and data size. It outperforms prior SOTA methods in the hardest zero-shot generalization setting by up to 2.9× task success rate given the same training data. With 10× less training data, VIMA still performs 2.7× better than the top competing approach. We open-source all code, pretrained models, dataset, and simulation benchmark at https://vimalabs.github.io.

StructDiffusion: Object-Centric Diffusion for Semantic Rearrangement of Novel Objects

“…ProgPrompt [7] and Code-As-Policies [15] require lots of prompt engineering and object detection capabilities, but they can complete fairly complex structures. Several of these works use an object-centric representation of the world [2], [14], [5], where objects are segmented or detected and encoded separately. For example, VIMA [2] used encoded object patches as input to a multimodal transformer.…”

“…Several of these works use an object-centric representation of the world [2], [14], [5], where objects are segmented or detected and encoded separately. For example, VIMA [2] used encoded object patches as input to a multimodal transformer. None of these works, however, look specifically at how we can ensure we are generating physically realistic structures: in our experiments, we show how these direct-regression-first approaches do not generate the same quality of structures, and that in particular, simply predicting placement poses or actions will lead to more failures.…”

“…Structures are everywhere in the real world: shelves are stocked, tables set, furniture assembled. Capturing all of these relationships requires models that can reason over both object geometry 1 Georgia Tech, 2 University of Utah and NVIDIA, 3 Meta AI and semantics at once, and also reason about physical validity when dealing with objects that have never been seen before.…”

“…Previous work in language-conditioned rearrangement has looked at 2D image representations [1], [2] instead of 3D ones. Other prior work which can perform some rearrange-ment tasks in 3D [3], [4] have limited generalization to novel objects.…”

“…Others, like StructFormer [5], regress to only a single solution, and cannot handle multiple challenging types of structures in a single model. To our knowledge, StructDiffusion is the first work that combines a multimodal transformer with a diffusion model to generalize over different objects given high-level language instructions; our method also overcomes limitations of prior techniques regarding voxelization and workspace size [1], [2], [3].…”