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

policy.py

Model

model.py

); shared

Config

FromConfig

types come from Runtime (

physicalai.config

) — see Runtime docs

use-from-config

and

configuration

Workflow

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.

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.

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.

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.

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.

Add a training config in

library/configs/physicalai/<name>.yaml

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

Done when:

physicalai fit --config configs/physicalai/<name>.yaml --trainer.fast_dev_run=true

completes one step.

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.
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 pytest tests/unit/policies -k <name>
```
  passes.
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 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.

physicalai-train-adding-a-policy

NPX Install

Tags

SKILL.md Content

Adding a Studio Policy

Workflow

Required checks

Verify

References