physicalai-train-adding-a-policy

Compare original and translation side by side

🇺🇸

Original

English

🇨🇳

Translation

Chinese

Adding a Studio Policy

添加Studio策略

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

策略存放在

library/src/physicalai/policies/<name>/

目录下。每个策略家族都是面向Lightning的

Policy

，封装了

torch.nn.Module

类型的

Model

，拆分到三个文件中。基类位于

policies/base/

（

policy.py

中的

Policy

、

model.py

中的

Model

）；通用的

Config

FromConfig

类型来自Runtime（

physicalai.config

）——可查看Runtime文档

use-from-config

和

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.

先参考现有策略家族。研究
```
policies/pi05/
```
（当前参考实现）：包含
```
config.py
```
（
```
Pi05Config(Config)
```
）、
```
model.py
```
（
```
Pi05Model(Model)
```
）、
```
policy.py
```
（
```
Pi05(ExportablePolicyMixin, Policy)
```
）、
```
preprocessor.py
```
，以及该架构所需的任何额外模块（如
```
pi_gemma.py
```
）。如果需要极简的策略家族，
```
policies/act/
```
是不含VLM栈的小型三文件结构。
- 完成标志：能够明确每个新文件对应的现有文件。

在
policies/<name>/
目录下创建三文件结构：

```
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

可正常导入。

实现基类
Policy
定义的、同时适用于训练和推理的策略接口：
- ```
forward(...)
```
  —— 训练路径；返回值需与
```
training_step
```
  兼容。
- ```
predict_action_chunk(...)
```
  —— 推理路径；返回符合配置动作时间跨度的张量。
- ```
select_action(...)
```
  —— 除非需要特殊流程，否则使用基类的动作队列逻辑。
- 完成标志：形状与下文合成批次的检查要求匹配。
注册策略家族，以便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>
```
  可正确解析。
在添加CLI配置前，先验证直接API构建：
python
```
from physicalai.policies import get_policy

policy = get_policy("<name>")
```
- 完成标志：直接构建、配置往返、合成
```
forward(...)
```
  /
```
predict_action_chunk(...)
```
  的形状检查均通过。

当策略支持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

可完成一个步骤。

仅当准备就绪时配置导出功能。添加
```
ExportablePolicyMixin
```
和有效的示例输入，然后遵循
```
physicalai-train-exporting-and-validating
```
技能流程。如果有意不支持导出，请在策略的文档字符串中明确说明。
在
library/tests/unit/policies/
目录下添加测试用例，与现有策略测试放在一起：至少包含一个构建/配置路径测试和一个形状验证测试。
- 完成标志：
```
uv run pytest tests/unit/policies -k <name>
```
  执行通过。
如果策略对用户可见，更新文档：更新
```
library/docs/explanation/policy/
```
中的内容以及任何配置/API示例。

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.

需逐一确认以下所有项（不能仅靠“看起来正常”判断）：

动作形状语义——batch、horizon/chunk length和动作维度需符合策略家族的约定，且保持一致。
观测特征——特征名称需与数据集/配置约定对齐（
```
data/observation.py
```
中的
```
Feature
```
、
```
FeatureType
```
）。
API构建路径——导入、
```
get_policy(...)
```
、直接构造函数调用以及合成形状检查无需CLI参与即可通过。
配置路径——当策略支持CLI访问时，通过
```
physicalai fit
```
使用的jsonargparse CLI路径（
```
class_path
```
/
```
init_args
```
）可正常构建策略。
重依赖处理——在
```
library/pyproject.toml
```
中将大型策略家族设为可选依赖，并延迟导入，与
```
pi05
```
/
```
pi0
```
/
```
groot
```
/
```
smolvla
```
的处理方式一致。
无静默契约变更——未经与导出/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/

在

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.

```
references/base-classes.md
```
——
```
Policy
```
/
```
Model
```
契约及文件拆分要求。",