2410.07864 (arxiv.org)

RDT-1B: a Diffusion Foundation Model for Bimanual Manipulation (rdt-robotics.github.io)

24年10月清华大学的论文

Bili-presentation

背景

双臂操作在人类生活是非常重要的，很多情况下，单手操作会比双手操作会有很大的优势。

双手操作在机器人技术中至关重要，但由于协调两个机械臂（导致多模态动作分布）的固有复杂性以及训练数据的稀缺性，开发基础模型极具挑战性。本文提出机器人扩散Transformer（RDT），一种用于双手操作的开创性扩散基础模型。

RDT 以扩散模型为基础，有效地表示多模态，以可扩展的Transformer来处理多模态输入的异质性并捕获机器人数据的非线性和高频性。

为了解决数据稀缺问题，引入物理上可解释的统一动作空间，它可以统一各种机器人的动作表示，同时保留原始动作的物理意义，促进学习可迁移的物理知识。

通过这些设计，在迄今为止最大的多机器人数据集上对 RDT 进行预训练，并将其扩展到 1.2B 参数，这是最大的基于扩散机器人操作基础模型。

在自建多任务双手数据集上对 RDT 进行微调，数据集包含超过 6K+ 集，改进其操作能力。在真实机器人上进行的实验表明，RDT 的表现明显优于现有方法。

它对未见过的目标和场景，表现出零样本泛化能力，理解并遵循语言指令，仅通过 1∼5 次演示即可学习新技能，并有效处理复杂、灵巧的任务。

输入：

  1. 文本的指令需求 Language Instruction
  2. 每一时刻都有
    1. 本体感觉 proprioception zt
    2. 带噪动作块 a~t
    3. 控制频率 c
    4. 扩散时间步长 k
    5. 图像输入 X={X1，X2，X3}
      1. 外部、右腕和左腕摄像头的一组图像

输出：

  1. 去噪后的动作块at：t+Ta
    1. 通常都是本体感受的一个子集

任务设定

作为一个具体的机器人任务

  1. Skill：verbs - 动作
  2. Objects：nouns - 动作要去发生的对象
  3. Scene： task enviorments 场景 动作发生的场合
  4. Modality： how skill is performed - 方式 例：用哪一只手去解决任务

本文旨在通过模仿学习来训练基础模型策略以实现可泛化性，

  • 一个有用的策略，我们需要其具备泛化性，对于一个没有见过的场景，物品能够正确完成既定目标
  • 在有限的数据集上训练，泛化到未曾见过的情况

为什么不是用RL来训练？

  1. 有Sim2Real 的 gap - 虚拟环境和现实的差距
  2. Reward design 奖励函数的设计 not dense
  3. not friendly to large models

挑战

  1. 一个强大的模型架构
    1. 有足够的表达能力
      1. 动作分布多峰的特点 - 所以不能用回归的方法去学习
    2. 灵活的处理不同的模态
      1. 大模型训练应该要比较稳定
  2. 如何在异构的数据上训练

那如果不能用回归的方法去训练，那分类呢？

Transformer + Discretzed Token

分类 - classification loss

  • lose information of the number
    • Cost（12，13） == Cost（12 ，120） 这两个的loss 是一样的
  • Uncoordinated Behaviors

Method

因为不能直接通过MSE回归去训练

所以想使用扩散模型来建模连续动作条件分布 p(at | l, ot)

  • 给定语言 l 和观察ot 然后建模出这个动作at
  • 扩散模型的好处
    • 表达能力很强：可以建模出动作空间的多峰分布
    • 质量很高：很多文生图 都是diffusion去完成的
  • 缺点：
    • 采样速度很差 - 对于image和video来说
      • 但是对于 action 是一个低维的物理量 不需要很频繁的采样

机器人策略学习中使用扩散模型来建模连续动作条件分布 p(at | l, ot)

问题：多模态性与确定性映射的不足

  • 给定语言指令 l 和观察 ot，存在多个合理的动作 at（多模态性）。
  • 如果使用确定性映射 (l, ot) → at 并通过回归学习训练数据 (l, ot, at) 元组，策略会学习动作的“平均值”。
  • 这个平均值可能落在可行动作分布之外（分布外），甚至完全不可行（例如，两个截然不同但都合理的动作的平均值可能是一个无效动作）。

解决方案：建模条件分布 p(at | l, ot)

  • 为了捕捉动作空间的多模态性，需要直接建模给定 l 和 ot 的动作分布 p(at | l, ot)。

选择扩散模型的原因

  • 在各种建模条件分布的方法中，扩散模型因其强大的表达能力和高质量的采样结果而脱颖而出。
  • 针对高维数据采样慢的顾虑：扩散模型对高维数据（如图像）采样慢是一个公认的缺点。然而，在机器人策略场景下：
    • 动作 at 的维度远低于图像（例如，关节角度/速度、末端执行器位姿等，维度通常在几十到几百）。
    • 因此，采样开销很小，足以满足实时或近实时控制的要求。
  • 结论：扩散模型是建模 p(at | l, ot) 的一个理想选择，这也正是 Diffusion Policy (Chi, 2023) 的核心思想。

但是：

将扩散模型应用于机器人数据的独特挑战

  • 虽然扩散模型在图像/视频上表现出色，但机器人数据（动作 at 和本体感受状态）具有根本不同的特性：
    • 非线性动力学 (de Wit, 2012)：机器人运动受复杂物理定律支配，动作效果非线性和耦合性强。
    • 高频变化可能性：物理交互（碰撞、接触摩擦、约束、材料阻尼/弹性）会导致状态和动作在极短时间内发生剧烈、非平滑的变化。
    • 不稳定的数值范围：传感器噪声、故障、极端工况（如高冲击力）可能导致物理量读数出现异常大或小的值（极端值）。
  • 与图像/视频数据的对比：
    • 图像/视频：虽然是高维，但天然具有时间和空间连续性/平滑性 (Chen, 2019; Liang, 2022)。帧间变化通常是渐进的、可预测的（物体移动、光照渐变）。像素值范围相对稳定（0-255 或 0-1）。
    • 机器人物理量：缺乏这种固有的平滑性，表现出不稳定性和非线性。数值范围可能剧烈波动。
  • 核心挑战：标准的扩散模型架构和训练流程（主要针对图像/视频设计）需要调整，以有效捕捉和学习机器人数据中这些特有的不稳定性、非线性动力学和高频变化模式。

多模态输入的异质性体现在结构上，即每种模态的格式和维数明显不同。这给多模态训练带来了挑战。为了解决这个问题，将这些不同的模态编码到统一的潜空间中。以下是编码方法：

  • 低维输入是代表机器人物理量的低维向量，包括本体感受、动作块和控制频率。为了对它们进行编码，用具有傅里叶特征的MLP去编码（Tancik，2020）），它可以有效捕捉低维空间中的高频变化。
  • 图像输入是高维的，包含丰富的空间和语义信息。为了提取紧凑的表示，用图像文本对齐的预训练视觉编码器 SigLIP（Zhai，2023）。在训练期间固定其权重以节省 GPU 内存。
  • 语言输入长度不一且高度抽象，由于其复杂性和模糊性而带来集成挑战。为了对它们进行编码，用预训练的基于 Transformer 语言模型 T5-XXL (Raffel et al., 2020)。还在训练期间固定其权重以节省 GPU 内存。

除了结构之外，异构性还具有不同输入中包含不同信息量的特征。

首先，不同模态的数据包含不同数量的信息。

例如，图像通常比文本包含更多信息，并且在编码后会产生更多tokens。其次，相同模态的不同输入可能包含非常不同的信息量。

例如，机器人的外部摄像头具有更全局的视图，并且比手腕摄像头包含更丰富的信息，如上图右上所示。在这种情况下，模型可能会学习一条捷径：只关注外部视图并忽略手腕视图，从而失去感知深度的能力。

为了解决这个问题，在编码过程中以一定的概率随机、独立地掩码每个多模态输入，以防止模型过度依赖特定的输入。

RDT 的网络结构。选择 Transformer 作为可扩展的主干网络 (Bao et al., 2023; Peebles & Xie, 2023)，并根据机器人问题的特点，对 Diffusion Transfomer (DiT) 进行了以下三个关键修改：

  • QKNorm & RMSNorm。输入的机器人物理量的不稳定数值范围，会导致梯度不稳定和数值溢出等问题，尤其是在训练大型基础模型时。为了解决这个问题，添加 QKNorm (Henry et al., 2020)，以避免在计算注意力时出现数值不稳定。此外，问题可以被视为时间序列预测任务，原始 DiTs 的 LayerNorm 中心化操作可能会导致 token 偏移和注意力偏移，从而破坏时间序列的对称性 (Huang et al., 2024)。因此，将 LayerNorm 替换为 RMSNorm（Zhang & Sennrich，2019），而无需进行居中操作。如图a显示，如果不进行这种修改，大规模预训练往往非常不稳定，甚至会爆炸。
  • MLP 解码器。为了提高非线性机器人动作的近似能力，用非线性 MLP 解码器替换最终的线性解码器，作为从潜空间到物理空间的投影。如图b中经验所示，如果没有这种设计，RDT 就无法有效捕捉非线性动态，从而失去完成需要精细操作的灵巧任务的能力→ fail to perform dexterous tasks
  • 交替条件注入 (ACI)。
    • 在模型中，图像和语言输入作为条件，它们是高维的并且长度可变，与传统 DiT（Xie et al.，2020）中的类标签条件形成对比。这些信息条件很难压缩成单个token，使得原始的自适应层规范化方法不合适。因此，采用交叉注意来适应不同长度的条件，避免进一步压缩时的信息丢失。
    • #Tokens(image) » #Tokens(language) 图像tokens通常比文本tokens多得多, 同时注入两种模态往往会掩盖与文本相关的信息，从而削弱指令跟踪的能力
    • 为了缓解这个问题，在连续层的交叉注意中策略性地交替注入图像和文本tokens，而不是在每一层都注入两者。

共享动作空间

在异构多机器人数据上进行训练。为了能够在异构多机器人数据上进行训练，需要一个在各个机器人之间共享的统一动作空间，该空间可以为多机器人动作提供统一的格式。

从机器人的原始动作空间到统一动作空间的映射应该是物理上可解释的，每个维度都应该具有明确的物理意义。

而且没有使用 Normalization

这可以鼓励模型从不同的机器人数据中学习共享的物理定律，从而提高从不同机器人的数据中学习的效率（Shah et al., 2023a）。

空间的设计包括两个步骤。首先，对于每个机器人，用一个空间来容纳其本体感受 zt 和动作 at。这是因为 at 通常是所需 zt+1 的子集（de Wit et al., 2012; Kouvaritakis & Cannon, 2016），因此 zt 的空间自然包含 at 的空间。

其次，设计一个统一的空间，其中包含大多数带有夹臂机器人的所有主要物理量。

将机器人的动作空间嵌入到这个统一空间中，方法是将原始动作向量的每个元素，根据其物理意义填充到统一动作空间向量的相应位置，其余位置则用填充符填充。

在这个统一空间下，对几乎所有带有夹臂的现代机器人数据进行 RDT 预训练，并大大扩展数据规模以满足基础模型的要求。具体来说，预训练数据集集合包括各种机器人的 46 个数据集，总大小为 1M+ 条轨迹和 21TB。