VIMA：基于多模态提示的通用机器人操作

ICML 2023

基于提示（prompt-based）的学习已在自然语言处理（NLP）中成为一种成功的范式，其中单一的通用语言模型可通过输入提示来执行任意指定任务。然而，不同的机器人任务目前仍由专用模型分别处理。本工作表明，我们可以使用多模态提示（multimodal prompts）——即文本与视觉 token 的交错序列——来表达广泛范围的机器人操作任务。

我们提出了 VIMA（VisuoMotor Attention agent），这是一种新颖的、可扩展的多任务机器人学习器，通过多模态提示实现了统一的序列输入输出（IO）接口。其架构遵循在 NLP 中已被证明有效且可扩展的编码器-解码器 Transformer 设计。VIMA 使用一个预训练的语言模型对交错的文本和视觉提示 token 序列进行编码，并在每个环境交互步骤中以自回归方式解码机器人控制动作。Transformer 解码器通过交叉注意力（cross-attention）层接收提示信息，这些层与常规的因果自注意力（causal self-attention）层交替堆叠。不同于直接处理原始像素，VIMA 采用以物体为中心（object-centric）的方法：我们使用现成的检测器将提示或观测中的所有图像解析为物体，并将其展平为物体 token 序列。上述设计共同构成了一个概念简洁、兼具强大模型与数据扩展能力的架构。

在本仓库中，我们提供了 VIMA 模型代码、覆盖多种模型规模的预训练检查点，以及演示和评估脚本。本代码库采用 MIT 许可证。

安装

VIMA 要求 Python ≥ 3.9。我们已在 Ubuntu 20.04 上完成测试。安装 VIMA 代码库非常简单：

pip install git+https://github.com/vimalabs/VIMA

预训练模型

我们在 Hugging Face 上托管了覆盖多种模型容量的预训练模型。下载链接如下所示。Mask R-CNN 模型可在此处获取：链接。

200M	92M	43M	20M	9M	4M	2M

基线方法实现

由于此前尚无方法能直接适用于我们的多模态提示设置，我们尽力选取了若干具有代表性的基于 Transformer 的智能体架构作为基线，并对其重新解释以兼容 VIMA-Bench。这些基线包括 VIMA-Gato、VIMA-Flamingo 和 VIMA-GPT。其实现代码位于 policy 文件夹中。

演示

要运行实时演示，请首先按照 VIMA-Bench 安装说明 安装 VIMA-Bench。随后可通过以下命令运行演示：

python3 scripts/example.py --ckpt={ckpt_path} --device={device} --partition={eval_level} --task={task}

其中 eval_level 表示四种评估级别之一，可选值包括 placement_generalization（位置泛化）、combinatorial_generalization（组合泛化）、novel_object_generalization（新物体泛化）和 novel_task_generalization（新任务泛化）。task 表示具体的任务模板。更多细节请参考 任务套件 和 基准评估。例如：

python3 scripts/example.py --ckpt=200M.ckpt --partition=placement_generalization --task=follow_order

运行上述命令后，应会弹出一个 PyBullet GUI 窗口，同时显示一个小窗口展示多模态提示。随后机械臂将开始移动以完成对应任务。请注意，由于 PyBullet GUI 需要图形显示，此演示在无头（headless）机器上可能无法运行。

论文与引用

我们的论文已发布于 arXiv。如果您觉得本工作对您有帮助，请考虑引用我们！

@inproceedings{jiang2023vima,
  title     = {VIMA: General Robot Manipulation with Multimodal Prompts},
  author    = {Yunfan Jiang and Agrim Gupta and Zichen Zhang and Guanzhi Wang and Yongqiang Dou and Yanjun Chen and Li Fei-Fei and Anima Anandkumar and Yuke Zhu and Linxi Fan},
  booktitle = {Fortieth International Conference on Machine Learning},
  year      = {2023}
}

VIMA 快速上手指南

环境准备

• 操作系统：推荐 Ubuntu 20.04（已在该系统测试）
• Python 版本：≥ 3.9
• 依赖框架：PyTorch
• 其他依赖：安装过程中会自动处理，无需手动配置

💡 提示：若在国内访问 GitHub 或 Hugging Face 较慢，建议配置镜像加速（如使用清华源）或通过代理下载模型。

安装步骤

直接通过 pip 安装 VIMA：

pip install git+https://github.com/vimalabs/VIMA

⚠️ 注意：如遇网络问题，可尝试替换为国内 Git 加速服务（如 https://ghproxy.com/https://github.com/vimalabs/VIMA）。

此外，运行演示需先安装 VIMA-Bench：

git clone https://github.com/vimalabs/VimaBench.git
cd VimaBench
pip install -e .

基本使用

1. 下载预训练模型（以 200M 模型为例）：

   wget https://huggingface.co/VIMA/VIMA/resolve/main/200M.ckpt
   

2. 运行示例任务：

   python3 scripts/example.py --ckpt=200M.ckpt --device=cpu --partition=placement_generalization --task=follow_order
   

   参数说明：

   • --ckpt：指定预训练模型路径
   • --device：运行设备（如 cpu 或 cuda）
   • --partition：评估级别，可选 placement_generalization、combinatorial_generalization、novel_object_generalization、novel_task_generalization
   • --task：具体任务名称，如 follow_order

🖥️ 注意：该演示依赖 PyBullet GUI，需在有图形界面的环境中运行，不支持纯命令行（headless）服务器。