跳到主要内容

Diffusion Policy 模型训练

Diffusion Policy 最初由 Columbia University 等提出，项目页为 diffusion-policy.cs.columbia.edu，官方代码在 columbia-ai-robotics/diffusion_policy。LeRobot 也提供了 PyTorch 版本实现，源码位于 LeRobot diffusion policy，并发布了 PushT 示例模型 lerobot/diffusion_pusht。

本文只覆盖 LeRobot 路径。若你需要严格复现原始论文中的 benchmark，请直接使用原始 Diffusion Policy 仓库和对应数据处理流程。

何时使用​

Diffusion Policy 将动作序列建模为条件去噪过程，适合连续、平滑、存在多种可行动作模式的控制任务。它通常比 ACT 更重，训练步数也更长；但在需要预测一段连续轨迹、动作平滑性很重要的任务中，它是合理选择。

优先考虑 Diffusion Policy 的情况：

  • 动作轨迹连续且对平滑性敏感。
  • 任务存在多种可行操作路径。
  • 数据量比 ACT 基线更充分。
  • 可以接受更长训练和较慢推理。

如果目标是先验证数据链路，建议先用 ACT 或 SmolVLA 建立基线，再训练 Diffusion Policy。

官方依据​

LeRobot v0.5.0 的公开示例包括：

  • examples/training/train_policy.py：演示如何用 DiffusionConfig 和 DiffusionPolicy 在 PushT 上训练。
  • lerobot/diffusion_pusht 模型卡：记录了 LeRobot 训练命令，核心参数为 --policy.type=diffusion、--dataset.repo_id=lerobot/pusht、--batch_size=64、--steps=200000。

按 LeRobot v0.5.0 的安装方式，训练应使用 lerobot-train：

lerobot-train \
  --output_dir=outputs/train/diffusion_pusht \
  --policy.type=diffusion \
  --dataset.repo_id=lerobot/pusht \
  --seed=100000 \
  --env.type=pusht \
  --batch_size=64 \
  --steps=200000 \
  --eval_freq=25000 \
  --save_freq=25000 \
  --wandb.enable=true

对自有机器人数据，通常不配置 --env.type，先做离线训练和 checkpoint 加载检查。

使用 LeRobot GPU 镜像​

艾欧智能已经构建并发布 ioaitech/lerobot-gpu:v0.5.0 Docker 镜像。该镜像包含 LeRobot v0.5.0、GPU 训练依赖、视频解码依赖和 lerobot-train 命令，可直接用于 Diffusion Policy 训练。中国大陆环境可改为 swr.cn-east-3.myhuaweicloud.com/ioaitech/lerobot-gpu:v0.5.0。

最小验证​

docker run --rm --gpus all --shm-size 16g \
  -v /path/to/lerobot_dataset:/data/input \
  -v /path/to/output:/outputs \
  ioaitech/lerobot-gpu:v0.5.0 \
  bash -lc 'lerobot-train \
    --policy.type=diffusion \
    --dataset.repo_id=local/my_dataset \
    --dataset.root=/data/input \
    --batch_size=8 \
    --steps=1000 \
    --output_dir=/outputs/diffusion_smoke \
    --job_name=diffusion_smoke \
    --policy.device=cuda \
    --wandb.enable=false'

正式训练模板​

docker run --rm --gpus all --shm-size 16g \
  -v /path/to/lerobot_dataset:/data/input \
  -v /path/to/output:/outputs \
  ioaitech/lerobot-gpu:v0.5.0 \
  bash -lc 'lerobot-train \
    --policy.type=diffusion \
    --dataset.repo_id=local/my_dataset \
    --dataset.root=/data/input \
    --batch_size=64 \
    --steps=100000 \
    --output_dir=/outputs/diffusion_policy \
    --job_name=diffusion_policy \
    --policy.device=cuda \
    --save_checkpoint=true \
    --save_freq=10000 \
    --wandb.enable=false'

如果训练命令提示某些 policy.* 参数不存在，先运行：

lerobot-train --help

LeRobot 的 CLI 参数随版本变化较快，复现实验时应记录镜像 tag 或 git commit。

从官方源码复现​

本地安装请先阅读 安装 LeRobot。若需要从源码复现，固定到 v0.5.0 tag：

git clone https://github.com/huggingface/lerobot.git
cd lerobot
git checkout v0.5.0

python -m venv .venv
source .venv/bin/activate
python -m pip install -U pip
pip install -e ".[all]"

本地数据训练：

lerobot-train \
  --policy.type=diffusion \
  --dataset.repo_id=local/my_dataset \
  --dataset.root=/path/to/lerobot_dataset \
  --batch_size=64 \
  --steps=100000 \
  --output_dir=outputs/train/diffusion_policy \
  --job_name=diffusion_policy \
  --policy.device=cuda \
  --wandb.enable=false

若要复现 PushT 示例，可直接使用 LeRobot 官方数据集和环境配置：

lerobot-train \
  --policy.type=diffusion \
  --dataset.repo_id=lerobot/pusht \
  --env.type=pusht \
  --batch_size=64 \
  --steps=200000 \
  --output_dir=outputs/train/diffusion_pusht \
  --job_name=diffusion_pusht \
  --policy.device=cuda

参数边界​

参数建议
--batch_size从 8 或 16 做 smoke test；正式训练再增加到显存允许的范围。
--steps扩散策略通常需要比 ACT 更长的训练；先建立 100k 量级基线。
--save_freq建议保存多个 checkpoint，便于比较不同训练阶段的动作质量。
--policy.device实际训练使用 cuda。
--wandb.enable长训练建议开启，用于观察 loss、学习率和异常中断。

高级参数如 horizon、n_action_steps、噪声调度、视觉 backbone 等应以当前 LeRobot 版本的 lerobot-train --help 和源码配置类为准，不建议照搬旧文档中的参数名。

评估建议​

Diffusion Policy 的离线 loss 不能直接代表实机成功率。建议至少比较：

  • 动作序列是否平滑，是否出现明显抖动或停顿。
  • 相同初始条件下多次 rollout 是否稳定。
  • 不同物体位置或视觉遮挡下是否仍能完成任务。
  • 推理延迟是否满足控制频率。

对于实时控制，推理步数与动作质量存在权衡。减少采样步数可以降低延迟，但可能损失动作质量，需要实机或仿真评估确认。

排错​

字段 key 不匹配​

先检查 meta/info.json 中的图像、状态和动作字段。Diffusion Policy 对时间窗口敏感，字段名和 fps 错误会直接影响监督信号。

显存不足​

降低 --batch_size，减少 DataLoader worker，并先缩短 --steps 做功能验证。不要在没有跑通 smoke test 前直接启动长训练。

与原始 Diffusion Policy 结果不一致​

LeRobot 实现、数据读取、图像预处理和评估环境都可能与原始仓库不同。若目标是复现实验论文，应使用 columbia-ai-robotics/diffusion_policy。

参考链接​