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.