RDT-1B

RDT-1B 是一个拥有 10 亿参数的模仿学习扩散Transformer模型，它在超过 100 万个多机器人 episodes 上进行了预训练。给定语言指令和最多三个视角的 RGB 图像，RDT 能够预测接下来的 64 个机器人动作。RDT 几乎兼容所有现代移动机械臂，从单臂到双臂、关节空间到末端执行器空间、位置控制到速度控制，甚至包括轮式移动。

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

更多信息请参阅我们的项目页面和论文。

模型详情

开发团队： RDT 团队，成员包括来自清华大学 TSAIL 小组的研究人员
任务类型： 视觉-语言-动作（语言、图像 => 机器人动作）
模型类型： 基于 Transformer 的扩散策略
许可证： MIT
语言（自然语言处理）： en
多模态编码器：
- 视觉主干网络： siglip-so400m-patch14-384
- 语言模型： t5-v1_1-xxl
预训练数据集： 46 个数据集，包括 RT-1 Dataset、RH20T、DROID、BridgeData V2、RoboSet 以及 Open X-Embodiment 的一个子集。详细列表请参见此链接。
代码仓库： https://github.com/thu-ml/RoboticsDiffusionTransformer
论文： https://arxiv.org/pdf/2410.07864
项目页面： https://rdt-robotics.github.io/rdt-robotics/

用途

RDT 接收语言指令、RGB 图像（最多三个视角）、控制频率（如有）和本体感觉作为输入，并预测接下来的 64 个机器人动作。

借助统一的动作空间，RDT 支持几乎所有机器人 manipulator 的控制，该空间涵盖了机器人 manipulator 的所有主要物理量（例如，末端执行器和关节、位置和速度，以及轮式移动）。要在您的机器人平台上部署，您需要将原始动作向量的相关量填充到统一空间向量中。更多信息请参见 our repository。

超出范围：由于具身差距，RDT 目前尚不能泛化到新的机器人平台（预训练数据集中未出现的平台）。在这种情况下，我们建议收集目标机器人的小型数据集，然后使用该数据集对 RDT 进行 fine-tune。教程请参见 our repository。

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

# Please first clone the repository and install dependencies
# Then switch to the root directory of the repository by "cd RoboticsDiffusionTransformer"

# Import a create function from the code base
from scripts.agilex_model import create_model

# Names of cameras used for visual input
CAMERA_NAMES = ['cam_high', 'cam_right_wrist', 'cam_left_wrist']
config = {
    'episode_len': 1000,  # Max length of one episode
    'state_dim': 14,      # Dimension of the robot's state
    'chunk_size': 64,     # Number of actions to predict in one step
    'camera_names': CAMERA_NAMES,
}
pretrained_vision_encoder_name_or_path = "google/siglip-so400m-patch14-384" 
# Create the model with the specified configuration
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,
)

# Start inference process
# Load the pre-computed language embeddings
# Refer to scripts/encode_lang.py for how to encode the language instruction
lang_embeddings_path = 'your/language/embedding/path'
text_embedding = torch.load(lang_embeddings_path)['embeddings']  
images: List(PIL.Image) = ... #  The images from last 2 frames
proprio = ... # The current robot state
# Perform inference to predict the next `chunk_size` actions
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}
}

谢谢！