RDT-1B:双手操作机器人的扩散基础模型

24年10月清华大学的论文"RDT-1B: A Diffusion Foundation Model for Bimanual Manipulation"。

双手操作在机器人技术中至关重要,但由于协调两个机械臂(导致多模态动作分布)的固有复杂性以及训练数据的稀缺性,开发基础模型极具挑战性。本文提出机器人扩散Transformer(RDT),一种用于双手操作的开创性扩散基础模型。RDT 以扩散模型为基础,有效地表示多模态,以可扩展的Transformer来处理多模态输入的异质性并捕获机器人数据的非线性和高频性。为了解决数据稀缺问题,引入物理上可解释的统一动作空间,它可以统一各种机器人的动作表示,同时保留原始动作的物理意义,促进学习可迁移的物理知识。通过这些设计,在迄今为止最大的多机器人数据集上对 RDT 进行预训练,并将其扩展到 1.2B 参数,这是最大的基于扩散机器人操作基础模型。在自建多任务双手数据集上对 RDT 进行微调,数据集包含超过 6K+ 集,改进其操作能力。在真实机器人上进行的实验表明,RDT 的表现明显优于现有方法。它对未见过的目标和场景,表现出零样本泛化能力,理解并遵循语言指令,仅通过 1∼5 次演示即可学习新技能,并有效处理复杂、灵巧的任务。

双手操作对于机器人完成现实世界的任务至关重要(Edsinger & Kemp,2007)。对于实际应用,有用的操作策略应该能够推广到未见过的场景,例如新的目标和场景。然而,当前的方法要么依赖于特定于任务的原语(Salehian,2017;Rakita,2019;Grannen,2023a),要么仅限于小规模模型、数据和简单任务(Krebs,2021;Franzese,2023;Grannen,2023b;Zhao,2023;Grotz,2024;Liu,2024),因此仅表现出狭窄的泛化能力,无法完成复杂的任务。继自然语言处理 (Achiam,2023;Touvron,2023) 和计算机视觉 (Radford,2021;Kirillov,2023) 取得成功之后,实现可泛化行为的一个有希望方向是通过对大规模数据集进行模仿学习来开发基础模型。
然而,开发双手操作基础模型并非易事。一个主要原因是由于硬件成本高昂,特定双臂机器人的可访问数据非常稀缺 (Sharma,2018;Collaboration,2023),从而破坏了训练基础模型的数据密集型要求。受到最近在单手操作方面的尝试(Brohan,2023;Kim,2024)的启发,先在广泛的多机器人数据集上进行预训练,然后在目标双臂机器人上收集的小数据集上进行微调。这可以帮助将数据大小扩展到三个数量级,并有可能从其他机器人的数据集中学习可迁移的物理知识。然而,仍存在两个关键的技术挑战。首先,可泛化的基础模型需要在表达力和可扩展性方面都具有高能力的架构。双手操作中动作空间维度是单手操作的两倍,从而为可行的动作分布带来了更高的多模态性(Li,2006;Jia,2024)。

因此,模型必须具有足够的表达力才能捕捉动作分布中的多模态性。以前的方法(Zhao,2023;Brohan,2023;Kim,2024)通常无法满足此标准,导致性能不令人满意。此外,该架构需要有效地处理来自不同模态的输入,包括文本、图像和动作。它必须具有可扩展性,才能在大规模机器人数据上稳定地进行训练。其次,数据异质性是由不同机器人的物理结构和动作空间定义的变化引起的,在多机器人数据训练期间会导致负迁移并阻碍策略泛化(Pan & Yang,2009)。现有的方法要么丢弃具有不同动作空间的机器人,要么仅保留结构在整个机器人中保持不变的数据部分,但代价是丢失宝贵的数据(Brohan,2023;Ghosh,2023;Shah,2023a)。

学习双手操作策略的一大挑战是动作空间的高维性,这加剧了数据稀缺性(Zollner,2004;Smith,2012;Lioutikov,2016;Stepputtis,2022)和多模态行为(Colome ́ & Torras,2018;2020;Figueroa & Billard,2017;Sharma,2018;Xie,2020;Franzese,2023)。一些研究已经开发出更具成本效益的数据收集接口(Zhao,2023;Aldaco,2024),但它们仅限于特定的硬件配置,仍然不足以弥补可泛化策略的数据缺口。其他人则试图通过引入归纳偏差来减少数据要求,例如区分用于稳定和功能的两个臂(Grannen,2023b)、参数化运动原语(Batinica,2017;Amadio,2019;Chitnis,2020;Franzese,2023),或使用体素表示(Grotz,2024;Liu,2024)。这些方法使用强先验或简化建模,成功地减少动作空间,但代价是应用范围缩小,无法表达双手行为的多模态性(Pearce,2023)。

基础模型通过在大型多任务机器人数据集 (Collaboration et al., 2023; Brohan et al., 2022; Fang et al., 2023) 上训练多任务“通才”模型 (Brohan et al., 2022; 2023; Ghosh et al., 2023; Kim et al., 2024),在实现可泛化行为方面展现出巨大潜力。大多数研究采用大视觉语言模型来直接预测动作 (Brohan et al., 2022; Driess et al., 2023; Brohan et al., 2023; Collaboration et al., 2023; Kim et al., 2024)。虽然它们展示了对新目标和任务的泛化能力,但在应用于双手操作时,它们面临着量化误差和不协调行为的问题(Pearce,2023),这主要是由于它们对动作空间进行了离散化。为了提高精度,扩散模型已用于连续控制(Ho,2020;Chi,2023;Pearce,2023;Ghosh,2023)。(Ghosh,2023)在 Open X-Embodiment(Collaboration,2023)数据集(25 个数据集)的子集上预训练基于 Transformer 的扩散策略,最多有 93M 个参数。

首先需要制定任务并阐述挑战。为了在硬件上评估模型,选择 ALOHA 双臂机器人作为目标机器人,因为它是最具代表性的双臂机器人之一,适合通过远程操作收集人类演示数据(Zhao,2023;Fu,2024;Aldaco,2024)。如图显示了目标机器人的示意图,该机器人由两个带夹持器的手臂和三个摄像头组成。注:设置和基础模型适用于任何双臂夹持机器人,不过双手操作可行的动作分布带来了更高的多模态性。

请添加图片描述

考虑语言为条件带视觉的双手操作具体任务,这是机器人技术的基础,在家庭等现实场景中具有重要价值(Stepputtis,2020;Brohan,2022;Zhao,2023)。形式化地说,给定一个语言指令 l,策略在时间 t 时呈现一个观察 ot;然后它产生一个动作 at 来控制两个机械臂来实现 l 指定的目标。

双手操作中的特定任务通常由多个元素组成:技能(例如,动词“挑选”或“擦拭”),目标(例如,名词“瓶子”或“桌子”),场景(即执行任务的环境)以及描述如何执行技能的模态(例如,副词“用左手挑选瓶子”)。给定一个新任务,需要一个实用的策略来泛化到训练数据中未出现的新元素。这对于以前的基于规则方法以及使用小模型/数据基于学习的方法来说都是具有挑战性的。

本文旨在通过模仿学习来训练基础模型策略以实现可泛化性。然而,由于硬件成本高昂,特定双臂机器人的可用数据特别稀缺(<10K 轨迹),远远达不到训练基础模型的常见要求。为了解决这个问题,建议采用预训练和微调流程(Radford,2018),利用来自多个机器人的数据,灵感来自单手操作的最新进展 Octo (Ghosh,2023), Open-X (Collaboration,2023), OpenVLA (Kim,2024)。通过这种方式,将数据大小扩大三个数量级。具体来说,首先在大型多机器人数据集 Dpre(主要是单臂)上对模型进行预训练,然后在目标机器人 D 的数据集上进行微调。

目前存在的挑战有两个:
挑战 1:如何设计一个强大的架构?一个可通用的基础模型需要一个强大的架构。这个要求包括两个主要方面。首先,架构必须具有足够的表现力来捕捉动作分布中的多模态性。上图展示了一个玩具示例,其中机器人试图抓住一个立方体。可以看到,有很多种模式可以完成这项任务,而单手操作只能控制一个机械臂。在收集演示时,人类操作员可能会随机选择其中一种,从而导致收集的动作数据具有多模态性。其次,可扩展性对于这种架构是必要的。作为基础模型,它应该能够有效地处理来自各种模态(文本、图像、动作等)的异构输入,同时可扩展以在大型数据集上进行稳定训练。
挑战 2:如何在异构数据上进行训练?在多机器人数据上进行训练提出了数据异构性的独特挑战。不同机器人的物理结构和动作空间可能有很大差异。先前的尝试要么局限于具有相似动作空间的机器人子集(Yang,2023;Ghosh,2023;Kim,2024),要么仅保留具有相同结构的输入子集(Collaboration,2023;Yang,2024),但代价是丢失了大量信息。如何在如此异构的数据上训练模型,这在很大程度上仍未得到充分解决。

如图所示是具有 1B 参数的机器人扩散Transformer (RDT-1B),是一种用于双手操作的语言为条件视觉运动策略,具有对未见过场景的泛化能力,如灵活性、少样本学习、指令跟随和零样本泛化等。

请添加图片描述

RDT框架如图所示:将各种机器人的异构动作空间嵌入到统一的动作空间中,以进行多机器人训练。

请添加图片描述

由于多模态性,给定语言指令 l 和观察 ot,可能有许多可能的操作 at 来继续执行任务。如果将其建模为确定性映射 (l, ot) → at 并回归训练数据中的 (l, ot, at) 元组,策略将学习动作模式的“平均值”。这可能会导致分布外的动作,例如多个模式的算术平均值,这可能完全不可行(Pearce,2023)。相反,选择对连续条件分布 p(at|l,ot) 进行建模。

在各种方法中,扩散模型在表达能力和采样质量方面都表现出色,但对高维数据(例如图像)的采样速度可能很慢。幸运的是,对于本文设置,缺点很小,因为 at 的维度比图像低得多,这只需要最小的采样开销。这使得扩散模型成为 diffusion policy(Chi 2023)中策略的理想选择。

然而,将扩散模型应用于机器人任务面临着独特的挑战,因为机器人物理量的固有属性(即动作和本体感受)与图像/视频数据不同。图像和视频数据虽然是高维的,但通常表现出一定程度的时间和空间连续性(Chen,2019;Liang,2022),帧之间的变化通常是增量的。相比之下,机器人物理量的特点是其非线性动力学(de Wit,2012)以及由物理相互作用(例如碰撞、约束和阻尼等材料特性)引起的高频变化的可能性。此外,这些量还具有不稳定的数值范围,可能是由于不可靠的传感器导致的极端值。这强调了调整当前扩散模型以有效捕捉机器人数据的不稳定性和非线性的必要性。

多模态输入的异质性体现在结构上,即每种模态的格式和维数明显不同。这给多模态训练带来了挑战。为了解决这个问题,将这些不同的模态编码到统一的潜空间中。以下是编码方法:
• 低维输入是代表机器人物理量的低维向量,包括本体感受、动作块和控制频率。为了对它们进行编码,用 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 就无法有效捕捉非线性动态,从而失去完成需要精细操作的灵巧任务的能力。
• 交替条件注入 (ACI)。在模型中,图像和语言输入作为条件,它们是高维的并且长度可变,与传统 DiT(Xie et al.,2020)中的类标签条件形成对比。这些信息条件很难压缩成单个token,使得原始的自适应层规范化方法不合适。因此,采用交叉注意来适应不同长度的条件,避免进一步压缩时的信息丢失。此外,鉴于图像tokens通常比文本tokens多得多,同时注入两种模态往往会掩盖与文本相关的信息,从而削弱指令跟踪的能力(定量结果见图 b)。为了缓解这个问题,在连续层的交叉注意中策略性地交替注入图像和文本tokens,而不是在每一层都注入两者。

请添加图片描述

在异构多机器人数据上进行训练。为了能够在异构多机器人数据上进行训练,需要一个在各个机器人之间共享的统一动作空间,该空间可以为多机器人动作提供统一的格式。从机器人的原始动作空间到统一动作空间的映射应该是物理上可解释的,每个维度都应该具有明确的物理意义。这可以鼓励模型从不同的机器人数据中学习共享的物理定律,从而提高从不同机器人的数据中学习的效率(Shah et al., 2023a)。

空间的设计包括两个步骤。首先,对于每个机器人,用一个空间来容纳其本体感受 zt 和动作 at。这是因为 at 通常是所需 zt+1 的子集(de Wit et al., 2012; Kouvaritakis & Cannon, 2016),因此 zt 的空间自然包含 at 的空间。其次,设计一个统一的空间,其中包含大多数带有夹臂机器人的所有主要物理量。将机器人的动作空间嵌入到这个统一空间中,方法是将原始动作向量的每个元素,根据其物理意义填充到统一动作空间向量的相应位置,其余位置则用填充符​​填充。

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

如图所示:微调数据集的特征有多样性目标、挑战性任务、和多模态特征。

请添加图片描述

收集全面的多任务双手数据集。尽管已经在大规模数据集上进行了预训练,但由于具身差距,RDT 仍然可能需要帮助才能零样本推广到目标双臂机器人。为了弥补这一差距,需要在目标机器人上收集多任务双手数据集以进行微调。大语言模型的最新进展(Ziegler,2019;Brown,2020;Touvron,2023)表明,高质量的微调数据集对于模型性能至关重要。

从三个方面确保数据集的高质量:(1)在数量方面,收集 6K+ 条轨迹,使数据集成为当今最大的双手数据集之一;(2)在全面性方面,考虑 300 多项具有挑战性的任务,涵盖了大多数操作任务类型,从拾取和放置、到插入电缆,甚至包括编写数学方程式;(3)在多样性方面,准备 100 多个具有各种尺寸和纹理的刚性和非刚性目标以及 15 多个具有不同照明条件的不同房间。此外,进一步利用 GPT-4-Turbo (Achiam et al., 2023) 重写人工注释的指令以增加文本多样性。

选取了 7 个具有挑战性的任务,从不同维度评估 RDT 的通用性和能力,包括模型在现实任务中可能遇到的复杂场景,例如各种看不见的元素和灵巧的操作。下表给出了每个任务维度的说明,而如图提供了详细的定义和可视化。

请添加图片描述
请添加图片描述
请添加图片描述

猜你喜欢

转载自blog.csdn.net/yorkhunter/article/details/143205775