RoboMamba：用于高效机器人推理和操作的多模态状态空间模型

24年6月的论文“RoboMamba: Multimodal State Space Model for Efficient Robot Reasoning and Manipulation”。

机器人操控的一个基本目标是使模型能够理解视觉场景并执行动作。尽管现有的机器人多模态大语言模型 (MLLM) 可以处理一系列基本任务，但它们仍然面临两个方面的挑战：1) 推理能力不足以处理复杂任务，2) MLLM 微调和推理的计算成本高。最近提出的状态空间模型 (SSM) 称为 Mamba，在具有线性推理复杂性的非平凡序列建模中展示了良好的能力。受此启发，推出 RoboMamba，这是一种端到端机器人 MLLM，它利用 Mamba 模型提供机器人推理和动作能力，同时保持高效的微调和推理。具体来说，首先将视觉编码器与 Mamba 集成，通过共同训练将视觉数据与语言嵌入对齐，为模型提供视觉常识和与机器人相关的推理。为了进一步使 RoboMamba 具备动作姿势预测能力，探索一种具有简单策略头的高效微调策略。一旦 RoboMamba 具备足够的推理能力，它就可以用最少的微调参数（模型的 0.1%）和时间（20 分钟）获得操作技能。

数据极大地推动了大语言模型 (LLM) [1–3] 的研究，展示了自然语言处理 (NLP) 中推理和泛化能力的显著进步。为了理解多模态信息，人们引入多模态大语言模型 (MLLM) [4–8]，使 LLM 具有视觉指令跟踪和场景理解的能力。受到 MLLM 在一般环境中强大能力的启发，最近的研究旨在将 MLLM 融入机器人操作中。一方面，一些研究 [9–12] 使机器人能够理解自然语言和视觉场景，自动生成任务规划。另一方面，[13–15] 有效地利用了 MLLM 的固有能力，使其能够预测操作姿势。

机器人操作涉及与动态环境中的目标交互，需要类似人类的推理能力来理解场景的语义信息 [11, 16]，以及强大的低层动作预测能力 [17, 18]。虽然现有基于 MLLM 的方法可以处理一系列基本任务，但它们仍然面临两个方面的挑战。首先，在机器人场景中，预训练 MLLM [6, 19] 的推理能力被发现不足。当微调机器人 MLLM 遇到复杂的推理任务时，这一缺陷带来挑战。其次，由于基于注意机制的 LLM 成本高昂，微调 MLLM 并使用它们生成机器人操作动作会产生更高的计算成本 [20, 21]。为了平衡推理能力和效率，NLP 领域出现一些研究 [22–24]。值得注意的是，Mamba [25] 引入创新的选择性状态空间模型 (SSM)，在保持线性复杂性的同时促进了上下文-觉察的推理。受此启发，提出一个问题：“能否开发出一种高效的机器人 MLLM，它既具有强大的推理能力，又能以非常经济的方式获得机器人操作技能？”
RoboMamba 概览如图所示：其作为一款高效的机器人 MLLM，集推理和操控能力于一身。首先，将视觉编码器与高效的 Mamba 语言模型集成并对齐，为模型赋予常识和与机器人相关的推理能力。引入一种极其高效的微调策略，使 RoboMamba 具备姿势预测能力，只需 20 分钟即可微调一个简单的策略头（3.7M 个参数）。

对于机器人视觉推理，RoboMamba 根据图像 I 和语言问题 Lq 生成语言答案 La，表示为 La = R(I, Lq)。推理答案通常包含针对一个问题 Lq 的单独子任务 (La →(L1a, L2a, …, Lna))。

对于动作预测，用高效而简单的策略头 π 来预测动作 a = π(R(I, Lq))。根据前人的研究[54, 15]，用 6-DoF 来表达Franka Emika Panda 机械臂的末端执行器姿势。6-DoF 包括表示3D坐标的末端执行器位置apos和表示旋转矩阵的方向adir。如果训练抓取任务，将夹持器状态添加到姿势预测中，从而产生7-DoF控制。

为了使 RoboMamba 具备视觉推理和操控能力，从预训练的大语言模型 (LLM) [25] 和视觉模型开始，构建有效的 MLLM 架构。如图所示，利用 CLIP 视觉编码器 [26] 从输入图像 I 中提取视觉特征 fv，其中 B 和 N 分别表示批量大小和tokens。与 [55, 56] 相比，不采用视觉编码器集成技术，该技术采用各种主干（即 DINOv2 [57]、CLIP-ConvNeXt [58]、CLIP-ViT）进行图像特征提取。集成引入了额外的计算成本，严重影响了机器人 MLLM 在现实世界中的实用性。因此，简单直接的模型设计、与高质量数据和适当的训练策略相结合，也可以实现强大的推理能力。为了让 LLM 理解视觉特征，用多层感知器 (MLP) 将视觉编码器连接到 LLM。通过这个简单的跨模态连接器，RoboMamba 可以将视觉信息转换为语言嵌入空间 fLv。注：模型效率在机器人领域至关重要，因为机器人需要根据人类指令快速响应。

因此，选择 Mamba 作为语言模型，因为它具有上下文-觉察理能力和线性计算复杂度。使用预训练的token化器将文本提示编码到嵌入空间 ft 中，然后与视觉tokens连接 (cat) 并输入到 Mamba 中。利用 Mamba 强大的序列建模来理解多模态信息，并利用有效的训练策略来开发视觉推理能力。然后对输出tokens Ta 进行去token化 (det) 以产生自然语言 La 中的响应。

构建 RoboMamba 架构后，下一个目标是训练模型学习通用视觉和与机器人相关的推理能力。为了使模型具有推理和操作能力，精心设计全面的训练流程，分为两个阶段。

第 1 阶段的训练分为两个步骤：对齐预训练（第 1.1 阶段）和指令协同训练（第 1.2 阶段）。具体来说，与之前的 MLLM 训练方法 [19, 59, 55] 不同，目标是让 RoboMamba 能够理解常见视觉和机器人场景。鉴于机器人领域涉及大量复杂和新的任务，RoboMamba 需要增强的泛化能力。因此，在第 1 阶段采用协同训练策略，将高级机器人数据（例如任务规划）与通用指令数据相结合。协同训练不仅可以产生更具泛化的机器人策略，而且还能增强通用场景推理能力，因为机器人数据中嵌入了复杂的推理任务。

阶段 1.1：对齐预训练。采用 LLaVA [4] 过滤的 558k 图像文本配对数据集进行跨模态对齐。冻结 CLIP 编码器和 Mamba 语言模型的参数，仅更新项目层。通过这种方式，可以将图像特征与预训练的 Mamba 词嵌入对齐。

状态 1.2：指令协同训练。在此阶段，首先遵循之前的 MLLM 工作 [4, 55, 56] 进行通用视觉指令数据收集。采用 655K LLaVA 混合指令数据集 [4] 和 400K LRV-Instruct 数据集 [60]，分别旨在学习视觉指令跟随和缓解幻觉。请注意，缓解幻觉在机器人场景中起着重要作用，因为机器人 MLLM 需要根据真实场景而不是想象的场景来生成任务规划。例如，现有的 MLLM 可能会用“步骤 1：找到把手”来公式化地回答“打开微波炉”，但许多微波炉没有把手。接下来，整合 800K RoboVQA 数据集 [27] 来学习高级机器人技能，例如长期规划、成功分类、判别和生成可供性、过去描述和未来预测。在联合训练期间，冻结 CLIP 编码器的参数，并在合并的 180 万个数据集上微调投影层和 Mamba。使用交叉熵损失对 Mamba 语言模型的所有输出进行监督。

在第 2 阶段，提出机器人操纵微调，有效地为 RoboMamba 提供低层操纵技能。现有基于 MLLM 的操控方法 [15, 14] 需要在操控微调阶段更新投影层和整个 LLM。虽然这种范式可以开发动作预测能力，但它也破坏了 MLLM 的固有能力并需要大量的训练资源。为了应对这些挑战，在第 2 阶段提出了一种高效的微调策略。冻结 RoboMamba 的所有参数并引入一个简单的策略头来模拟 Mamba 的输出tokens。策略头包含两个 MLP，分别学习末端执行器的位置 apos 和方向 adir，总共占整个模型参数的 0.1%。

RoboMamba 仅预测图像中接触像素的 2D 位置 (x, y)，然后使用深度信息将其转换为 3D 空间。为了评估这种微调策略，用 SAPIEN 模拟 [28] 生成一个包含 10k 个末端执行器姿势预测的数据集。经过操控微调后，一旦 RoboMamba 拥有足够的推理能力，它就可以通过极其高效的微调获得姿势预测技能。由于微调参数最少 (7MB) 且模型设计高效，只需 20 分钟即可实现新操控技能。这一发现凸显了推理能力对于学习操控技能的重要性，并提出了一个新视角：可以有效地为 MLLM 配备操控能力，而不会损害其固有的推理能力。最后，RoboMamba 可以使用语言响应进行常识和与机器人相关的推理，并使用策略头进行动作姿势预测。

数据集（第 1 阶段）情况。在对齐预训练阶段，用 LLaVA-LCS 558K 数据集 [59]，它是 LAION-CC-SBU 数据集的精选子集，并附有说明。在指令协同训练阶段，创建一个总计 180 万个样本的组合数据集，包括 LLaVA-v1.5 655K 混合 [59]、LRV-INSTRUCT 400K [60] 和 RoboVQA 800K 数据集 [27]。

数据集（第 2 阶段）情况。对于机器人操作微调阶段使用的数据集，利用 SAPIEN [28] 设置一个带有来自 PartNet [45] 的铰接式物体交互式模拟环境。为了生成数据，使用 Franka Panda 机器人与目标随机交互。当成功操作时，记录末端执行器的成功 6-DOF 姿势，作为训练的真实标签。在训练集中，收集 20 个类别的 10K 图像。对于测试，使用一组 1.1K 图像，其中包括来自训练集的见过类别和未见过类别。
在训练之前，RoboMamba 加载预训练的 CLIP/SigLIP ViT-Large [26, 61] 作为视觉编码器，加载 2.8/1.4B Mamba [1] 模型作为语言模型。在对齐预训练和指令联合训练期间，分别进行 1 个 epoch 和 2 个 epoch 的训练。使用 AdamW 优化器，其中 (β1, β2) = (0.9, 0.999) 和 2e-5 的学习率 (LR)。浮点计算的精度设置为 16 位。对于操纵微调，对模型进行了 5 个epochs的训练，将 LR 设置为 1e-5 并应用 0.1 的权重衰减。浮点精度设置为 32 位。所有实验均在 NVIDIA A100 GPU 上进行。