physicalai-train-training-a-policy

Training a policy (library)

Training uses physicalai.train.Trainer (library/src/physicalai/train/trainer.py, a lightning.Trainer subclass) with a Policy and a DataModule. The library deliberately supports two equal entry points:

• CLI — physicalai fit (and validate, test, predict): jsonargparse YAML under library/configs/, overrides on the command line; checkpoints under experiments/{name}/version_N/ by default. See library/docs/how-to/training/cli.md.
• Python API — construct Policy, LeRobotDataModule (or another datamodule), and Trainer, then trainer.fit(model=policy, datamodule=datamodule) (and validate / test / predict with a checkpoint as needed). See library/docs/getting-started/quickstart.md and library/docs/explanation/trainer/README.md.

The CLI subcommands and the Python API share the same objects; YAML class_path / init_args should match what you would wire in code.

The four CLI subcommands share the same --model / --data / --trainer.* shape (see cli/_dispatch.py); validate/test/predict additionally take --ckpt_path. When a task is about library behavior rather than shell usage, prefer the Python API path first and then verify CLI parity if the change is user-facing.

Anatomy of a config

A config wires three pieces via class_path / init_args:

• model — a Policy subclass (e.g. physicalai.policies.ACT).
• data — a DataModule, usually physicalai.data.lerobot.LeRobotDataModule with a repo_id (e.g. lerobot/pusht).
• trainer — Lightning args (max_epochs, accelerator, devices, callbacks…).

First-party configs live under library/configs/physicalai/<policy>/<embodiment>/; LeRobot-wrapped configs live in library/configs/lerobot/. Compose with __base__ and override any field on the CLI (--trainer.max_epochs 200 --data.train_batch_size 64).

Python API workflow

Use this path when the user asks for code, notebooks, tests, direct library integration, or changes to Trainer, Policy, or datamodules.

from physicalai.data import LeRobotDataModule
from physicalai.policies import ACT
from physicalai.train import Trainer

datamodule = LeRobotDataModule(repo_id="lerobot/pusht", train_batch_size=2)
policy = ACT()
trainer = Trainer(fast_dev_run=True)
trainer.fit(model=policy, datamodule=datamodule)

1. Construct the same objects the CLI would instantiate: a Policy, a DataModule, and Trainer.
   • Done when: construction works without relying on jsonargparse YAML.
2. Smoke-test the API wiring with Trainer(fast_dev_run=True).
   • Done when: one train + one val batch complete without shape or feature errors.
3. Validate / test / predict from Python with the corresponding Trainer method and ckpt_path when needed.
   • Done when: the API call and the equivalent CLI command agree on checkpoint/config behavior.

CLI workflow

Use this path when the user asks for terminal commands, docs under library/docs/how-to/, YAML configs, reproducible experiments, or entry-point behavior.

1. Start from an existing config matching your policy family; copy it rather than writing from scratch.
   • Done when: physicalai fit --config <your.yaml> --print_config renders the fully-resolved config with no errors.
2. Smoke-test the wiring before a real run:

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

   • Done when: one train + one val batch complete without shape or config errors.
3. Run training, overriding on the CLI as needed:

   physicalai fit --config configs/physicalai/<policy>/<embodiment>/<config>.yaml --trainer.max_epochs 200
   

   • Done when: checkpoints appear under experiments/{name}/version_N/.
4. Validate / test / predict from a checkpoint:

   physicalai validate --config configs/physicalai/<policy>/<embodiment>/<config>.yaml --ckpt_path experiments/<name>/version_0/checkpoints/last.ckpt
   

5. Iterate on metrics, not just loss — confirm the val metric relevant to the task moves, and record the config + checkpoint that produced it.

Debugging a run

• API: construct Policy, DataModule, and Trainer directly in a short script or test to isolate whether failure is in object construction, dataloading, or CLI parsing.
• --trainer.fast_dev_run=true — one batch each stage; the first thing to try on any failure.
• --print_config — see the exact resolved config jsonargparse built.
• Shape/feature mismatches usually mean the datamodule's Feature names or action dim disagree with the policy — cross-check against the physicalai-train-adding-a-policy skill.
• Dataset download stalls: the run is pulling a LeRobot repo_id; see the physicalai-train-working-with-datasets skill.

Required checks

• Config resolves (--print_config) and fast_dev_run passes before any long run.
• The equivalent Python API construction path passes for library-facing changes.
• accelerator/devices match the installed backend extra (xpu/cuda/cpu).
• New or renamed config fields stay consistent with the policy's Config class.
• Doc code blocks that show training commands still pass tests/test_docs.py.

Verify

# from library/
physicalai fit --config configs/physicalai/<policy>/<embodiment>/<config>.yaml --trainer.fast_dev_run=true
uv run --no-sync pytest tests/unit/train

For API-facing changes, add or run an equivalent Python smoke test (not a shell heredoc) that constructs Policy, DataModule, and Trainer directly and calls trainer.fit(...).

Related skills

• physicalai-train-adding-a-policy — when the model itself needs changes.
• physicalai-train-working-with-datasets — for the data half of the config.
• physicalai-train-benchmarking-a-policy — to evaluate a trained checkpoint in a gym.