Pi0 モデルのファインチューニング：カスタムデータセットを使ったロボットタスクへの適応
Pi0 は先進的な視覚-言語-行動（Vision-Language-Action, VLA）モデルで、Few-Shot Fine-Tuning を通じて特定のロボットタスクに素早く適応できます。
本ドキュメントは LeRobot フレームワークに基づき、事前学習済みの Pi0 モデル（lerobot/pi0）をカスタムロボットデータセットでファインチューニングします。
前提条件
- LeRobot 形式のデータセットを保有（前節の LeRobot データセットエ クスポートガイドを参照）
- Pi0 モデルを特定のロボットハードウェア、タスク、または制御戦略に適応させる必要がある
- PyTorch と Hugging Face エコシステムのトレーニングプロセスに慣れている
- 深層学習モデルトレーニングの基本的な経験がある
ファインチューニングの概要
ファインチューニング（Fine-Tuning）は、事前学習モデルを基にドメイン固有のデータでさらにトレーニングするプロセスで、通常少ないイテレーションを伴います。
Pi0 モデルは多様なロボットタスクで事前学習されており、一般的な視覚知覚、言語理解、行動生成能力を備えています。
ファインチューニングにより Pi0 は達成可能：
- 環境適応：特定のカメラ視点、照明条件、機械構造への調整
- タスク特化：特定のタスク（例：物体把持、分類、配置）の性能最適化
- 精度向上：目標タスクでの制御精度と成功率の大幅向上
簡単に言うと：事前学習モデルは多くを知っていますが、あなたのことを必ずしも知りません。ファインチューニング後、それはあなたを理解します。
環境準備
システム要件
環境が以下の要件を満たしていることを確認してください：
- Python ≥ 3.8（推奨 3.10 以上）
- GPU：少なくとも 32GB VRAM（Pi0 は大規模モデル、NVIDIA V100 またはより高性能の GPU を推奨）
- メモリ：少なくとも 64GB システム RAM
- ストレージ：データセットとモデルチェックポイントのための十分なディスクスペース
依存関係のインストール
# LeRobot リポジトリをクローン（プライマリ）
git clone https://github.com/lerobot-ai/lerobot.git
cd lerobot/
# LeRobot フレームワークをインストール（Pi0 サポートを含む）
pip install -e ".[pi0]"
# インストールを確認
python -c "from lerobot.policies import Pi0Policy; print('Pi0 が正常にインストールされました！')"
ミラー（必要に応じて）：https://github.com/huggingface/lerobot
データセットの準備
LeRobot 形式データのエクスポート
IO データプラットフォームのエクスポートページを通じて、アノテーション済みデータを LeRobot 形式データセットにエクスポートできます。 LeRobot データセットのエクスポート を参照。
重要なお知らせ：データを HuggingFace にアップロードせずにトレーニング可能。--dataset.root= パラメータでローカルディレクトリのデータを指定してファインチューニング。
データセット構造の例
エクスポートされたデータが ~/DualPiper_Pickup_Pen に保存されていると仮定：
$ ls ~/DualPiper_Pickup_Pen
data meta videos
$ ls ~/DualPiper_Pickup_Pen/data
episode_000 episode_001 episode_002 ...
$ ls ~/DualPiper_Pickup_Pen/data/episode_000
observations actions.npy language_instruction.txt metadata.json
ファインチューニングの開始
基本トレーニングコマンド
train.py スクリプトを使用してトレーニング：
# CUDA メモリ割り当てポリシーを設定（推奨）
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
# トレーニングを開始
python3 -m lerobot.scripts.train \
--dataset.root=~/DualPiper_Pickup_Pen \
--dataset.repo_id=io-ai-data/DualPiper_Pickup_Pen \
--policy.path=lerobot/pi0 \
--policy.repo_id=lerobot/pi0 \
--policy.device=cuda \
--output_dir=./checkpoints/pi0_finetuned \
--batch_size=1 \
--policy.attention_implementation=flash_attention_2 \
--training.learning_rate=1e-5 \
--training.num_epochs=10
Pi0 モデルはパラメータ数が多く、シングル GPU トレーニングは数十時間かかる可能性があります。例えば、2 時間のデータセットのファインチューニングはシングル NVIDIA V100 で約 50 時間、8 枚の V100 で約 10 時間に短縮されます。
パラメータの詳細
| パラメータ | 意味 | 推奨値 | 説明 |
|---|---|---|---|
--dataset.root | ローカルデータセットパス | ~/your_dataset | LeRobot 形式データセットディレクトリを指す |
--dataset.repo_id | Hugging Face データセット ID | your-username/dataset | メタデータ識別用 |
--policy.path | 事前学習モデルパス | lerobot/pi0 | 公式事前学習モデルを指定 |
--policy.repo_id | モデルリポジトリ ID | lerobot/pi0 | Hugging Face モデルリポジトリ |
--policy.device | トレーニングデバイス | cuda | GPU アクセラレーションを有効化 |
--output_dir | 出力ディレクトリ | ./checkpoints/pi0_finetuned | ファインチューニングモデル保存パス |
--batch_size | バッチサイズ | 1 | VRAM に基づいて調整 |
--policy.attention_implementation | 注意機構 | flash_attention_2 | 効率的な実装を使用してトレーニングを加速 |
--training.learning_rate | 学習率 | 1e-5 | ファインチューニング時に低い値を使用して壊滅的な忘却を避ける |
--training.num_epochs | トレーニングエポック数 | 10 | データセット規模に基づいて調整 |
マルチ GPU アクセラレーテッドトレーニング
複数の GPU がある場合、並列トレーニングで速度を大幅に向上させられます。
# 8 GPU でトレーニング
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
torchrun --nproc_per_node=8 --master_port=29500 \
-m lerobot.scripts.train \
--dataset.root=/data/nfs2/export/lerobot/DualPiper_Pickup_Pen/ \
--dataset.repo_id=io-ai-data/DualPiper_Pickup_Pen \
--policy.path=lerobot/pi0 \
--policy.repo_id=lerobot/pi0 \
--policy.device=cuda \
--output_dir=/home/puxk/DualPiper_Pickup_Pen_pi0_model \
--batch_size=1 \
--num_workers=4 \
--policy.attention_implementation=flash_attention_2
総バッチサイズは nproc_per_node × batch_size。--batch_size=1 で 8 カードは総バッチサイズ 8 に相当。メモリが十分なら増やせます。
デバッグ時は --nproc_per_node=2 でテストし、全 8 カードに展開。
特別注意事項
- ネットワークアクセス：Hugging Face モデルのダウンロードに hf-mirror などのプロキシツールが必要な場合あり
- VRAM 監視：トレーニング中に GPU 使用率をリアルタイム監視
- チェックポイント回復：トレーニング中断時、
--output_dirの最新チェックポイントから再開
Pi0 モデルは言語処理に Google PaliGemma を統合。最近の Hugging Face モデルとして：
バージョン互換性
- Transformers：≥ 4.37.0 必要
- 互換性チェック：
embed_tokensエラーが発生したら transformers をアップグレード
解決策
# transformers を最新版にアップグレード
pip install -U transformers
# またはバージョン指定
pip install transformers>=4.37.0
方針選択（Pi0 vs SmolVLA）
- 単一タスク・短ホライゾン・高制御周波数で言語条件の要求が高くない場合：Pi0/ACT 系ポリシーは計算効率と応答性に優れる傾向。
- マルチタスク/クロスシーンの汎化や単一 GPU/コンシューマ環境の親和 性を重視する場合：SmolVLA を推奨。
smolvla_baseの少量微調整で安定した利得が得られることが多い。 - 資源制約下では：混合精度、勾配累積、入力解像度とバッチサイズの適正化を優先。
トレーニング結果の検査
モデル保存
トレーニング完了後、モデルファイルは以下に位置：
./checkpoints/pi0_finetuned/
├── config.json
├── pytorch_model.bin
├── tokenizer.json
└── training_args.bin
ファインチューニングモデルのロード
from lerobot import Pi0Policy # 注意：インポートパスは lerobot.policies の可能性あり
# モデルをロード
policy = Pi0Policy.from_pretrained("./checkpoints/pi0_finetuned")
# 推論例
action = policy.select_action(observation_batch)
ログ分析
トレーニングログには以下を含む：
- 損失曲線：トレーニング/検証損失トレンド
- 学習率：スケジューラ変化
- 勾配統計：ノルムとクリッピング情報
- リソース利用：GPU/CPU メモリ使用
パフォーマンス評価
IO データプラットフォーム評価
IO プラットフォームはモデル推論品質評価を提供し 、実データとモデル出力の視覚比較をサポート、各関節のコマンド比較を含む。
その他の評価方法
定量メトリクス
- 行動エラー：予測行動と ground truth の平均絶対誤差 (MAE) または平均二乗誤差 (MSE)
- 軌跡類似度：動的時間伸縮 (DTW) または Fréchet 距離を使用して軌跡を比較
- 検証損失：検証セット上の損失関数値を監視
定性評価
- 実機テスト：物理ロボットにデプロイし、タスク成功率を測定
- 行動分析：生成された行動シーケンスが合理的かを視覚化
- 人間評価：モデル出力に対する専門家の主観スコアを収集
評価スクリプト例
import numpy as np
from lerobot import Pi0Policy
from torch.utils.data import DataLoader # test_dataset が DataLoader であると仮定
# モデルをロード
policy = Pi0Policy.from_pretrained("./checkpoints/pi0_finetuned")
# 評価関数
def evaluate_model(policy, test_loader):
total_mae = 0.0
num_batches = 0
for batch in test_loader:
with torch.no_grad():
predicted_action = policy.select_action(batch)
true_action = batch['actions']
# MAE を計算
mae = np.mean(np.abs(predicted_action.cpu().numpy() - true_action.cpu().numpy()))
total_mae += mae
num_batches += 1
return total_mae / num_batches
# 評価を実行（test_loader が準備されていると仮定）
avg_mae = evaluate_model(policy, test_loader)
print(f"平均絶対誤差 (MAE): {avg_mae:.4f}")