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／クラスターのユーザー向け

Slurm 管理のクラスターでは、リソース割り当てを使ってコマンドを実行します。

srun --partition=gpu --gpus-per-node=1 --mem=64G --cpus-per-task=8 --pty bash

/home がいっぱいにならないよう、キャッシュの保存先をスクラッチ領域に変更してください。

export HF_HOME=/scratch/$USER/.cache/huggingface
export XDG_CACHE_HOME=/scratch/$USER/.cache
export PIP_CACHE_DIR=/scratch/$USER/.cache/pip
export UV_CACHE_DIR=/scratch/$USER/.cache/uv

uv 管理の環境を使う場合は、クラスターの Python モジュールを重ねて読み込まないでください。通常は module load cuda だけで十分です。

---

高度なトピック

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

Resources

• OpenPI repository: https://github.com/Physical-Intelligence/openpi
• OpenPI client package: https://github.com/Physical-Intelligence/openpi/tree/main/packages/openpi-client
• pi0 paper: https://www.physicalintelligence.company/blog/pi0
• LeRobot dataset format: https://huggingface.co/docs/lerobot