RDT-1B: A DIFFUSION FOUNDATION MODEL FOR BIMANUAL MANIPULATION

RDT-1B: A DIFFUSION FOUNDATION MODEL FOR BIMANUAL MANIPULATION。双臂机器人操作,扩散模型

abstract

  1. 解决coordination和traning data lack问题
  2. Robotics Diffusion Transformer
  3. 基于扩散基础模型,bimanual manipulation
  4. diffusion model+Transformer to deal with multi-modal
  5. for data scarcity: Physically Interpretable Unified Action Space
  6. pre-train on large and scale up
  7. fine-tune on self multi-task dataset
  8. experiments on real robot show outperform,good generalization

introduction

  1. 问题:narrow generalization and failing in complex tasks
  2. 解决:“imitation learning on large-scale datasets” (选中模仿学习)
  3. 问题:数据缺少
  4. 解决:首先对广泛的多机器人数据集进行预训练,然后对目标双臂机器人上收集的小数据集进行微调。
  5. 有两个难点
    1. 需要一个架构,有表现力和可扩展性
    2. 数据异构
  6. RDT
    1. 最大的双手操控基础模型,通用性强,用“Diffusion Transformers (DiTs)”作为骨干网络
    2. 设计思路
      1. expressiveness:利用扩散模型的能力来建模复杂的分布
      2. scalability:利用 Transformer 主干并精心设计多模式编码
      3. characterize:对原有DiT结构进行重要修改,包括 MLP 解码、改进的归一化和交替条件注入
      4. heterogeneous:提出物理上可解释的统一行动空间
  7. 实验
    1. 在最大集合上预训练 RDT 模型
    2. 在自收集的多任务双手数据集上微调 RDT

PROBLEM FORMULATION AND CHALLENGES

  1. 选择
    1. ALOHA双臂机器人
    2. 语言条件下的视觉双手操作
  2. instruction,observation,time,action,RGB observation,low-dimensional proprioception of the robot,control frequency的设置
    1. low-dimensional proprioception of the robot 不知道是什么
  3. task:给定一项新任务,需要一种实用的策略来泛化到训练数据中未出现的未见过的元素。
  4. policy:预训练和微调管道
  5. 开发这样的“foundation model” 有俩challenge(introduction提到的)
    1. architecture
      1. expressiveness
      2. scalability
    2. heterogeneous data

RDT

整体的框架见论文图3

RDT MODEL

为了解决第一个challenge。

Diffusion Modeling

将扩散模型应用于机器人任务的2个挑战

  1. 机器人物理量的特点是其非线性动力学和物理相互作用(如碰撞、约束和阻尼等材料属性)可能引起的高频变化
  2. 不稳定的数值范围,可能是传感器导致的

Encoding of Heterogeneous Multi-Modal Inputs

这个异构是SD带来的,和第二个challenge不一样。

为了将这些不同的模态编码到一个统一的潜在空间中,用以下编码方法

  1. 针对表示物理量的低维输入,MLP(具有傅里叶特征可以有效地捕捉低维空间中的高频变化
  2. 图像输入是高维的,包含丰富的空间和语义信息,用图像-文本对齐的预训练视觉编码器 SigLIP
  3. 语言输入的长度各不相同,并且高度抽象,用预先训练的基于 Transformer 的语言模型 T5-XXL

除了结构之外,异构性还意味着不同输入中包含的信息量不同。首先,不同模态的数据包含的信息量不同。例如,图像通常比文本包含更多的信息,并且在编码后,它们会产生更多的标记。其次,同一模态的不同输入可能包含非常不同的信息量。例如,与腕式摄像头相比,机器人的外部摄像头具有更全局的视野,包含的信息也更丰富,如图 3 的右上方所示。在这种情况下,模型可能会学习一条捷径:只关注外部视图而忽略手腕视图,从而失去感知深度的能力。为了解决这个问题,我们在编码过程中以一定的概率随机独立地屏蔽每个多模态输入,以防止模型过度依赖特定的输入。

Network Structure

  1. Transformer 作为可扩展的骨干网络
  2. 对Diffusion Transfomer (DiT)做3个改进
    1. QKNorm & RMSNorm
      1. 输入机器人物理量的数值范围不稳定,会导致梯度不稳定和数值溢出等问题,尤其是在训练大型基础模型时。
      2. 为了解决这个问题,我们添加了 QKNorm以避免在计算注意力时数值不稳定。
      3. 此外,我们还注意到,我们的问题可以被认为是一个时间序列预测任务,原始 DiT 的 LayerNorm 中的居中操作可能会导致标记转移和注意力转移,从而破坏时间序列的对称性。
      4. 因此,我们用RMSNorm替换LayerNorm,而无需进行居中操作。
    2. MLP Decoder
      1. 为了提高非线性机器人动作的逼近能力
      2. 将最终的线性解码器替换为非线性 MLP 解码器,作为从潜在空间返回物理空间的投影
    3. Alternating Condition Injection (ACI)
      1. 条件(图像和语言)高维且长度会变,难以压缩为单个令牌
      2. 采用交叉注意corss attention适应不同长度的条件,避免在进一步压缩中丢失信息
      3. 图像标记通常比文本标记多得多
      4. 在连续层的交叉注意力中策略性地交替注入图像和文本标记

data

在异构多机器人数据上训练

一个统一的action space

  1. 使用单个空间来容纳其本体感觉 zt 和动作 at
  2. 包含所有主要物理量

收集数据集

为了进行微调。

为了确保数据集的高质量

  1. 数量:6K+ 轨迹
  2. 综合性:300+ challenging tasks
  3. 差异性: 100+ 个不同大小和纹理的刚体和非刚体的物体,以及 15+ 个具有不同照明条件的不同房间。此外,我们进一步利用 GPT-4-Turbo (Achiam et al., 2023) 重写人工注释的指令,以增加文本多样性。

实验

想解决5个问题

  1. generalize
  2. effective
  3. facilitate
  4. capable
  5. helpful

看视频学的

  1. 前任工作里面有一个code as policies[5] 通过编程 rule base 会限制应用范围
  2. 另外一些工作就是end-to-end,泛化性差
  3. 所以本文模仿学习 end-to-end,不用rl
  4. [10-13]是用rl的,可以看看

没搞清楚的问题

  1. 预训练也是用RDT这个框架的吗?输入要求的图片也有?
    1. 之前视频里面说直接看代码就行,先看代码里面有没有说预训练的部分
    2. 太好了,找到pretrain的文档了
  2. 扩散模型在RDT起到的作用是
    1. 扩散模型做什么的:原始数据加噪声,去噪声,输出没加噪声的
    2. 在RDT里面,DM有两个步骤
      1. 加噪声:从真实action生成带噪声的动作
      2. 利用RDT那个网络框架得到去噪的action
      3. 训练目标就是最小化MSE
  3. DiT在这中间也没做什么
    1. DiT就相当于锦(扩散模型)上添花
  4. 重新看的时候才注意到图像处理的部分SigLip反而没那么重要……它后面还提到一个petr进行3d编码,那个才是比较重要的……那个能把多张2d图片编码成3d位置(自动驾驶里面的),不然单纯的图像景深什么的也没法做得很准确吧。如果我搞这部分我肯定会想死磕2d图片到3d位置定位这里(梦回本科学图形学。。)

流程如下

  1. 预训练
    1. 准备数据集
      1. 下载
      2. 清洗,转换格式
      3. 统一动作空间
    2. 输入:
      1. 图片经过SigLIP处理、指令T5-XXL处理。ACI交替注入token
      2. 经过噪声处理的action、proprio,映射到unified action space中,MLP。其他如控制频率等,MLP。二者连接
    3. 有cross-attention的DiT Block
    4. norm、mlp
    5. 损失函数,最小化MSE
  2. 微调
    1. 输入数据:目标机器人的任务指令(语言)、视觉观察(图片)、机器人状态和动作
    2. 训练过程和预训练一致,数据来源为自己调整的机器人(原文中使用ALOHA)
  3. 推理
    1. 输入:本体感觉proprio,相机获取的图像images,指令instruction;输出action
    2. 用DPM-Solver++加快训练速度
  4. 验证模型对相机位置的健壮性:修改相机位置、角度后收集捕获的图片作为模型输入,评价标准为能否生成正确的动作轨迹
    1. 说明:原来的模型提供了3个相机:场景相机、左臂、右臂。其中场景相机固定,左右臂跟随双臂移动。
    2. 验证已训练好的模型
      1. 左右相机跟随机械臂
        1. 修改左右相机相对机械臂的位置和视角
        2. 修改场景相机的位置和观察角度
      2. 左右臂不跟随机械臂运动,
        1. 对比固定左右相机的位置和角度后输出的运动轨迹与之前的差异
        2. 改变左右相机相对物体、机械臂的位置,比较输出的运动轨迹
    3. 修改相机位置,重新微调模型
      1. 改变3个相机的位置和角度(仿照4.2)
      2. 修改原文网络结构中和图像相关的方法
        1. 直接处理图像的:siglip
        2. 位置嵌入:PETR