Pi0モデルでデータをファインチューニング

Pi0は強力な視覚-言語-アクション（VLA）モデルで、少量のデータでファインチューニングすることで、ロボットタスクに素早く適応できます。

このドキュメントはLeRobotフレームワークに基づいており、事前学習済みのPi0モデル（lerobot/pi0）を使用して独自のロボットデータセットでファインチューニングを行います。

準備条件

• LeRobotデータフォーマットのデータセットをお持ちです（前節のLeRobotデータセットのエクスポートを参照）
• Pi0モデルをロボットハードウェア、タスク、または制御スタイルに適応させたい
• PyTorchとHugging Faceエコシステムでの学習に精通している
• 基本的なディープラーニング学習経験がある

ファインチューニングとは？

ファインチューニングとは、事前学習済みモデルの基盤上で、独自のデータを使用して短時間モデルを継続学習することです。

Pi0は複数の汎用ロボットタスクで事前学習されたモデルで、視覚理解、指示理解、アクション予測の多くの汎用能力を既に学習しています。

ファインチューニングにより、Pi0を以下にすることができます：

• ロボット環境への適応：カメラ視点、照明環境、機械構造に適応
• 特定タスクへの適応：特定のタスク（ペンの取り上げ、分類、配置など）に最適化
• 制御精度の向上：タスクでの制御精度と成功率を大幅に向上

簡単に言えば：事前学習モデルは多くのことを知っていますが、必ずしもあなたを理解しているわけではありません。ファインチューニング後、あなたを理解するようになります。

環境セットアップ

システム要件

環境が以下の要件を満たしていることを確認してください：

• Python ≥ 3.8（推奨3.10または3.11）
• GPU：少なくとも32GB VRAMを推奨（Pi0は大きなモデルで、V100以上のレベルのグラフィックカードを推奨）
• メモリ：少なくとも64GB RAMを推奨
• ストレージ：データセットとモデルを保存する十分なディスク容量

依存関係のインストール

# LeRobotリポジトリをクローン
git clone https://github.com/huggingface/lerobot.git
cd lerobot/

# LeRobotフレームワークをインストール（Pi0サポートを含む）
pip install -e ".[pi0]"

# インストールを確認
python -c "from lerobot.policies import Pi0Policy; print('Pi0インストール成功！')"

データセットの準備

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モデル自体がかなり大きいため、単一カードでの学習には少なくとも数十時間かかります。例えば、収集した2時間のデータをファインチューニングした場合、単一V100カードで50時間かかりました。8枚のV100カードで速度が向上しますが、それでも約10時間かかります。

パラメータ詳細

パラメータ	意味	推奨値	説明
--dataset.root	ローカルデータセットパス	~/your_dataset	LeRobotデータセットディレクトリを指す
--dataset.repo_id	Hugging Faceデータセット名	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	Flash Attention 2で加速
--training.learning_rate	学習率	1e-5	ファインチューニングには小さな学習率を使用
--training.num_epochs	学習エポック数	10	データ量に応じて調整

マルチ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枚すべてのカードに展開することをお勧めします。

特別な注意事項

1. ネットワーク加速：HuggingFaceからPi0モデルをダウンロードするためにネットワーク加速が必要な場合があります。hg-mirrorなどのツールを参照してください
2. VRAM監視：学習中はGPU VRAM使用量を監視してください
3. 学習再開：学習が中断された場合、--output_dirのチェックポイントから継続できます

重要な注意事項：Gemmaモデル互換性

Pi0モデルは内部でGoogle PaLI-Gemmaモデル構造を言語処理に使用しています。最近追加されたHugging Faceモデルであるため：

バージョン要件

• Transformers：最新版のtransformersを使用してください（推奨≥4.37.0）
• 互換性チェック：embed_tokens関連のエラーが発生した場合は、transformersをアップグレードしてください

解決策

# 最新版のtransformersにアップグレード
pip install -U transformers

# またはバージョンを指定
pip install transformers>=4.37.0

学習結果の確認

モデル保存場所

学習完了後、モデルは以下に保存されます：

./checkpoints/pi0_finetuned/
├── config.json
├── pytorch_model.bin
├── tokenizer.json
└── training_args.bin

ファインチューニングされたモデルの読み込み

from lerobot.policies import Pi0Policy

# ファインチューニングされたモデルを読み込み
policy = Pi0Policy.from_pretrained("./checkpoints/pi0_finetuned")

# 推論を実行
action = policy.select_action(batch)

学習ログ分析

学習中は、以下の詳細なログが生成されます：

• 損失曲線：学習損失と検証損失の変化
• 学習率スケジューリング：学習率の変化
• 勾配情報：勾配ノルムと勾配クリッピング
• メモリ使用量：GPU VRAMとシステムメモリ使用量

効果をどのように評価しますか？

IOデータプラットフォーム評価

IOデータプラットフォームはモデル推論品質評価をサポートしています。任意の実データとモデル推論結果を比較し、各関節の指示の比較を可視化できます。

その他の評価方法

以下の方法でタスクでのPi0の効果を評価することもできます：

1. 定量的評価

• アクションエラー：予測アクションと実アクションを比較（角度エラー、位置エラーなど）
• 軌跡比較：予測vs ground truth行動軌跡を可視化
• 損失指標：検証セットでの損失変化トレンドを記録

2. 定性的評価

• 実機展開：実ロボットでタスク成功率を検証
• 行動観察：モデル生成アクションが期待通りかどうかを観察
• ユーザーフィードバック：実際の使用からのユーザーフィードバックを収集

3. 評価スクリプト例

import numpy as np
from lerobot.policies import Pi0Policy

# モデルを読み込み
policy = Pi0Policy.from_pretrained("./checkpoints/pi0_finetuned")

# 評価関数
def evaluate_model(policy, test_dataset):
    total_error = 0
    num_samples = 0
    
    for batch in test_dataset:
        predicted_action = policy.select_action(batch)
        true_action = batch['actions']
        
        # アクションエラーを計算
        error = np.mean(np.abs(predicted_action - true_action))
        total_error += error
        num_samples += 1
    
    return total_error / num_samples

# 評価を実行
avg_error = evaluate_model(policy, test_dataset)
print(f"平均アクションエラー: {avg_error:.4f}")

関連リソース

• LeRobotプロジェクトホームページ
• Hugging Face Pi0モデル
• オリジナルPi0実装（JAX）
• Pi0論文
• LeRobotドキュメント

よくある質問

Q: ファインチューニング中にモデルがembed_tokensが存在しないとエラーを報告しますか？

A: これはtransformersのバージョンが古いためです。アップグレードしてください：

pip install -U transformers

Q: データセットがHugging Faceにアップロードされていませんが、ローカルデータを使用できますか？

A: はい、--dataset.root=your_local_pathパラメータを使用してください。

Q: 学習中にVRAMが不足した場合はどうすればよいですか？

A: 以下の方法を試してください：

# 1. バッチサイズを減らす
--batch_size=1

実際の状況に基づいて決定する必要がある他のパラメータもあります。

Q: 学習時間は通常どのくらいかかりますか？

A: 学習時間は以下に依存します：

• データ量：通常10時間以上
• ハードウェア構成：GPU性能が大きく影響
• 学習エポック数：5-20エポックを推奨

Q: 学習が収束したかどうかをどのように判断しますか？

A: 以下の指標を観察してください：

• 損失曲線：学習損失と検証損失が安定
• 検証指標：検証セットでの性能が向上しなくなる
• 過学習：検証損失が上昇し始めたら学習を停止

Q: ファインチューニングされたモデルをどのように展開しますか？

A: IOデータプラットフォームは既にPi0やSmolVLA、その他のロボティクス分野で一般的なモデルの自動展開をサポートしています。詳細は技術専門家にお問い合わせください。

以下の手順も参考にできます：

# 1. モデルを読み込み
policy = Pi0Policy.from_pretrained("./checkpoints/pi0_finetuned")

# 2. 入力データを準備
observation = {
    'images': camera_images,
    'language_instruction': "赤いペンを取ってください"
}

# 3. アクションコマンドを生成
action = policy.select_action(observation)

学習が順調に進むことをお祈りしています！