许可证: mit
语言:

• 英文
  任务标签: 机器人学
  库名称: transformers
  标签:
• 机器人学
• pytorch
• 多模态
• 预训练
• 视觉语言动作
• 扩散模型
• RDT

RDT-170M

RDT-170M 是一个拥有1.7亿参数的模仿学习扩散Transformer模型（消融实验中的RDT(small)）。其隐藏层大小为1024，深度为14，均为RDT-1B模型的一半。在给定语言指令和最多三个视角的RGB图像输入时，RDT能够预测未来64个机器人动作。该模型兼容几乎所有现代移动机械臂，包括单臂/双臂、关节空间/末端执行器空间、位置控制/速度控制，甚至轮式移动平台。

所有代码、预训练模型权重及数据均采用MIT许可证授权。

更多信息请访问我们的项目主页和论文。

模型详情

• 开发团队: 清华大学TSAIL课题组的RDT研究团队
• 任务类型: 视觉-语言-动作（语言+图像→机器人动作）
• 模型类型: 基于Transformer的扩散策略
• 许可证: MIT
• 支持语言: 英文
• 多模态编码器:
  • 视觉骨干网络: siglip-so400m-patch14-384
  • 语言模型: t5-v1_1-xxl
• 预训练数据集: 包含RT-1数据集、RH20T、DROID、BridgeData V2、RoboSet及Open X-Embodiment子集在内的46个数据集，完整列表见此处。
• 代码仓库: https://github.com/thu-ml/RoboticsDiffusionTransformer
• 论文: https://arxiv.org/pdf/2410.07864
• 项目主页: https://rdt-robotics.github.io/rdt-robotics/

使用说明

RDT接收语言指令、RGB图像（最多三个视角）、控制频率（如有）和本体感知作为输入，预测未来64个机器人动作。通过统一动作空间的支持，RDT可控制几乎所有机械臂平台，该空间包含机械臂所有主要物理量（如末端执行器/关节空间、位置/速度控制、轮式移动等）。部署时需将原始动作向量映射到统一空间向量，具体操作请参考代码仓库。

使用限制: 由于本体差异，RDT暂不能泛化至预训练数据未涵盖的新型机器人平台。建议针对目标机器人收集小规模数据集进行微调，教程详见代码仓库。

以下是在机器人上使用RDT-1B模型进行推理的示例：

# 请先克隆代码库并安装依赖
# 通过"cd RoboticsDiffusionTransformer"进入代码库根目录

# 从代码库导入创建函数
from scripts.agilex_model import create_model

# 视觉输入使用的摄像头名称
CAMERA_NAMES = ['cam_high', 'cam_right_wrist', 'cam_left_wrist']
config = {
    'episode_len': 1000,  # 单次任务最大长度
    'state_dim': 14,      # 机器人状态维度
    'chunk_size': 64,     # 单步预测动作数
    'camera_names': CAMERA_NAMES,
}
pretrained_vision_encoder_name_or_path = "google/siglip-so400m-patch14-384" 
# 根据配置创建模型
model = create_model(
    args=config,
    dtype=torch.bfloat16, 
    pretrained_vision_encoder_name_or_path=pretrained_vision_encoder_name_or_path,
    pretrained='robotics-diffusion-transformer/rdt-1b',
    control_frequency=25,
)

# 开始推理流程
# 加载预计算的语言嵌入
# 语言指令编码方法参见scripts/encode_lang.py
lang_embeddings_path = 'your/language/embedding/path'
text_embedding = torch.load(lang_embeddings_path)['embeddings']  
images: List(PIL.Image) = ... # 最近2帧图像
proprio = ... # 当前机器人状态
# 推理预测未来chunk_size个动作
actions = policy.step(
    proprio=proprio,
    images=images,
    text_embeds=text_embedding 
)

引用

如果我们的工作对您有帮助，请引用：

@article{liu2024rdt,
  title={RDT-1B: a Diffusion Foundation Model for Bimanual Manipulation},
  author={Liu, Songming and Wu, Lingxuan and Li, Bangguo and Tan, Hengkai and Chen, Huayu and Wang, Zhengyi and Xu, Ke and Su, Hang and Zhu, Jun},
  journal={arXiv preprint arXiv:2410.07864},
  year={2024}
}

感谢支持！