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

Pi0 / Pi0.5 のファインチューニング

Pi0 / Pi0.5 は Physical Intelligence の OpenPI プロジェクト由来です。公式リポジトリ Physical-Intelligence/openpi に pi0_base や pi05_base などの基盤重みがあり、カスタムデータの主入口として LeRobot データセットを想定しています。Pi0.5 は現時点の公式説明では flow matching ヘッド を中心に学習・推論が説明されています。より階層的な計画が必要なタスクでは、公式実装の範囲を先に確認してください。

IO-AI は OpenPI の学習ラインを基に ioaitech/train_openpi:pi0 と ioaitech/train_openpi:pi05 の Docker イメージを公開しています。OpenPI 実行環境・基盤重み取得・正規化統計計算・LeRobot データアダプタが同梱されています。まずイメージでループを閉じ、データマッピングや設定を深く変える必要が出たら公式ソースに戻してください。

使い分け​

高品質なマルチモーダルデータがあり、大規模モデル上でファインチューニングしたい場合向けです。ACT や Diffusion Policy よりスタックが重く、JAX・FSDP・LoRA・基盤重み・正規化統計・言語タスク列が揃っている必要があります。

モデルイメージ基盤チェックポイント用途の目安
Pi0ioaitech/train_openpi:pi0gs://openpi-assets/checkpoints/pi0_baseOpenPI のファインチューニングを最初に通す。言語一般化は中程度まで。
Pi0.5ioaitech/train_openpi:pi05gs://openpi-assets/checkpoints/pi05_baseより広い環境一般化。データと計算資源が十分な場合。

データ要件​

イメージのエントリは /data/input/meta/info.json を確認します。最低限のディレクトリ構成：

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

ラッパーは概ね次を行います。

  • LeRobot のスキーマ読み取り
  • コンテナ内 LeRobot キャッシュへのリンク
  • v2.1 で episodes_stats.jsonl が欠ける場合の補完
  • Hugging Face parquet メタデータの一部互換処理
  • OpenPI 用の正規化統計計算
  • GPU 枚数に応じた単 GPU LoRA / 多 GPU FSDP の選択

タスク文はデータセット優先。データにタスク文字列が無い場合のみ --prompt がデフォルトになります。

イメージで学習​

GPU 確認：

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

スモークテスト​

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

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 4 \
  --steps 20000 \
  --learning_rate 2.5e-5 \
  --save_interval 1000 \
  --action_horizon 50

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 30000 \
  --learning_rate 2.5e-5 \
  --save_interval 1000 \
  --action_horizon 50

マルチ 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

出力の既定：

/path/to/output/docker_train/train/

ラッパー内 TrainConfig(name="docker_train", exp_name="train") に対応します。

よく使うフラグ​

フラグ既定意味
--batch_size1グローバルバッチ。JAX デバイス数で割り切れるよう調整。
--steps1000学習ステップ数。
--gpusall全 GPU または 0,1 など。
--prompt空データにタスク文が無いときのデフォルト指示。
--save_interval500チェックポイント間隔。
--learning_rate2.5e-5未指定時のピーク LR。
--fsdp_devicesauto多 GPU では可視 GPU 数。単 GPU は 1。
--loraauto単 GPU は既定で LoRA、多 GPU は FSDP 既定。
--ema_decay空LoRA 時はメモリ節約のため EMA 無効。
--action_horizon50動作系列長。
--num_workers8DataLoader ワーカー数。

公式 OpenPI での再現​

データマッピングやモデル設定を変える場合は公式へ。README の流れは uv で依存同期 → 設定作成/変更 → 正規化統計 → 学習です。

git clone --recurse-submodules https://github.com/Physical-Intelligence/openpi.git
cd openpi

GIT_LFS_SKIP_SMUDGE=1 uv sync
GIT_LFS_SKIP_SMUDGE=1 uv pip install -e .

公式の LIBERO 例：

uv run scripts/compute_norm_stats.py --config-name pi05_libero

XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 \
  uv run scripts/train.py pi05_libero \
  --exp-name=my_experiment \
  --overwrite

独自ロボットでは src/openpi/training/config.py のデータ設定と TrainConfig、入出力 transform を実装します。上記イメージには汎用 LeRobot アダプタが入っており、まず学習を通すのに便利です。深いマッピング変更は公式側で行ってください。

再現性の注意​

  • 基盤重みはビルド時に /models/openpi-assets/checkpoints/*_base/params に配置。
  • 既定では W&B と Hugging Face Hub へのネットワークをオフにしがち（オフライン向け）。
  • 単 GPU は LoRA 既定、多 GPU は FSDP 既定。比較実験では --lora・--fsdp_devices・--batch_size・シードを記録。
  • Pi0.5 は flow matching ヘッド中心の説明が中心。内部の全段階が公開されているとは限りません。

FAQ​

起動時に LeRobot データセットが無い​

/data/input へのマウントと /data/input/meta/info.json の存在を確認。親ディレクトリだけマウントしていないか。

単 GPU と多 GPU で結果が違う​

LoRA と FSDP で最適化経路が異なります。比較時はフラグとシードを固定。

正規化統計を手で？​

ioaitech/train_openpi:* ではラッパーが計算。公式ソースでは先に scripts/compute_norm_stats.py。

出力パスが固定に見える​

name=docker_train、exp_name=train が固定です。ホスト側の出力パスで実験を分けてください（例：/outputs/pi05_pick_block_v1）。

参考リンク​