RDT-1B: A DIFFUSION FOUNDATION MODEL FOR BIMANUAL MANIPULATION
RDT-1B: A DIFFUSION FOUNDATION MODEL FOR BIMANUAL MANIPULATION。双臂机器人操作,扩散模型
abstract
- 解决coordination和traning data lack问题
- Robotics Diffusion Transformer
- 基于扩散基础模型,bimanual manipulation
- diffusion model+Transformer to deal with multi-modal
- for data scarcity: Physically Interpretable Unified Action Space
- pre-train on large and scale up
- fine-tune on self multi-task dataset
- experiments on real robot show outperform,good generalization
introduction
- 问题:narrow generalization and failing in complex tasks
- 解决:“imitation learning on large-scale datasets” (选中模仿学习)
- 问题:数据缺少
- 解决:首先对广泛的多机器人数据集进行预训练,然后对目标双臂机器人上收集的小数据集进行微调。
- 有两个难点
- 需要一个架构,有表现力和可扩展性
- 数据异构
- RDT
- 最大的双手操控基础模型,通用性强,用“Diffusion Transformers (DiTs)”作为骨干网络
- 设计思路
- expressiveness:利用扩散模型的能力来建模复杂的分布
- scalability:利用 Transformer 主干并精心设计多模式编码
- characterize:对原有DiT结构进行重要修改,包括 MLP 解码、改进的归一化和交替条件注入
- heterogeneous:提出物理上可解释的统一行动空间
- 实验
- 在最大集合上预训练 RDT 模型
- 在自收集的多任务双手数据集上微调 RDT
PROBLEM FORMULATION AND CHALLENGES
- 选择
- ALOHA双臂机器人
- 语言条件下的视觉双手操作
- instruction,observation,time,action,RGB observation,low-dimensional
proprioception of the robot,control frequency的设置
- low-dimensional proprioception of the robot 不知道是什么
- task:给定一项新任务,需要一种实用的策略来泛化到训练数据中未出现的未见过的元素。
- policy:预训练和微调管道
- 开发这样的“foundation model” 有俩challenge(introduction提到的)
- architecture
- expressiveness
- scalability
- heterogeneous data
- architecture
RDT
整体的框架见论文图3
RDT MODEL
为了解决第一个challenge。
Diffusion Modeling
将扩散模型应用于机器人任务的2个挑战
- 机器人物理量的特点是其非线性动力学和物理相互作用(如碰撞、约束和阻尼等材料属性)可能引起的高频变化
- 不稳定的数值范围,可能是传感器导致的
Encoding of Heterogeneous Multi-Modal Inputs
这个异构是SD带来的,和第二个challenge不一样。
为了将这些不同的模态编码到一个统一的潜在空间中,用以下编码方法
- 针对表示物理量的低维输入,MLP(具有傅里叶特征可以有效地捕捉低维空间中的高频变化
- 图像输入是高维的,包含丰富的空间和语义信息,用图像-文本对齐的预训练视觉编码器 SigLIP
- 语言输入的长度各不相同,并且高度抽象,用预先训练的基于 Transformer 的语言模型 T5-XXL
除了结构之外,异构性还意味着不同输入中包含的信息量不同。首先,不同模态的数据包含的信息量不同。例如,图像通常比文本包含更多的信息,并且在编码后,它们会产生更多的标记。其次,同一模态的不同输入可能包含非常不同的信息量。例如,与腕式摄像头相比,机器人的外部摄像头具有更全局的视野,包含的信息也更丰富,如图 3 的右上方所示。在这种情况下,模型可能会学习一条捷径:只关注外部视图而忽略手腕视图,从而失去感知深度的能力。为了解决这个问题,我们在编码过程中以一定的概率随机独立地屏蔽每个多模态输入,以防止模型过度依赖特定的输入。
Network Structure
- Transformer 作为可扩展的骨干网络
- 对Diffusion Transfomer (DiT)做3个改进
- QKNorm & RMSNorm
- 输入机器人物理量的数值范围不稳定,会导致梯度不稳定和数值溢出等问题,尤其是在训练大型基础模型时。
- 为了解决这个问题,我们添加了 QKNorm以避免在计算注意力时数值不稳定。
- 此外,我们还注意到,我们的问题可以被认为是一个时间序列预测任务,原始 DiT 的 LayerNorm 中的居中操作可能会导致标记转移和注意力转移,从而破坏时间序列的对称性。
- 因此,我们用RMSNorm替换LayerNorm,而无需进行居中操作。
- MLP Decoder
- 为了提高非线性机器人动作的逼近能力
- 将最终的线性解码器替换为非线性 MLP 解码器,作为从潜在空间返回物理空间的投影
- Alternating Condition Injection (ACI)
- 条件(图像和语言)高维且长度会变,难以压缩为单个令牌
- 采用交叉注意corss attention适应不同长度的条件,避免在进一步压缩中丢失信息
- 图像标记通常比文本标记多得多
- 在连续层的交叉注意力中策略性地交替注入图像和文本标记
- QKNorm & RMSNorm
data
在异构多机器人数据上训练
一个统一的action space
- 使用单个空间来容纳其本体感觉 zt 和动作 at
- 包含所有主要物理量
收集数据集
为了进行微调。
为了确保数据集的高质量
- 数量:6K+ 轨迹
- 综合性:300+ challenging tasks
- 差异性: 100+ 个不同大小和纹理的刚体和非刚体的物体,以及 15+ 个具有不同照明条件的不同房间。此外,我们进一步利用 GPT-4-Turbo (Achiam et al., 2023) 重写人工注释的指令,以增加文本多样性。
实验
想解决5个问题
- generalize
- effective
- facilitate
- capable
- helpful
看视频学的
- 前任工作里面有一个code as policies[5] 通过编程 rule base 会限制应用范围
- 另外一些工作就是end-to-end,泛化性差
- 所以本文模仿学习 end-to-end,不用rl
- [10-13]是用rl的,可以看看
没搞清楚的问题
- 预训练也是用RDT这个框架的吗?输入要求的图片也有?
- 之前视频里面说直接看代码就行,先看代码里面有没有说预训练的部分
- 太好了,找到pretrain的文档了
- 扩散模型在RDT起到的作用是
- 扩散模型做什么的:原始数据加噪声,去噪声,输出没加噪声的
- 在RDT里面,DM有两个步骤
- 加噪声:从真实action生成带噪声的动作
- 利用RDT那个网络框架得到去噪的action
- 训练目标就是最小化MSE
- DiT在这中间也没做什么
- DiT就相当于锦(扩散模型)上添花
- 重新看的时候才注意到图像处理的部分SigLip反而没那么重要……它后面还提到一个petr进行3d编码,那个才是比较重要的……那个能把多张2d图片编码成3d位置(自动驾驶里面的),不然单纯的图像景深什么的也没法做得很准确吧。如果我搞这部分我肯定会想死磕2d图片到3d位置定位这里(梦回本科学图形学。。)
流程如下
- 预训练
- 准备数据集
- 下载
- 清洗,转换格式
- 统一动作空间
- 输入:
- 图片经过SigLIP处理、指令T5-XXL处理。ACI交替注入token
- 经过噪声处理的action、proprio,映射到unified action space中,MLP。其他如控制频率等,MLP。二者连接
- 有cross-attention的DiT Block
- norm、mlp
- 损失函数,最小化MSE
- 准备数据集
- 微调
- 输入数据:目标机器人的任务指令(语言)、视觉观察(图片)、机器人状态和动作
- 训练过程和预训练一致,数据来源为自己调整的机器人(原文中使用ALOHA)
- 推理
- 输入:本体感觉proprio,相机获取的图像images,指令instruction;输出action
- 用DPM-Solver++加快训练速度
- 验证模型对相机位置的健壮性:修改相机位置、角度后收集捕获的图片作为模型输入,评价标准为能否生成正确的动作轨迹
- 说明:原来的模型提供了3个相机:场景相机、左臂、右臂。其中场景相机固定,左右臂跟随双臂移动。
- 验证已训练好的模型
- 左右相机跟随机械臂
- 修改左右相机相对机械臂的位置和视角
- 修改场景相机的位置和观察角度
- 左右臂不跟随机械臂运动,
- 对比固定左右相机的位置和角度后输出的运动轨迹与之前的差异
- 改变左右相机相对物体、机械臂的位置,比较输出的运动轨迹
- 左右相机跟随机械臂
- 修改相机位置,重新微调模型
- 改变3个相机的位置和角度(仿照4.2)
- 修改原文网络结构中和图像相关的方法
- 直接处理图像的:siglip
- 位置嵌入:PETR