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）と、2 つのレイアウト間の双方向フィールドマッピング。
• 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 でデータセットを選択し、その特徴（画像キー、state の次元、action の次元）が対象ポリシーの Config と一致することを確認します。
   • 完了条件：ポリシーが想定する Feature 名と action の次元がデータセットと一致していること。
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
   

   • 完了条件：特徴の不足エラーがなく、正しい形状でバッチが 1 つ流れること。
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、action の次元、正規化が一致していること。
• 変換を往復しても、フィールドが失われたり名前が変更されたりしないこと。
• データモジュールを直接 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 に揃えます。