physicalai-train-adding-a-policy

Studio ポリシーの追加

ポリシーは library/src/physicalai/policies/<name>/ にあります。各ファミリーは、3 つのファイルに分割された torch.nn.Module の Model をラップする Lightning 対応の Policy です。基底クラスは policies/base/ にあります（policy.py の Policy、model.py の Model）。共有の Config／FromConfig 型は Runtime（physicalai.config）から取得します。Runtime のドキュメント use-from-config と configuration を参照してください。

ワークフロー

1. まず、近いファミリーを読みます。 現在のリファレンス実装である policies/pi05/ を確認します。config.py（Pi05Config(Config)）、model.py（Pi05Model(Model)）、policy.py（Pi05(ExportablePolicyMixin, Policy)）、preprocessor.py、およびアーキテクチャに必要な追加モジュール（例：pi_gemma.py）を調べてください。意図的に最小構成にするファミリーでは、VLM スタックを含まない、より小規模な 3 ファイル構成の policies/act/ を参考にできます。

   • 完了条件：新しい各ファイルが、既存のどのファイルに対応するかを説明できます。

2. policies/<name>/ に 3 ファイル構成を作成します。

   • config.py — <Name>Config(Config)。すべてのハイパーパラメーターを型付きフィールドとして定義します。
   • model.py — <Name>Model(Model)。純粋な torch.nn.Module のロジックを実装します。
   • policy.py — <Name>(Policy)（エクスポートを実装する場合に限り、ExportablePolicyMixin を追加します）。
   • 完了条件：from physicalai.policies.<name> import <Name>, <Name>Config, <Name>Model を問題なくインポートできます。

3. 基底 Policy を通して、トレーニングと推論の両方で使われる ポリシーインターフェースを実装します。

   • forward(...) — トレーニング経路。training_step と互換性のある値を返します。
   • predict_action_chunk(...) — 推論経路。設定されたアクションホライズンに対応するテンソルを返します。
   • select_action(...) — 特殊なフローが正当化される場合を除き、基底クラスのアクションキュー動作を使用します。
   • 完了条件：合成バッチで、以下のチェックと形状が一致します。

4. API と CLI の両方のユーザーがファミリーを利用できるように 登録します。

   • policies/__init__.py にエクスポートを追加します（__all__ とインポート。例：<Name>、<Name>Config、<Name>Model）。
   • policies/__init__.py の get_physicalai_policy_class(...)／get_policy(...) ディスパッチに、小文字の名前を追加します。
   • 完了条件：from physicalai.policies import <Name>, get_policy が機能し、get_policy("<name>") がインスタンスを返し、--model physicalai.policies.<Name> が解決されます。

5. CLI 設定を追加する前に、直接 API での構築を確認します。

   from physicalai.policies import get_policy
   
   policy = get_policy("<name>")
   

   • 完了条件：直接構築、設定のラウンドトリップ、合成入力に対する forward(...)／predict_action_chunk(...) の形状チェックが通ります。

6. CLI からユーザーが利用するポリシーの場合は、library/configs/physicalai/<name>.yaml に トレーニング設定を追加します。model.class_path、data.class_path（通常は physicalai.data.lerobot.LeRobotDataModule）、および trainer.* を設定します。configs/physicalai/pi05.yaml を参考にしてください。

   • 完了条件：physicalai fit --config configs/physicalai/<name>.yaml --trainer.fast_dev_run=true が 1 ステップ完了します。

7. 準備が整った場合に限り、エクスポートを組み込みます。 ExportablePolicyMixin と有効なサンプル入力を追加し、physicalai-train-exporting-and-validating スキルの手順に従います。エクスポートを意図的にサポートしない場合は、その旨をポリシーの docstring に明記してください。

8. 既存のポリシーテストの近くにある library/tests/unit/policies/ に テストを追加します。少なくとも、構築／設定のテストを 1 つ、形状検証テストを 1 つ追加してください。

   • 完了条件：uv run pytest tests/unit/policies -k <name> が成功します。

9. ポリシーがユーザー向けの場合は ドキュメントを更新します。対象は library/docs/explanation/policy/ と、関連する設定／API の例です。

必須チェック項目

以下のすべての項目を確認してください（「問題なさそう」だけでは不十分です）。

• アクション形状のセマンティクス — バッチ、ホライズン／チャンク長、アクション次元が正しく、ファミリーの慣例から変更されていないこと。
• 観測特徴量 — 特徴量名がデータセット／設定の規約（data/observation.py：Feature、FeatureType）と一致していること。
• API の構築経路 — CLI を介さずに、インポート、get_policy(...)、コンストラクターの直接使用、合成形状チェックが成功すること。
• 設定の経路 — ポリシーが CLI から利用可能な場合、physicalai fit が使用する jsonargparse CLI 経路（class_path／init_args）で構築できること。
• 重量級の依存関係 — library/pyproject.toml のオプションの extra で大規模なファミリーを制御し、pi05／pi0／groot／smolvla と同様に遅延インポートすること。
• 明示されない契約変更を行わないこと — エクスポート／Runtime と調整せずに、アクション次元、特徴量名、前処理を変更しないこと。

検証

library/ から実行してください。

uv run pytest tests/unit/policies -k <name>
physicalai fit --config configs/physicalai/<name>.yaml --trainer.fast_dev_run=true
prek run --all-files library/

参考資料

• references/base-classes.md — Policy／Model の契約とファイル分割に関する要件。