Skip to content

教程raw.githubusercontent.com6 天前

RDT-1B 数据准备、微调与真机部署实操 ​

> 本文依据以下官方文档整理:raw.githubusercontent.com/thu-ml/RoboticsDiffusionTransforme、raw.githubusercontent.com/thu-ml/RoboticsDiffusionTransforme。命令与版本以文档当前版本为准。

前置条件 ​

  • Conda;RDT 环境使用 python=3.10.0,TensorFlow Dataset 环境使用 python=3.10
  • PyTorch 示例依赖:torch==2.1.0、torchvision==0.16.0、--index-url https://download.pytorch.org/whl/cu121
  • 已下载 t5-v1_1-xxl 与 siglip-so400m-patch14-384 多模态编码器
  • 预训练 buffer 路径,文档要求 at least 400GB
  • 微调数据集磁盘目录,例如 my_cool_dataset
  • 如做文档中的真机示例,对应 ALOHA dual-arm robot 场景

1. 1. 克隆仓库并安装 RDT 环境 ​

先在本机克隆 RoboticsDiffusionTransformer,创建名为 rdt 的 Conda 环境,并安装 PyTorch、packaging、flash-attn 与 requirements.txt 中依赖。

bash
## Clone this repo
git clone git@github.com:thu-ml/RoboticsDiffusionTransformer.git
cd RoboticsDiffusionTransformer

## Create a Conda environment
conda create -n rdt python=3.10.0
conda activate rdt

## Install pytorch
## Look up https://pytorch.org/get-started/previous-versions/ with your cuda version for a correct command
pip install torch==2.1.0 torchvision==0.16.0  --index-url https://download.pytorch.org/whl/cu121

## Install packaging
pip install packaging==24.0

## Install flash-attn
pip install flash-attn --no-build-isolation

## Install other prequisites
pip install -r requirements.txt

2. 2. 链接多模态编码器并填写 buffer ​

下载 t5-v1_1-xxl 与 siglip 后,在仓库根目录创建 google 子目录并建立软链接。随后在 configs/base.yaml#L22 填写 dataset 下的 buf_path。

bash
## Under the root directory of this repo
mkdir -p google

## Link the downloaded encoders to this repo
ln -s /path/to/t5-v1_1-xxl google/t5-v1_1-xxl
ln -s /path/to/siglip-so400m-patch14-384 google/siglip-so400m-patch14-384

3. 3. 安装 TensorFlow Dataset 环境 ​

若要走预训练数据管线,需要在另一个 Conda 环境中安装 TensorFlow Dataset 相关依赖。文档给出 requirements_data.txt 与手动安装两种方式。

bash
## Under the root directory of this repo
conda create -n rdt-data python=3.10
conda activate rdt-data

## Install all the prequisites
pip install -r requirements_data.txt
## Or you can manually install each package (please refer to requirements_data.txt for specific versions) 
pip install tfds-nightly gsutil tensorflow Pillow pyyaml opencv-python tensorflow-graphics imageio[ffmpeg]
## If the speed is too slow, you can specify alternative sources (tfds-nightly is not available in Tsinghua mirror)
pip install -i https://pypi.tuna.tsinghua.edu.cn/simple gsutil tensorflow Pillow pyyaml opencv-python tensorflow-graphics imageio[ffmpeg]

4. 4. 连接数据集根目录并下载 Open X-Embodiment ​

预训练或复用预训练管线前,先把磁盘上的 dataset 目录链接到仓库 data/datasets。Open X-Embodiment 还需在 ../data/openx_embod/download.sh#L72 指定 gsutil 路径。

bash
ln -s /path/to/dataset /path/to/repo/RoboticsDiffusionTransformer/data/datasets
bash
## Under the root directory of this repo
cd data/openx_embod
## Download all datasets
bash download_openx_embod.sh

5. 5. 准备其他预训练数据集 ​

文档还给出 Mobile ALOHA、Bridgev2、Calvin、RH20T 与 RoboSet 的处理命令。按需要的数据集进入对应目录执行即可。

bash
cd data/aloha
cat tfrecord_part_a* > tfrecord_part_a.zip
unzip tfrecord_part_a.zip
bash
cd data/bridgev2
## Download and preprocess the dataset
sh download.sh
bash
cd data/calvin
## Download and preprocess the dataset
sh download.sh
## Convert the dataset to TFRecord format
python hdf5totfrecords.py
bash
cd data/rh20t
## Convert the dataset to TFRecord
python hdf5totfrecords.py
bash
cd data/roboset
## Download and preprocess the dataset
sh download.sh

6. 6. 链接自己的微调数据集 ​

如果微调数据在本地磁盘,文档示例命名为 my_cool_dataset。进入 data 目录创建 datasets,并把该数据集软链接到仓库中。

bash
## Under the root directory of this repo
cd data
mkdir -p datasets

## Link the downloaded dataset to this repo
ln -s /path/to/my_cool_dataset datasets/my_cool_dataset

7. 7. 注册微调数据并实现加载器 ​

在 configs/dataset_control_freq.json 添加 my_cool_dataset 的控制频率;在 configs/finetune_datasets.json 与 configs/finetune_sample_weights.json 中用 my_cool_dataset 替换 agilex。然后改 data/hdf5_vla_dataset.py:设置 HDF5_DIR、DATASET_NAME,并实现 parse_hdf5_file() 与 parse_hdf5_file_state_only()。

8. 8. 确认真机部署示例入口 ​

仓库说明中给出的真机部署示例是 scripts/agilex_inference.py。文档还说明 RDT 已在 ALOHA dual-arm robot 上完成 fine-tuned 与 deployed;本段未给出额外运行命令。

验证是否跑通 ​

确认 rdt 与 rdt-data 环境依赖安装完成;google/t5-v1_1-xxl、google/siglip-so400m-patch14-384 与 data/datasets 软链接存在;configs/base.yaml#L22 已填写 buf_path;所选数据集下载、解压或 TFRecord 转换命令无报错;configs/finetune_datasets.json、configs/finetune_sample_weights.json、configs/dataset_control_freq.json 已包含 my_cool_dataset;data/hdf5_vla_dataset.py 的加载函数可按你的数据结构返回所需字段。

文档里点名的注意事项 ​

  • configs/base.yaml#L22 的 buffer will only be used during pre-training,且文档要求 at least 400GB。
  • 微调集如果属于 Open X-Embodiment 或预训练数据集合,可以通过 pre-trained pipeline 微调,但需要移除参数中其他 redundant datasets。
  • HDF5VLADataset 名称不代表必须使用 HDF5,只要 class is correctly implemented。
  • 如果机器人是 single-arm,动作要填入 unified action vector 的 right-arm 部分,以对齐预训练数据集。
  • RDT 对 EEF rotation 使用 6D representation。
  • 新数据集预训练管线中,branch 或 loop operations should only use TensorFlow functions,例如 use tf.cond instead of if。
  • configs/dataset_img_keys.json 中要按 exterior-cam, right-wrist-cam, left-wrist-cam, any-cam 顺序指定四张图像。

来源

本页事实对应的公开页面,以原文为准。

  1. raw.githubusercontent.com/thu-ml/RoboticsDiffusionTransformer/main/README.md访问 2026-10-22
  2. raw.githubusercontent.com/thu-ml/RoboticsDiffusionTransformer/main/docs/pretrain.md访问 2026-10-22