メインコンテンツまでスキップ

Pi0 / Pi0.5 モデルファインチューニングガイド

本稿では、IO-AI.TECH が公開する Docker イメージを用いて、pi0_base / pi05_base などの事前学習済み重みの上で Pi0 と Pi0.5 モデルをファインチューニングする手順を説明します。コマンド、マウント規約、引数名はイメージ内の OpenPI ラッパー train_lerobot.py と一致します。

公式 OpenPI のファインチューニングフローを、環境構築なしで使いたい場合に適したパスです。

この経路を推奨する理由​

Pi0 / Pi0.5 は OpenPI（内部は JAX）と公開ベースチェックポイント上でファインチューニングします。LeRobot データセットが用意できていれば、次のイメージが最も手早いです。

  • ioaitech/train_openpi:pi0
  • ioaitech/train_openpi:pi05

いずれもファインチューニングに必要な依存関係を含み、次の規約を使います。

  • 入力：/data/input
  • 出力（チェックポイントなど）：/data/output
Docker イメージレジストリ

イメージは Docker Hub に公開されています（ioaitech/train_openpi:pi0、ioaitech/train_openpi:pi05）。

ワンコマンドでファインチューニング​

前提条件​

  • Linux ホスト
  • NVIDIA ドライバ正常
  • Docker
  • docker run --gpus all が動作

GPU 確認：

docker run --rm --gpus all nvidia/cuda:12.1.0-base-ubuntu22.04 nvidia-smi

Pi0​

docker run --rm --gpus all \
  -v /path/to/lerobot_dataset:/data/input \
  -v /path/to/output:/data/output \
  ioaitech/train_openpi:pi0 \
  --batch_size 8 \
  --steps 20000 \
  --save_interval 1000 \
  --learning_rate 2.5e-5 \
  --action_horizon 50 \
  --prompt "pick up the object"

Pi0.5​

docker run --rm --gpus all \
  -v /path/to/lerobot_dataset:/data/input \
  -v /path/to/output:/data/output \
  ioaitech/train_openpi:pi05 \
  --batch_size 8 \
  --steps 20000 \
  --save_interval 1000 \
  --learning_rate 2.5e-5 \
  --action_horizon 50 \
  --prompt "pick up the object"

主な違いはイメージタグです。ビルド時の MODEL_TYPE に応じてベース重みと Pi0Config が選ばれます。

最小スモーク実行​

docker run --rm --gpus all \
  -v /path/to/lerobot_dataset:/data/input \
  -v /path/to/output:/data/output \
  ioaitech/train_openpi:pi0 \
  --batch_size 1 \
  --steps 1000

成功後に batch_size や steps、その他のファインチューニング用ハイパーパラメータを増やします。

データ要件​

/data/input/meta/info.json が必須です。ルート構造の例：

your_dataset/
├── meta/
│   └── info.json
├── data/
└── videos/

ファインチューニング用ラッパーは次を自動処理します。

  • LeRobot データセット版の検出
  • 必要時に v3 をパイプライン互換の v2.1 レイアウトへ
  • v2.1 で episodes_stats.jsonl が無い場合の生成
  • 正規化統計の計算
  • LeRobot キャッシュへのシンボリックリンク

Pi0 と Pi0.5 の違い​

同一のファインチューニングラッパーですが設定と重みが異なります。

  • Pi0 は pi0_base
  • Pi0.5 は pi05_base
  • Pi0.5 は Pi0Config(pi05=True)（状態・トークンまわりが Pi0 と異なる）
  • Pi0.5 はより大きいデフォルト max_token_len

パイプライン検証は Pi0 から。Pi0.5 のベースを明示的に使う場合に ioaitech/train_openpi:pi05 に切り替えます。

主な引数​

以下はイメージ内 train_lerobot.py の定義と一致します。

引数デフォルト説明
--batch_size1各最適化ステップのバッチ。JAX デバイス数で割り切れるよう調整される場合あり
--steps1000最適化ステップ数
--gpusall全 GPU または 0,1 など
--prompt空タスクテキストが無いデータセット用のデフォルト指示
--save_interval500チェックポイント間隔
--learning_rate空未指定時ピーク LR 2.5e-5
--fsdp_devicesautoFSDP デバイス数（自動）
--loraauto単一 GPU では既定で LoRA 有効、複数 GPU では無効になりやすい
--ema_decay空LoRA 時は VRAM 節約のため既定でオフになりやすい
--action_horizon50アクション列長

ファインチューニング時の挙動​

  • 単一 GPU：--lora auto で LoRA が有効になりやすい
  • 複数 GPU：--fsdp_devices auto で FSDP 系の経路
  • batch_size がデバイス数より小さい、または割り切れない場合は引き上げられる

そのため、まずは「動く」コマンドを優先し、JAX の細部は後から調整できます。

--prompt について​

エピソードにタスク文字列があれば、ファインチューニング時はそちらが優先されます。タスクフィールドが無いデータセットでのみ --prompt がフォールバックとして効きます。

出力​

TrainConfig で次が固定されています。

  • name = docker_train
  • exp_name = train

したがってチェックポイントの既定パスは：

/path/to/output/docker_train/train/

正規化統計はアセットディレクトリに保存され、以降のファインチューニングや推論で再利用されます。

推奨ワークフロー​

1. 新データセットの短い検証​

docker run --rm --gpus all \
  -v /path/to/lerobot_dataset:/data/input \
  -v /path/to/output:/data/output \
  ioaitech/train_openpi:pi0 \
  --batch_size 1 \
  --steps 1000 \
  --save_interval 200

2. 単一 GPU の保守的設定​

docker run --rm --gpus all \
  -v /path/to/lerobot_dataset:/data/input \
  -v /path/to/output:/data/output \
  ioaitech/train_openpi:pi0 \
  --batch_size 4 \
  --steps 20000 \
  --save_interval 1000 \
  --action_horizon 50

3. マルチ GPU​

docker run --rm --gpus all \
  -v /path/to/lerobot_dataset:/data/input \
  -v /path/to/output:/data/output \
  ioaitech/train_openpi:pi05 \
  --gpus 0,1,2,3 \
  --batch_size 16 \
  --steps 30000 \
  --fsdp_devices 4 \
  --save_interval 1000

よくある質問​

1. 起動時にデータセット無し​

/data/input のマウントと meta/info.json を確認してください。

2. 単一 GPU と複数 GPU で挙動が違う​

LoRA・FSDP・バッチ調整がデバイス数依存のため、意図した設計です。

3. 出力ディレクトリ名が固定なのはなぜ​

現行ラッパーは /data/output/docker_train/train に固定です。実験名分けは別レイヤで行う想定です。

4. 正規化統計を手動で先に回す必要は​

ありません。ラッパーがファインチューニング開始前に計算・保存します。

5. Pi0 と Pi0.5 の選び方​

まず Pi0 でパイプラインを安定させ、Pi0.5 ベースが必要なら ioaitech/train_openpi:pi05 を使います。

実践的な提案​

  • 長時間ジョブの前にデータ構造を確認
  • 短い実行の後に長い実行
  • 単一 GPU では、まずデフォルトの自動戦略を尊重
  • Pi0 と Pi0.5 は出力ルートを分けて比較しやすくする

参考資料​