physicalai-train-working-with-datasets

使用 Studio 数据集

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（StrEnum：physicalai、lerobot），以及两种布局之间的双向字段映射。
• data/observation.py — Observation、Feature、FeatureType、NormalizationParameters。
• data/datamodules.py — 基础 DataModule（Lightning LightningDataModule，自动选择 num-workers 的启发式逻辑）。
• data/dataset.py — 基础 Dataset；data/gym.py — 用于 Gym 生成数据的 GymDataset。

Python API 用法

在笔记本、测试、直接检查批次或调试数据加载时使用此路径，无需调用训练 CLI。

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()))

完成条件：批次包含策略所需的 observation/action 字段，且批次维度和动作维度符合预期。

将数据接入训练配置

在 physicalai fit 配置中，data 块用于选择数据模块及其 repo_id：

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 技能。

工作流

1. 选择数据集：根据 repo_id 选择数据集，并确认其特征（图像键、状态维度、动作维度）与目标策略的 Config 匹配。
   • 完成条件：策略所需的 Feature 名称和动作维度与数据集一致。
2. 训练前通过 Python API 验证一个批次：

   datamodule.prepare_data()
   datamodule.setup("fit")
   batch = next(iter(datamodule.train_dataloader()))
   

   • 完成条件：批次的键和形状正确，且无需调用 CLI。
3. 通过 YAML 配置数据集时验证 CLI 一致性：

   physicalai fit --config <config.yaml> --trainer.fast_dev_run=true
   

   • 完成条件：一个批次能够正常通过，形状正确且没有缺少特征的错误。
4. 仅在需要时通过 converters.py 转换布局（DataFormat.physicalai ↔ DataFormat.lerobot）；保持字段名称稳定，因为它们会传递到训练和导出流程中。
5. 通过 NormalizationParameters/Feature 设置归一化，并确保与策略推理时的预期一致。

调试数据加载

• 缺少或重命名特征 → 配置中的数据集特征与策略不一致；按照 data/observation.py 中的约定对齐 Feature 名称。
• 首个批次缓慢或卡住 → LeRobot repo_id 正在下载；首次运行时这是预期情况（需要下载的测试请参阅 requires_download 测试标记）。
• 批次维度错误 → 修改策略前，先检查 train_batch_size 和数据模块的 collate/observation 处理逻辑。

必需检查

• 数据集、Config 和任何导出元数据中的特征名称、FeatureType、动作维度和归一化设置一致。
• 转换可往返进行，且不会丢弃或重命名字段。
• 直接构造数据模块 API 和通过 YAML 配置构造数据模块所生成的批次兼容。
• 需要下载的测试标记为 requires_download；确保默认的 uv run pytest 可在离线状态下运行。

验证

# from library/
uv run pytest tests/unit/data tests/unit/datamodules

相关技能

• physicalai-train-training-a-policy — data 块是训练配置的一部分。
• physicalai-train-adding-a-policy — 将 observation 特征与策略 Config 对齐。