fine-tuning-serving-openpi

OpenPI のファインチューニングとサービング

公開 openpi リポジトリの Physical Intelligence 製 OpenPI モデル（pi0、pi0-fast、pi0.5）を、ロボットマニピュレーションタスク向けにファインチューニングしてサービングするエンドツーエンドのワークフローです。初期状態のマシンのセットアップ、JAX トレーニング、PyTorch トレーニング、チェックポイントの変換、ポリシーの推論サービングについて説明します。

クイックスタート

公開リポジトリをクローンし、ワークスペースをインストールしてから、事前学習済みポリシーをサービングします。

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 .
uv run scripts/serve_policy.py --env DROID

from openpi_client import websocket_client_policy

client = websocket_client_policy.WebsocketClientPolicy(host="localhost", port=8000)
result = client.infer(observation)
actions = result["actions"]  # numpy array of shape (chunk_size, action_dim)

主な概念

モデルファミリー：OpenPI は Physical Intelligence の 3 種類のモデルを実装しています。

モデル	アーキテクチャ	速度	品質	代表的な用途
pi0	フローマッチング VLA	ベースライン	最高	研究、複雑なタスク
pi0-fast	自己回帰アクショントークン	2～5 倍高速	良好	リアルタイム制御
pi0.5	pi0 + 改良されたビジョンエンコーダー	ベースライン	最良	最新のデフォルト

主な設計上の選択：

• デュアルバックエンド：JAX（主要、公式トレーニング）と PyTorch（コミュニティ、デプロイに適した構成）
• 設定駆動型：すべてのトレーニング／サービングパラメーターは src/openpi/training/config.py で定義されています
• 正規化統計：トレーニング前に、各設定で事前計算済みの正規化統計が必要です
• WebSocket サービング：ポリシーサーバーは、低レイテンシーの推論に対応する WebSocket API を公開します

トレーニングループの不変条件：設定またはデータセットを変更した後は、必ず次のサイクルを再実行してください。

1. 正規化統計を計算 → 2. トレーニング → 3. チェックポイントをサービング → 4. 推論を検証

計算要件

タスク	GPU	VRAM	備考
pi0.5 のサービング（推論）	1x A100/H100	約 24 GB	GPU 1 基で十分です
pi0.5 のファインチューニング（JAX）	1x A100 80GB	約 60 GB	マルチ GPU には fsdp_devices を使用
pi0 のファインチューニング（JAX）	1x A100 80GB	約 40 GB	モデルのフットプリントが小さくなります
ファインチューニング（PyTorch DDP）	1-8x A100	約 40 GB/GPU	torchrun ランチャー
正規化統計の計算	CPU または 1x GPU	約 8 GB	高速で、ログインノード上でも実行できます

ワークフロー 0：初期状態のマシンのセットアップ

このチェックリストをコピーして、進捗を記録してください。

Setup Progress:
- [ ] Step 1: Clone the public openpi repo with submodules
- [ ] Step 2: Install uv and sync the workspace
- [ ] Step 3: Install the editable package
- [ ] Step 4: Verify core imports and serving entrypoint

ステップ 1：リポジトリをクローン

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

サブモジュールなしで既にクローンしている場合：

git submodule update --init --recursive

ステップ 2：依存関係を同期

GIT_LFS_SKIP_SMUDGE=1 uv sync

ステップ 3：編集可能パッケージをインストール

GIT_LFS_SKIP_SMUDGE=1 uv pip install -e .

ステップ 4：インストールを確認

uv run python -c "from openpi.training import config as _config; print(_config.get_config('pi05_droid').name)"
uv run scripts/serve_policy.py --help

使用するケースと代替手段

次の場合にこのスキルを使用します：

• LeRobot または RLDS のデータセットで pi0、pi0-fast、または pi0.5 をファインチューニングする場合
• ALOHA、DROID、または LIBERO の評価向けに OpenPI ポリシーをサービングする場合
• JAX チェックポイントを PyTorch 形式に変換する場合
• OpenPI のトレーニングに関する問題（正規化統計、メモリ、設定）をデバッグする場合

次の場合は代わりに fine-tuning-openvla-oft を使用します：

• 連続アクションヘッドと LoRA を使用して OpenVLA をファインチューニングする場合
• LIBERO または ALOHA で OpenVLA-OFT 論文の結果を再現する場合

次の場合は代わりに evaluating-cosmos-policy を使用します：

• シミュレーションベンチマークで NVIDIA Cosmos Policy を評価する場合

---

ワークフロー 1：LeRobot データでの JAX ファインチューニング

このチェックリストをコピーして、進捗を記録してください。

JAX Fine-Tuning Progress:
- [ ] Step 1: Select and copy closest training config
- [ ] Step 2: Update dataset mapping and base checkpoint
- [ ] Step 3: Compute normalization statistics
- [ ] Step 4: Launch JAX training
- [ ] Step 5: Serve checkpoint and run inference sanity check

ステップ 1：設定を選択

src/openpi/training/config.py から最も近い設定をコピーします。

設定	用途
pi05_libero	pi0.5 の LIBERO ファインチューニング
pi0_libero	LIBERO での pi0 のフルファインチューニング
pi0_fast_libero	LIBERO での pi0-fast
pi0_aloha_pen_uncap	ALOHA のカスタムデータ
pi05_droid_finetune	小規模なカスタム DROID データセット（LeRobot 形式）
pi05_full_droid_finetune	大規模なフル DROID RLDS トレーニング

ステップ 2：データセットとトランスフォームを更新

# In src/openpi/training/config.py, modify your config:
TrainConfig(
    name="my_custom_config",
    model_type="pi05",
    data=LeRobotDataConfig(
        repo_id="your-org/your-dataset",
        # Adjust transforms to match your data format
    ),
    weight_loader=Pi05WeightLoader(),  # Match model type
)

データセットの repo_id を設定し、weight_loader がモデルの種類（pi0 または pi0.5）と一致していることを確認します。

ステップ 3：正規化統計を計算

uv run scripts/compute_norm_stats.py --config-name <config_name>

設定、データセット、またはトランスフォームを変更した場合は、トレーニングを開始するたびに、事前に実行する必要があります。

ステップ 4：JAX トレーニングを開始

XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 uv run scripts/train.py <config_name> \
  --exp-name=<run_name> \
  --overwrite

フル DROID RLDS トレーニングでは、rlds 依存関係グループを追加します。

uv run --group rlds scripts/compute_norm_stats.py \
  --config-name pi05_full_droid_finetune \
  --max-frames 10000000

XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 uv run --group rlds scripts/train.py \
  pi05_full_droid_finetune \
  --exp-name=<run_name> --overwrite

ステップ 5：サービングして検証

uv run scripts/serve_policy.py policy:checkpoint \
  --policy.config=<config_name> \
  --policy.dir=checkpoints/<config_name>/<run_name>/<step>

テストクライアントで確認します。

from openpi_client import websocket_client_policy

client = websocket_client_policy.WebsocketClientPolicy(host="localhost", port=8000)
# Build observation matching your config's expected keys
obs = {"image": img_array, "state": state_array, "prompt": "pick up the cup"}
result = client.infer(obs)
print(f"Action shape: {result['actions'].shape}")  # (chunk_size, action_dim)

---

ワークフロー 2：PyTorch のトレーニングとチェックポイントの変換

このチェックリストをコピーして、進捗を管理してください：

PyTorch Setup Progress:
- [ ] Step 1: Sync dependencies and verify transformer version
- [ ] Step 2: Apply OpenPI transformer patches
- [ ] Step 3: Convert JAX checkpoint to PyTorch format
- [ ] Step 4: Launch PyTorch training or serve converted checkpoint

ステップ 1：依存関係を同期する

uv sync
uv pip show transformers

ステップ 2：必要なパッチを適用する

OpenPI PyTorch では、インストール済みの transformers パッケージにカスタム変更を加える必要があります：

cp -r ./src/openpi/models_pytorch/transformers_replace/* \
  .venv/lib/python3.11/site-packages/transformers/

ステップ 3：JAX チェックポイントを変換する

uv run examples/convert_jax_model_to_pytorch.py \
  --checkpoint_dir <jax_checkpoint_dir> \
  --config_name <config_name> \
  --output_path <pytorch_checkpoint_dir>

ステップ 4：トレーニングまたはサービングを行う

単一 GPU でのトレーニング：

uv run scripts/train_pytorch.py <config_name> --exp_name <run_name>

マルチ GPU 分散トレーニング：

uv run torchrun --standalone --nnodes=1 --nproc_per_node=<num_gpus> \
  scripts/train_pytorch.py <config_name> --exp_name <run_name>

変換済みチェックポイントを使ったプログラムからの推論：

from openpi.training import config as _config
from openpi.policies import policy_config

config = _config.get_config("pi05_droid")
policy = policy_config.create_trained_policy(config, "<pytorch_checkpoint_dir>")
result = policy.infer(example)
actions = result["actions"]  # numpy array

チェックポイントは、次の形式に従います：checkpoints/<config_name>/<exp_name>/<step>/。

---

ワークフロー 3：ポリシー推論のサービング

このチェックリストをコピーして、進捗を管理してください：

Inference Server Progress:
- [ ] Step 1: Choose target environment and checkpoint
- [ ] Step 2: Start policy server
- [ ] Step 3: Confirm server is reachable
- [ ] Step 4: Integrate client into robot or simulation code

ステップ 1：環境を選択する

デフォルトの環境プリセット：

環境	設定	デフォルトのチェックポイント
ALOHA	pi05_aloha	gs://openpi-assets/checkpoints/pi05_base
ALOHA_SIM	pi0_aloha_sim	gs://openpi-assets/checkpoints/pi0_aloha_sim
DROID	pi05_droid	gs://openpi-assets/checkpoints/pi05_droid
LIBERO	pi05_libero	gs://openpi-assets/checkpoints/pi05_libero

ステップ 2：サーバーを起動する

デフォルトモード（プリセットのチェックポイントを使用）：

uv run scripts/serve_policy.py --env ALOHA

明示的なチェックポイントモード（カスタムモデルまたはローカルモデル）：

uv run scripts/serve_policy.py policy:checkpoint \
  --policy.config=pi05_libero \
  --policy.dir=checkpoints/pi05_libero/my_run/20000

実行時の観測にプロンプトが含まれない場合は、--default_prompt "task description" を追加してください。

ステップ 3：接続を確認する

uv run examples/simple_client/main.py --env DROID

ステップ 4：リモートクライアントをロボットコードに組み込む

ロボット環境に軽量クライアントをインストールします：

pip install "openpi-client @ git+https://github.com/Physical-Intelligence/openpi.git#subdirectory=packages/openpi-client"

完全な統合例：

from openpi_client import websocket_client_policy
import numpy as np

# Connect to remote policy server
client = websocket_client_policy.WebsocketClientPolicy(
    host="gpu-server.local", port=8000
)

# Build observation (keys must match policy transforms)
observation = {
    "image": np.random.rand(224, 224, 3),  # RGB image
    "state": np.zeros(7),                   # Joint positions
    "prompt": "pick up the red block",
}

# Get actions
result = client.infer(observation)
actions = result["actions"]  # shape: (action_chunk_size, action_dim)

# Execute first action on robot
robot.step(actions[0])

---

よくある問題

問題：norm stats が見つからないというエラー

対処方法：トレーニングの前に scripts/compute_norm_stats.py --config-name <config_name> を実行してください。

問題：JAX のトレーニング中にメモリ不足になる

対処方法：XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 を設定する、バッチサイズを小さくする、または fsdp_devices を設定してください：

# In config: use model-parallel sharding
TrainConfig(
    ...
    fsdp_devices=4,  # Shard across 4 GPUs
)

問題：PyTorch チェックポイントの読み込み中に OOM が発生する

対処方法：export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

問題：設定が見つからない

対処方法：設定名が src/openpi/training/config.py に存在することを確認してください（_CONFIGS 辞書との完全一致が必要です）。

問題：ライブラリの変更後、PyTorch のトレーニングが収束しなくなる

対処方法：transformer のパッチを再適用してください。uv cache clean transformers を実行してキャッシュをリセットしてから、再適用してください。

問題：serve_policy.py が ModuleNotFoundError でクラッシュする

対処方法：まずパブリックワークスペースを再同期してください：

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

不足しているモジュールがシミュレーター関連の場合は、その例で必要となる追加のランタイム依存関係をインストールしてください：

uv pip install pytest robosuite==1.4.0 gym bddl easydict matplotlib

問題：uv sync が rerun-sdk のホイールの不一致で失敗する

対処方法：

uv sync --no-dev
# or
uv sync --no-dev --no-install-package rerun-sdk

問題：チェックポイントのダウンロードがタイムアウトする

対処方法：gsutil をインストールして、手動で先行ダウンロードしてください：

pip install gsutil
gsutil -m cp -r gs://openpi-assets/checkpoints/pi05_libero /local/cache/

以前のダウンロードが中断されていた場合は、古い .lock ファイルを削除してください。

問題：ポリシーサーバーが終了コード 137 で終了する

対処方法：OOM による強制終了です。JAX のメモリ変数を設定してください：

export XLA_PYTHON_CLIENT_PREALLOCATE=false
export XLA_PYTHON_CLIENT_ALLOCATOR=platform

---

HPC／クラスターのユーザー向け

詳細、コード例、パラメータ表については、references/for-hpc-cluster-users.md を参照してください。このステップが該当する場合は、内容を確認してください。

高度なトピック

設定レシピとベースライン： references/config-recipes.md を参照してください トレーニングのデバッグガイド： references/training-debugging.md を参照してください チェックポイントと環境の対応表： references/checkpoints-and-env-map.md を参照してください リモートクライアントの統合： references/remote-client-pattern.md を参照してください PyTorch の精度とパッチ適用に関する注意点： references/pytorch-gotchas.md を参照してください

リソース

• OpenPI リポジトリ：https://github.com/Physical-Intelligence/openpi
• OpenPI クライアントパッケージ：https://github.com/Physical-Intelligence/openpi/tree/main/packages/openpi-client
• pi0 論文：https://www.physicalintelligence.company/blog/pi0
• LeRobot データセット形式：https://huggingface.co/docs/lerobot

エージェントの作業手順

1. 環境を確認します。 GPU メモリ、モデルのチェックポイントとライセンス、フレームワークのバージョン、入力形式を確認します。
2. 入力を明確にします。 データまたはユーザーに確認し、形式、識別子、パラメーターを確定します。結果に影響する値は推測せず、確認してください。
3. まず小規模なバージョンを実行します。 まず、低解像度または短い長さで単一サンプルの推論を実行します。
4. 上記の手順と参考資料に従って、タスク全体を実行します。
5. 結果を検証します。 出力を定性的に確認し、タスク指標でも評価します。生成コンテンツのライセンス制限も確認してください。
6. 報告します。 実行した内容（バージョン、コマンド、パラメーター）、確認した内容、未解決の点を明記します。

このような場合	対応方法
大きな入力でメモリ不足が発生する	解像度またはシーケンス長を下げる、半精度を使用する、または入力をタイルに分割します。
これらの手順に記載されている関数、フラグ、またはエンドポイントがインストール済みのバージョンにない	インストール済みバージョンのドキュメント（help()、--help、公式ドキュメント）を確認して対応し、ユーザーに伝えます。API を捏造してはいけません。
必須の入力、識別子、またはパラメーターが曖昧である	ユーザーに確認するか、実行前に前提を明示します。

誠実さに関するルール

• 結果、パラメーター、識別子、引用、統計情報を決して捏造しないでください。実行または検証できないことは、明確に伝えてください。
• 商用利用または公開に先立ち、モデルとデータセットのライセンスを確認してください。
• ここに記載されたバージョン固有の詳細は古くなっている可能性があります。インストール済みバージョンの公式ドキュメントで確認してください。
• 費用が発生する操作、共有 GPU またはクラウドのクォータを消費する操作、個人データまたは患者データを扱う操作、元に戻せない操作を行う前に、確認を求めてください。

関連スキル

• evaluating-cosmos-policy：LIBERO および RoboCasa のシミュレーション環境で NVIDIA Cosmos Policy を評価します。
• fine-tuning-openvla-oft：連続アクションヘッドと LoRA 適応を用いて、ロボットのアクション生成向けに OpenVLA-OFT および OpenVLA-OFT+ ポリシーをファインチューニングし、評価します。…
• pufferlib：PufferLib の強化学習環境、ベクトル化、ポリシー、PuffeRL のトレーニング、評価、安全な c…に関する、バージョンを考慮したガイダンスです。