physicalai-train-working-with-datasets

# Working with Studio Datasets Studio data lives in `library/src/physicalai/data/`. Datasets use the **LeRobot format** and are consumed through Lightning datamodules. The datamodules are first-class Python API objects; YAML/CLI configs are a serialization of the same construction path. Key modules: - `data/lerobot/datamodule.py` — `LeRobotDataModule` (the class configs reference as `physicalai.data.lerobot.LeRobotDataModule`). - `data/lerobot/dataset.py` — LeRobot dataset wrapper. - `data/lerobot/converters.py` — `DataFormat` (StrEnum: `physicalai`, `lerobot`) and bidirectional field mapping between the two layouts. - `data/observation.py` — `Observation`, `Feature`, `FeatureType`, `NormalizationParameters`. - `data/datamodules.py` — base `DataModule` (Lightning `LightningDataModule`, auto num-workers heuristic). - `data/dataset.py` — base `Dataset`; `data/gym.py` — `GymDataset` for gym-generated data. ## Python API usage Use this path for notebooks, tests, direct batch inspection, or debugging dataloading without involving the training CLI. ```python from physicalai.data import LeRobotDataModule datamodule = LeRobotDataModule(repo_id="lerobot/pusht", train_batch_size=2) datamodule.prepare_data() datamodule.setup("fit") batch = next(iter(datamodule.train_dataloader())) ``` Done when: the batch contains the observation/action fields the policy expects, with the expected batch/action dimensions. ## Wiring data into a training config In a `physicalai fit` config, the `data` block selects the datamodule and its `repo_id`: ```yaml data: class_path: physicalai.data.lerobot.LeRobotDataModule init_args: repo_id: lerobot/pusht train_batch_size: 64 ``` `repo_id` points at a LeRobot/HuggingFace dataset; the datamodule pulls it on first use. See the `physicalai-train-training-a-policy` skill for the full config. ## Workflow 1. **Pick the dataset** by `repo_id` and confirm its features (image keys, state dim, action dim) match the target policy's `Config`. - Done when: the policy's expected `Feature` names and action dimension line up with the dataset. 2. **Verify a batch through the Python API** before training: ```python datamodule.prepare_data() datamodule.setup("fit") batch = next(iter(datamodule.train_dataloader())) ``` - Done when: the batch has correct keys and shapes without invoking the CLI. 3. **Verify CLI parity** when the dataset is configured through YAML: ```bash physicalai fit --config <config.yaml> --trainer.fast_dev_run=true ``` - Done when: one batch flows through with correct shapes and no missing-feature errors. 4. **Convert layouts** only when needed via `converters.py` (`DataFormat.physicalai` ↔ `DataFormat.lerobot`); keep field names stable, since they propagate to training and export. 5. **Set normalization** through `NormalizationParameters`/`Feature` consistently with what the policy expects at inference. ## Debugging dataloading - Missing/renamed feature → the config's dataset features disagree with the policy; align `Feature` names in `data/observation.py` conventions. - Slow/stalled first batch → the LeRobot `repo_id` is downloading; expected on first run (see the `requires_download` test marker for tests that need this). - Wrong batch dimensions → check `train_batch_size` and the datamodule's collate/observation handling before changing the policy. - OOM or heavy swapping during training (common on smaller policies like ACT/SmolVLA on low-RAM machines) → try `pin_memory=False` and/or `persistent_workers=False` on the `DataModule`; see `library/docs/explanation/data/datamodules.md`. ## Required checks - Feature names, `FeatureType`, action dim, and normalization match between dataset, `Config`, and any export metadata. - Conversions round-trip without dropping or renaming fields. - Direct datamodule API construction and YAML config construction produce compatible batches. - Tests that require downloads are marked `requires_download`; keep default `uv run --no-sync pytest` runnable offline. ## Verify ```bash # from library/ uv run --no-sync pytest tests/unit/data tests/unit/datamodules ``` ## Related skills - `physicalai-train-training-a-policy` — the `data` block is one half of a training config. - `physicalai-train-adding-a-policy` — align observation features with the policy `Config`.