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 / FromConfig types come from Runtime (physicalai.config) — see docs/how-to/config/use-from-config.md and docs/explanation/configuration.md in the physicalai repo.

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.

1. 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.

1. Implement the policy interface used by both training and inference through the base Policy:

- forward(...) — training path; return values compatible with trainingstep. - predictactionchunk(...) — inference path; return a tensor with the configured action horizon. - selectaction(...) — use base-class action-queue behavior unless a specialized flow is justified. - Done when: shapes match the checks below for a synthetic batch.

1. 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 getphysicalaipolicyclass(...) / getpolicy(...) dispatch in policies/init.py. - Done when: from physicalai.policies import <Name>, getpolicy works, getpolicy("<name>") returns an instance, and --model physicalai.policies.<Name> resolves.

1. 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(...) / predictactionchunk(...) shape checks pass.

1. Add a training config in library/configs/physicalai/<name>.yaml when the policy is user-facing from the CLI. Wire model.classpath, a data.classpath (usually physicalai.data.lerobot.LeRobotDataModule), and trainer.*. Mirror configs/physicalai/pi05.yaml.

- Done when: physicalai fit --config configs/physicalai/<name>.yaml --trainer.fastdevrun=true completes one step.

1. 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.
2. 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.

1. 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 (classpath/initargs) 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/:

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

References

• references/base-classes.md — the Policy/Model contract and file-split expectations.