physicalai-train-adding-a-policy

# Adding a Studio Policy Policies live in `library/src/physicalai/policies/<name>/`. Each family is a Lightning-facing `Policy` wrapping a `torch.nn.Module` `Model`, split across three files. Base classes are in `policies/base/` (`Policy` in `policy.py`, `Model` in `model.py`); shared `Config` types come from Runtime (`physicalai.config`), while class construction and CLI configuration use `jsonargparse` (`FromConfigMixin`, `class_path`, and `init_args`) — see the Runtime configuration documentation. ## Workflow 1. **Read a nearby family first.** Study `policies/pi05/` (current reference implementation): `config.py` (`Pi05Config(Config)`), `model.py` (`Pi05Model(Model)`), `policy.py` (`Pi05(ExportablePolicyMixin, Policy)`), `preprocessor.py`, and any extra modules the architecture needs (e.g. `pi_gemma.py`). For a deliberately minimal family, `policies/act/` is a smaller three-file layout without the VLM stack. - Done when: you can name which existing file each new file mirrors. 2. **Create the three-file split** in `policies/<name>/`: - `config.py` — `<Name>Config(Config)`, all hyperparameters as typed fields. - `model.py` — `<Name>Model(Model)`, pure `torch.nn.Module` logic. - `policy.py` — `<Name>(Policy)` (add `ExportablePolicyMixin` only when export is implemented). - Done when: `from physicalai.policies.<name> import <Name>, <Name>Config, <Name>Model` imports cleanly. 3. **Implement the policy interface** used by both training and inference through the base `Policy`: - `forward(...)` — training path; return values compatible with `training_step`. - `predict_action_chunk(...)` — inference path; return a tensor with the configured action horizon. - `select_action(...)` — use base-class action-queue behavior unless a specialized flow is justified. - Done when: shapes match the checks below for a synthetic batch. 4. **Register the family** so both API and CLI users can find it: - Add exports to `policies/__init__.py` (`__all__` and imports, e.g. `<Name>`, `<Name>Config`, `<Name>Model`). - Add the lowercase name to the `get_physicalai_policy_class(...)` / `get_policy(...)` dispatch in `policies/__init__.py`. - Done when: `from physicalai.policies import <Name>, get_policy` works, `get_policy("<name>")` returns an instance, and `--model physicalai.policies.<Name>` resolves. 5. **Prove direct API construction** before adding CLI config: ```python from physicalai.policies import get_policy policy = get_policy("<name>") ``` - Done when: direct construction, config round-trip, and synthetic `forward(...)` / `predict_action_chunk(...)` shape checks pass. 6. **Add a training config** under `library/configs/physicalai/<policy>/<embodiment>/` when the policy is user-facing from the CLI. Wire `model.class_path`, a `data.class_path` (usually `physicalai.data.lerobot.LeRobotDataModule`), and `trainer.*`. Mirror `configs/physicalai/pi05/aloha/default.yaml`. - Done when: `physicalai fit --config configs/physicalai/<policy>/<embodiment>/<config>.yaml --trainer.fast_dev_run=true` completes one step. 7. **Wire export only when ready.** Add `ExportablePolicyMixin` and a valid sample input, then follow the `physicalai-train-exporting-and-validating` skill. If export is intentionally unsupported, say so explicitly in the policy docstring. 8. **Add tests** under `library/tests/unit/policies/` next to existing policy tests: at least one construction/config path and one shape-validation test. - Done when: `uv run --no-sync pytest tests/unit/policies -k <name>` passes. 9. **Update docs** if the policy is user-visible: `library/docs/explanation/policy/` and any config/API examples. ## Required checks Account for every item below (not just "looks fine"): - **Action shape semantics** — batch, horizon/chunk length, and action dimension are correct and unchanged from the family's convention. - **Observation features** — feature names align with dataset/config conventions (`data/observation.py`: `Feature`, `FeatureType`). - **API construction path** — imports, `get_policy(...)`, direct constructor use, and synthetic shape checks pass without CLI involvement. - **Config path** — construction works through the jsonargparse CLI path used by `physicalai fit` (`class_path`/`init_args`) when the policy is CLI-visible. - **Heavy dependencies** — gate large families behind an optional extra in `library/pyproject.toml` and import lazily, matching `pi05`/`pi0`/`groot`/`smolvla`. - **No silent contract changes** — do not alter action dims, feature names, or preprocessing without coordinating export/Runtime. ## Verify From `library/`: ```bash uv run --no-sync pytest tests/unit/policies -k <name> physicalai fit --config configs/physicalai/<policy>/<embodiment>/<config>.yaml --trainer.fast_dev_run=true prek run --all-files library/ ``` ## References - `references/base-classes.md` — the `Policy`/`Model` contract and file-split expectations.