企业官网建设流程全解析

openpi 在 DROID 数据集上训练 pi0.5：全量 RLDS 微调与自定义数据微调实战指南

【免费下载链接】openpi项目地址: https://gitcode.com/GitHub_Trending/op/openpi

本文是 openpi 项目中针对DROID 数据集训练的完整实战指南，覆盖两条核心路线：其一，使用 RLDS 数据格式在完整DROID 数据集（约 75k 条演示）上复现 pi0.5-DROID 训练管线；其二，将自定义的小规模 DROID 平台采集数据转换为 LeRobot 格式后进行专家模型微调。读完本文，你将掌握从环境安装、数据集下载、归一化统计计算、启动训练到数据过滤（idle filter）的全流程操作，并能结合仓库源码理解 RLDS 数据加载器、动作空间选择与采样索引过滤的底层实现。

本文以 examples/droid/README_train.md 为骨架展开，所有命令均基于当前仓库实际可用的配置与脚本，训练完成后的部署与真机运行请继续参考 examples/droid/README.md。

背景：为什么全量 DROID 训练需要切换 RLDS

openpi 的大部分训练流程（如 Aloha、Libero 示例）默认使用LeRobot加载数据。但对于 DROID 这类超大规模真实机器人数据集，当前版本的 LeRobot 数据加载器在可扩展性上尚不足以支撑训练（官方注释也指出其团队正在改进）。因此，仓库为 DROID 全量训练专门实现了基于RLDS（Relational Learning Description Schema，即 TFDS 使用的数据格式）的数据加载器，其核心实现位于 droid_rlds_dataset.py，并在 data_loader.py 中通过create_rlds_dataset接入统一训练入口——该函数目前只支持 DROID 数据集。

第一步：安装 RLDS 依赖组（强制 Python 3.11）

RLDS 依赖组以tensorflow-cpu==2.15.0为核心，而该版本在 PyPI 上只发布 Python 3.11 的 wheel（见 pyproject.toml 中的[dependency-groups] rlds注释）。因此，在创建虚拟环境之前，必须先确保使用 Python 3.11，否则uv sync会因找不到对应 wheel 而失败：

uv python install 3.11 uv venv --python 3.11 uv sync --group rlds

rlds依赖组共包含三个包：

• dlimp：由 DLataset 提供的 RLDS → tf.data 的高效包装层（仓库通过 git 依赖固定到特定 revision，见 pyproject.toml）；
• tensorflow-cpu==2.15.0：仅用于 CPU 数据管线；
• tensorflow-datasets==4.9.9：用于构建 TFDS dataset builder。

从源码看，droid_rlds_dataset.py 在初始化时会将 TensorFlow 配置为不注册任何 GPU 设备（tf.config.set_visible_devices([], "GPU")），避免与 JAX/PyTorch 抢占显存。同时，RLDS 数据加载要求TrainConfig.num_workers=0，因为多进程数据加载由 TF 数据管线内部自行处理（见 config.py 的注释）。

第二步：下载 DROID 数据集（1.0.1 版本，1.8TB）

安装gsutil（Google Cloud CLI）后，使用以下命令下载 DROID 数据集：

gsutil -m cp -r gs://gresearch/robotics/droid/1.0.1 <your_download_path>/droid/1.0.1

务必下载 1.0.1 版本，原因如下：

• 1.0.1 包含完整的语言标注（约75k 条演示都有指令标注）；
• 1.0.0 仅有 3 万条演示的标注。

如果你确实需要其他版本，可以修改 droid_rlds_dataset.py 中RLDSDataset的version字段，或直接修改 config.py 中RLDSDroidDataConfig.datasets的默认定义（当前默认即version="1.0.1"、weight=1.0）。

下载该数据集需要约1.8TB磁盘空间，请提前规划存储。gsutil -m参数启用并行拷贝，可显著缩短下载时间。

第三步：配置训练路径并启动全量训练

3.1 修改rlds_data_dir

在 config.py 中找到你使用的训练配置（全量 DROID 对应pi05_full_droid_finetune），将RLDSDroidDataConfig中的rlds_data_dir改为你下载 DROID 数据的目录（即droid目录的父目录）。需要说明的是，仓库中该字段示例值指向/mnt/pi-data/kevin这类内部路径，实操时必须替换为你的本地路径。

3.2 计算归一化统计量（约 10 分钟）

uv run --group rlds scripts/compute_norm_stats.py --config-name pi05_full_droid_finetune --max-frames 10_000_000

该脚本（compute_norm_stats.py）会：

1. 按配置创建 RLDS 数据加载器（当data_config.rlds_data_dir非空时走create_rlds_dataloader分支）；
2. 流式遍历数据，对"state"与"actions"两个键分别用normalize.RunningStats在线累积均值/方差统计量（由于是增量更新，不要求一次性加载全部数据，这也是--max-frames 10_000_000只采样前 1000 万帧即可的原因）；
3. 将统计结果写入assets/<config_name>/<repo_id>目录，供后续训练时加载。

注意：计算过程中会移除字符串字段（如prompt），因为字符串无法被 JAX 处理且对归一化统计无用（见脚本中的RemoveStringstransform）。

3.3 启动训练

XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 uv run --group rlds scripts/train.py pi05_full_droid_finetune --exp-name=my_experiment --overwrite

各参数含义：

• XLA_PYTHON_CLIENT_MEM_FRACTION=0.9：限制 JAX 预分配显存比例为 90%，为 TensorFlow 数据管线和系统预留余量；
• pi05_full_droid_finetune：训练配置名（在 config.py 中定义）；
• --exp-name=my_experiment：实验名，用于命名 checkpoint 与 wandb 日志目录；
• --overwrite：若 checkpoint 目录已存在则覆盖（与--resume互斥，二者同时设置会抛出ValueError，见 config.py）。

pi05_full_droid_finetune配置的关键超参数如下（完整定义见 config.py）：

3.4 关于动作空间的重要说明

原文档明确指出：原始 pi0.5-DROID 模型使用关节速度动作训练，但关节速度动作与仿真评估环境不兼容（速度难以在仿真中精确复现）。因此，openpi 的训练配方默认改用关节位置动作。

这一设计在源码中有三处体现：

1. droid_rlds_dataset.py 定义了DroidActionSpace枚举：JOINT_POSITION与JOINT_VELOCITY，数据加载器默认取JOINT_POSITION；
2. 数据管线中，joint_position + gripper_position被拼接为动作（droid_rlds_dataset.py）；
3. 由于加载器产出的是绝对位置动作，而 pi0 系列模型期望的是相对首帧状态的 delta 动作，RLDSDroidDataConfig在动作空间为JOINT_POSITION时会自动追加DeltaActions(delta_action_mask)变换（对前 7 维做 delta，夹爪保持绝对值），见 config.py。

计算资源需求

根据仓库训练配方的注释（config.py）：

• 从 pi0 初始化开始训练：约2 天 × 8× H100即可收敛（100k 迭代、batch size 256，约 1 个 epoch）；
• 若从PaliGemma初始化：约5 天 × 8× H100（240k 迭代，即 3 个 epoch）。

仓库还提到曾尝试用LoRA做低成本微调，但目前为止未发现 LoRA 策略能达到足够好的性能，因此全量 DROID 训练默认采用全参数微调。

数据过滤：处理 DROID 中的 idle 时间步

与任何多样的真实机器人数据集一样，DROID 数据并不"干净"——其中包含大量机器人静止不动的idle时间步（部分原因是数据采集使用 VR 遥操作界面）。对这些 idle 过渡帧做合理过滤，可以显著提升策略性能。openpi 的 DROID 训练配方默认实现了与所有 pi-DROID 模型一致的 idle 过滤器，其原理是预先计算训练时应采样的数据索引。

过滤脚本的工作方式

索引列表由 compute_droid_nonidle_ranges.py 生成，核心逻辑（L36-L38）包含三个可调阈值：

• min_idle_len = 7：连续 idle 帧数超过 7 则整段剔除（多数 DROID 动作分块策略每块会执行前 8 个动作，这样过滤可避免策略"卡在"静止输出上）；
• min_non_idle_len = 16：非 idle 连续段少于 16 帧（约 1 秒）则整段剔除；
• filter_last_n_in_ranges = 10：每个保留段末尾再裁掉 10 帧（这些帧对应的动作块中含有大量 idle 动作）。

idle 判定依据是相邻两帧关节速度差的绝对值均小于1e-3（compute_droid_nonidle_ranges.py）。脚本输出一个 JSON 文件，将每个 episode 的唯一 ID（recording_folderpath--file_path）映射到需要保留的[start, end)帧区间列表，并支持断点续跑（每处理 1000 个 episode 自动写盘一次）。

过滤器在训练时的应用

训练时，DroidRldsDataset会加载过滤器字典（filter_dict_path指向gs://openpi-assets/droid/droid_sample_ranges_v1_0_1.json，见 config.py），构建一个tf.lookup.StaticHashTable（droid_rlds_dataset.py）。每个时间步通过拼接recording_folderpath--file_path--<step_index>生成唯一step_id（droid_rlds_dataset.py），查表命中则保留、未命中则丢弃。

如果你想修改 idle 过滤器或实现自定义采样逻辑，可以：

1. 修改 compute_droid_nonidle_ranges.py 重新生成索引列表；
2. 通过RLDSDroidDataConfig中的filter_dict_path="<path_to_filter_dict>"参数指向新文件（在 config.py 中配置）。

注意：仓库预生成的过滤索引列表仅对droid/1.0.1数据集有效，对其他版本不适用。如果你使用自定义 DROID 版本，必须重新运行上述脚本生成新的采样索引。

RoboArena：真机评测你的策略

训练完成后，可以把自己的 DROID 策略提交到RoboArena 基准，该基准允许在真实世界的多样任务与场景中评估策略表现。评测流程与基线策略部署命令见 examples/droid/README.md，其中也提供了从 pi0、pi0-FAST、PaliGemma 等不同初始化训练的 RoboArena 基线策略的部署方式（对应推理配置在 roboarena_config.py）。

自定义 DROID 数据集微调（LeRobot 路线）

如果你在 DROID 平台上采集了较小的自定义数据集（一般 < 数十小时），推荐走 LeRobot 转换路线：先把数据转为 LeRobot 格式，再对 pi0.5-DROID 做微调。数据量大时则回到上文介绍的 RLDS 路线（RLDS 效率更高）。

Step 1：将自定义 DROID 数据集转换为 LeRobot 格式

本示例使用真实 DROID 数据的一个小子集（仅 30 条演示，1.6GB）来演示流程——实操时请替换为你自己的数据集：

gsutil -m cp -r gs://gresearch/robotics/droid_raw/1.0.1/IRIS/success/2023-12-04 <your_target_path>

接着下载 DROID 语言标注（12MB），以便将演示与语言指令配对（使用自有数据时可手动录入指令，无需下载）：

gsutil -m cp -r gs://gresearch/robotics/droid_raw/1.0.1/aggregated-annotations-030724.json <your_target_dir>

前提条件：每个 episode 的目录中必须包含recordings/MP4文件夹（存放 MP4 视频文件）。如果只有 SVO 原始视频，需要先执行 DROID 官方的svo_to_mp4.py脚本完成 MP4 视频提取。

然后运行转换脚本（30 条演示约 <5 分钟）：

uv run examples/droid/convert_droid_data_to_lerobot.py --data_dir <your_target_path>

该脚本（convert_droid_data_to_lerobot.py）的关键行为：

• 使用LeRobotDataset.create创建数据集，robot_type="panda"、fps=15（DROID 数据通常按 15fps 记录）；
• 定义的 feature 与 DROID RLDS 命名约定一致：三路图像（exterior_image_1_left、exterior_image_2_left、wrist_image_left，分辨率 180×320×3）、joint_position（7 维）、gripper_position（1 维）；
• 动作采用 8 维 = 7 维关节速度 + 1 维夹爪位置（与 pi05-droid 预训练使用的动作空间一致，注意与 RLDS 全量路线的关节位置动作不同，因此LeRobotDROIDDataConfig不再追加 delta 变换，见 config.py）；
• 图像从 BGR 翻转为 RGB 并缩放到 (320, 180)（L128-L134）；
• 从aggregated-annotations-030724.json中按 episode id 匹配语言指令；
• 结果保存到$LEROBOT_HOME目录；如需上传到 Hugging Face Hub，可追加--push_to_hub参数。

Step 2：使用自定义数据集启动微调

仓库提供了现成的微调配置pi05_droid_finetune（config.py），其要点包括：

• model=Pi0Config(pi05=True, action_dim=32, action_horizon=16)：32 维动作是 pi0.5 的固定维度；
• data=LeRobotDROIDDataConfig(repo_id="your_hf_username/my_droid_dataset", ...)：把repo_id替换为你转换得到的数据集 ID；
• assets复用 pi05_droid 预训练 checkpoints 中保存的DROID 原始归一化统计量（gs://openpi-assets/checkpoints/pi05_droid/assets），微调时不应重新计算；
• weight_loader从gs://openpi-assets/checkpoints/pi05_droid/params加载 pi0.5-DROID 基座权重；
• 训练规模：num_train_steps=20_000、batch_size=32。

如需换用其他基座模型，只需修改上述配置中的模型与权重加载器部分（可在 config.py 中搜索pi05_droid_finetune附近的其他 DROID 配置作为参考）。

启动微调：

uv run scripts/train.py pi05_droid_finetune --exp-name=my_experiment --overwrite

训练完成后，按照 examples/droid/README.md 中的步骤启动策略服务器（如uv run scripts/serve_policy.py --env=DROID）并在 DROID 真机上运行推理。

附：RLDS 数据加载器内部机制速览

理解 droid_rlds_dataset.py 有助于你调试自定义过滤逻辑，其流水线依次执行：

1. 只保留成功轨迹：按文件路径正则匹配.*success.*（L72-L77）；
2. 重复数据集：dataset.repeat()保证训练永不出界；
3. 重结构化：拼接动作、随机二选一取两个外部相机中的左侧画面（仅训练单路外部相机）、从 3 条语言指令中随机采样 1 条（L115-L169）；
4. 动作分块：将轨迹切分为长度为action_chunk_size（默认 16）的动作块，末尾不足时重复最后一个动作（对绝对位置动作是合理行为，L173-L194）；
5. 扁平化 + 过滤：展开为单帧数据集并按passes_filter查表过滤；
6. 延迟解码图像：RLDS 中图像以编码形式存储，训练时才解码以节省内存（L212-L222）；
7. 混合多数据集：sample_from_datasets按weight加权采样（权重之和必须为 1.0），随后做 25 万规模的 shuffle 缓冲（L232-L236）。

值得注意的两个数值：默认shuffle_buffer_size=250_000（若内存紧张可调小，但低于约 10 万会显著削弱打乱随机性）；DroidRldsDataset.__len__硬编码返回 2000 万，代表过滤后 DROID 的近似样本数（L245-L248）。

小结

围绕 DROID 数据集，openpi 提供了两条互补的训练路径：RLDS 全量训练面向大规模数据，需要 Python 3.11 环境、1.8TB 磁盘与 H100 级别的算力，并通过预计算的 idle 过滤索引显著提升策略质量；LeRobot 自定义微调面向数十小时以内的小规模数据，一条命令完成格式转换后即可在消费级配置上微调 pi0.5-DROID 专家策略。无论选择哪条路径，都建议严格使用 1.0.1 版本数据，并优先采用关节位置动作以保持与仿真评估环境的兼容性。训练产出模型后的推理与真机部署细节，请继续阅读 examples/droid/README.md。

【免费下载链接】openpi项目地址: https://gitcode.com/GitHub_Trending/op/openpi