physicalai-train-working-with-datasets

Compare original and translation side by side

🇺🇸

Original

English

🇨🇳

Translation

Chinese

Working with Studio Datasets

处理Studio数据集

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.

Studio数据存储于

library/src/physicalai/data/

目录下。数据集采用LeRobot格式，并通过Lightning数据模块调用。这些数据模块是一等Python API对象；YAML/CLI配置是同一构建路径的序列化形式。

核心模块：

data/lerobot/datamodule.py

—

LeRobotDataModule

（配置文件中引用的类为

physicalai.data.lerobot.LeRobotDataModule

）。

```
data/lerobot/dataset.py
```
— LeRobot数据集封装器。
```
data/lerobot/converters.py
```
—
```
DataFormat
```
（字符串枚举类型：
```
physicalai
```
、
```
lerobot
```
）以及两种格式间的双向字段映射。

data/observation.py

—

Observation

、

Feature

、

FeatureType

、

NormalizationParameters

。

```
data/datamodules.py
```
— 基础
```
DataModule
```
（基于Lightning
```
LightningDataModule
```
，自动计算工作进程数的启发式实现）。
```
data/dataset.py
```
— 基础
```
Dataset
```
；
```
data/gym.py
```
— 用于gym生成数据的
```
GymDataset
```
。

Python API usage

Python API 使用方法

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.

在Notebook、测试、直接批处理检查，或不涉及训练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()))

完成标志：批处理包含策略预期的观测/动作字段，且批处理/动作维度符合预期。

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.

在

physicalai fit

配置中，

data

块用于选择数据模块及其

repo_id

：

yaml

data:
  class_path: physicalai.data.lerobot.LeRobotDataModule
  init_args:
    repo_id: lerobot/pusht
    train_batch_size: 64

repo_id

指向LeRobot/HuggingFace数据集；数据模块会在首次使用时拉取该数据集。完整配置可参考

physicalai-train-training-a-policy

技能。

Workflow

工作流程

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.
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.
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.
Convert layouts only when needed via
```
converters.py
```
(
```
DataFormat.physicalai
```
↔
```
DataFormat.lerobot
```
); keep field names stable, since they propagate to training and export.
Set normalization through
```
NormalizationParameters
```
/
```
Feature
```
consistently with what the policy expects at inference.

通过
repo_id
选择数据集，确认其特征（图像键、状态维度、动作维度）与目标策略的
```
Config
```
匹配。
- 完成标志：策略预期的
```
Feature
```
  名称和动作维度与数据集一致。
训练前通过Python API验证批处理：
python
```
datamodule.prepare_data()
datamodule.setup("fit")
batch = next(iter(datamodule.train_dataloader()))
```
- 完成标志：无需调用CLI即可获取具有正确键和形状的批处理。
验证YAML配置下的CLI一致性：
bash
```
physicalai fit --config <config.yaml> --trainer.fast_dev_run=true
```
- 完成标志：一批数据可正常流转，形状正确且无特征缺失错误。
仅在必要时通过
converters.py
转换格式（
```
DataFormat.physicalai
```
↔
```
DataFormat.lerobot
```
）；保持字段名称稳定，因为它们会传播到训练和导出环节。
通过
NormalizationParameters
/
Feature
设置归一化，确保与策略推理时的预期一致。

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.

特征缺失/重命名 → 配置中的数据集特征与策略不匹配；按照
```
data/observation.py
```
的约定对齐
```
Feature
```
名称。
首个批处理加载缓慢/停滞 → LeRobot
```
repo_id
```
正在下载；首次运行时属于正常现象（需要下载的测试会标记
```
requires_download
```
）。
批处理维度错误 → 在修改策略前，检查
```
train_batch_size
```
以及数据模块的整理/观测处理逻辑。

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 pytest
```
runnable offline.

数据集、
```
Config
```
以及任何导出元数据之间的特征名称、
```
FeatureType
```
、动作维度和归一化设置保持一致。
格式转换可往返进行，不会丢失或重命名字段。
直接通过数据模块API构建和通过YAML配置构建可生成兼容的批处理。
需要下载的测试标记
```
requires_download
```
；确保默认的
```
uv run pytest
```
可离线运行。

Verify

验证

bash

undefined