Skip to content

[python] Support custom dataset readers in PaimonLeRobotDataset - #9785

Merged
JingsongLi merged 16 commits into
apache:masterfrom
XiaoHongbo-Hope:codex/lerobot-indexed-reader
Sep 14, 2026
Merged

[python] Support custom dataset readers in PaimonLeRobotDataset#9785
JingsongLi merged 16 commits into
apache:masterfrom
XiaoHongbo-Hope:codex/lerobot-indexed-reader

Conversation

@XiaoHongbo-Hope

@XiaoHongbo-Hope XiaoHongbo-Hope commented Sep 13, 2026

Copy link
Copy Markdown
Contributor

Purpose

Let PaimonLeRobotDataset read logical frames through custom dataset readers.

Changes

  • Add PaimonDatasetReader: metadata, logical frame schema, validated batched row reads, delta windows, media decoding, and Torch conversion.
  • Make PaimonLeRobotDataset a Dataset facade over PaimonDatasetReader.
  • Accept a standard Paimon table or a configured PaimonDatasetReader through the same constructor.

Custom readers implement read_indices. schema enables eager validation; file_io is required only for BLOB or video descriptors. Reader instances must be picklable for DataLoader workers.

Relationship to LeRobot

LeRobot main selects a BaseDatasetReader by meta.storage_format. The reader owns row access, delta windows, media decoding, and returns assembled samples.

The BaseDatasetReader source docstring explicitly lets subclasses define their own constructors because storage inputs differ. schema and file_io are PyPaimon-specific inputs, not LeRobot parameters.

# LeRobot main
LeRobotDataset
└─ BaseDatasetReader
   ├─ DatasetReader          (Parquet / MP4)
   └─ LanceDatasetReader     (Lance)

# This PR
PaimonLeRobotDataset
└─ PaimonDatasetReader
   ├─ _PaimonTableDatasetReader   (standard table group)
   └─ CustomDatasetReader         (custom storage layout)

# Future LeRobot backend
LeRobotDataset
└─ PaimonDatasetReader       (BaseDatasetReader)

PaimonDatasetReader mirrors the official Reader responsibilities. Once the official contract is available in PyPaimon's supported LeRobot range, it can implement BaseDatasetReader without changing custom row-resolution code.

Official references: BaseDatasetReader, storage_format reader lookup.

Review focus

Is the minimum custom-reader contract (read_indices, plus optional schema and file_io) sufficient while keeping sample assembly in one Reader?

Tests

  • 80 passed in multimodal_lerobot_test.py.
  • Python syntax compilation, flake8 for changed files, and git diff --check pass.

@XiaoHongbo-Hope XiaoHongbo-Hope changed the title [python] Support logical frame readers for LeRobot datasets [python] Support custom frame readers in PaimonLeRobotDataset Sep 13, 2026
@XiaoHongbo-Hope XiaoHongbo-Hope changed the title [python] Support custom frame readers in PaimonLeRobotDataset [python] Support custom dataset sources in PaimonLeRobotDataset Sep 13, 2026
@XiaoHongbo-Hope XiaoHongbo-Hope changed the title [python] Support custom dataset sources in PaimonLeRobotDataset [python] Support custom dataset readers in PaimonLeRobotDataset Sep 13, 2026
@XiaoHongbo-Hope
XiaoHongbo-Hope marked this pull request as ready for review September 13, 2026 15:13

@JingsongLi JingsongLi left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewed 1cf87f1. Requirement fit: SUPPORTED. Implementation: CLEAN in the reviewed reader contracts.

Custom storage layouts can now provide logical frames while reusing the existing episode selection, delta windows and sample assembly. The boundary is small enough to be useful: read_indices supplies rows, and the shared reader validates completeness, uniqueness and control columns before assembling samples. The facade also preserves the existing table-backed path and explicitly avoids forwarding pickle special methods. The LeRobot reader reference supports this division of responsibilities; direct integration with LeRobot's backend registry remains future work.

I checked batch ordering, episode remapping, schema/media contracts and serialization ownership. The available environment passed 39 tests and additional real-Arrow custom-reader probes; 38 optional tests were skipped. Three full-file failures reproduce unchanged on the exact base. Torch/media/DataLoader end-to-end execution was unavailable locally; current PR CI is successful. No actionable introduced regression found.

@JingsongLi
JingsongLi merged commit 3723288 into apache:master Sep 14, 2026
13 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants