Diffusion Policy の学習（LeRobot）

Diffusion Policy は Columbia 大学等による手法で、diffusion-policy.cs.columbia.edu と columbia-ai-robotics/diffusion_policy が参照実装です。LeRobot には PyTorch 実装 LeRobot diffusion policy と PushT の例 lerobot/diffusion_pusht があります。

本ページは LeRobot 経路のみ を扱います。論文ベンチマークの厳密再現には元リポジトリを使ってください。

使い分け

条件付きデノイズで動作系列をモデル化します。連続で滑らかな制御や、複数の有効軌道があるタスクに向きます。ACT より重く長い学習になりがちです。

公式参照（LeRobot v0.5.0）

• examples/training/train_policy.py — PushT での学習例
• lerobot/diffusion_pusht モデルカード — lerobot-train と --policy.type=diffusion など

学習はインストール済み CLI：

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 なしでオフライン学習と CKPT ロードを確認。

LeRobot GPU イメージ

IO-AI の ioaitech/lerobot-gpu:v0.5.0 に LeRobot v0.5.0・GPU 依存・動画デコード・lerobot-train が入っています。

スモーク

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'

フラグが拒否される場合は、使用中のイメージ／コミットで lerobot-train --help を確認。

公式ソースからの再現

インストール のあと v0.5.0 を checkout：

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-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 でスモーク、その後 VRAM に合わせて増やす。
--steps	ACT より長め。まず 100k 前後の基準。
--save_freq	複数 CKPT を残して動作品質を比較。
--wandb.enable	長時間ジョブではオン推奨。

ホライズンや n_action_steps などはリリースごとに変わり得るため、固定したバージョンのヘルプと設定クラスを正とします。

評価の注意

オフライン損失はロボット成功率の弱い代理指標です。滑らかさ・再現性・遮蔽耐性・推論遅延を実機またはシムで見てください。

トラブルシューティング

キー不一致

meta/info.json の画像・状態・動作キーと FPS を確認。

OOM

バッチとワーカーを下げ、短い steps でスモークしてから本番。

元実装と結果が合わない

前処理・データローダ・環境が異なります。論文レベルなら columbia-ai-robotics/diffusion_policy を使用。

参考リンク

• Diffusion Policy プロジェクト
• 元リポジトリ
• LeRobot 実装
• diffusion_pusht モデルカード
• インストール