Skip to content

[python] Add LeRobot Dataset v3 import - #9446

Merged
JingsongLi merged 19 commits into
apache:masterfrom
XiaoHongbo-Hope:codex/lerobot-import
Aug 31, 2026
Merged

[python] Add LeRobot Dataset v3 import#9446
JingsongLi merged 19 commits into
apache:masterfrom
XiaoHongbo-Hope:codex/lerobot-import

Conversation

@XiaoHongbo-Hope

@XiaoHongbo-Hope XiaoHongbo-Hope commented Aug 28, 2026

Copy link
Copy Markdown
Contributor

Purpose

Add a one-time image-based LeRobot v3 importer:

LeRobot Dataset v3 -> Arrow batches -> Paimon table

snapshot_id = conn.load_from_lerobot(
    "robot_data",
    "/path/to/lerobot_dataset",
)
  • Supports local directories, FileIO URIs such as OSS through source_options, and Hugging Face repo_id.
  • Derives the schema from meta/info.json, preserves feature names and Episode boundaries, writes Arrow batches, and creates one snapshot.
  • Maps scalars to scalar types, vectors to VECTOR, higher-rank tensors to nested ARRAY, and compressed images to BLOB.
  • Reuses the HDF5 structural Arrow validation and adds strict LeRobot dtype, shape, component-name, BLOB, and VECTOR compatibility checks for existing tables.
  • Keeps LeRobot optional through pypaimon[lerobot] and rejects v2.x clearly.

Remote FileIO sources are read directly without materializing the dataset locally. Episode metadata projects only the five routing columns; the current implementation loads one active data Parquet file into memory before slicing it into batches.

This PR supports image features only. Video features are rejected with a clear error; packed video import will be a separate follow-up after its storage model is finalized. This PR does not depend on packed video support.

Tests

  • Official LeRobot API creates a two-Episode v3 image dataset in a temporary directory without network access.
  • Covers metadata-derived schema, scalar/vector/tensor values, ordering, image BLOB round trip, append, snapshot, shape/dtype/component-name/BLOB/VECTOR incompatibility, safe numeric conversion, empty datasets, OSS reads, missing dependency, v2 rejection, and video rejection.
  • multimodal_lerobot_test.py: 17 passed.
  • Flake8, py_compile, and git diff --check: passed.

@XiaoHongbo-Hope
XiaoHongbo-Hope marked this pull request as ready for review August 30, 2026 23:32
raise ValueError(
"LeRobot string feature %s must be scalar." % name)
arrow_type = _tensor_type(scalar_type, shape)
description = "LeRobot dtype=%s, shape=%s" % (dtype, list(shape))

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.

[P1] Preserve LeRobot component names in the append contract. This description records only dtype and shape, so two valid features such as names=["x", "y"] and names=["y", "x"] produce identical schemas and _validate_lerobot_schema accepts the append. The loader then writes vector positions unchanged, silently mixing different component semantics in one training table. Please persist a canonical representation of feature["names"] in the field metadata and compare it for existing tables; add a regression test with reordered names.

rows = []
for path in paths:
rows.extend(_read_remote_parquet(
self._file_io, path).to_pylist())

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.

[P1] Project the episode metadata columns before materializing rows. _read_remote_parquet currently reads every column and to_pylist() retains all of them for the full import. Normal LeRobot v3 episode rows include per-feature stats/* arrays, while this class only uses five scalar lookup fields. Large FileIO datasets can therefore consume GBs of Python heap and OOM before the first data batch. Please read only episode_index, dataset_from_index, dataset_to_index, data/chunk_index, and data/file_index.

_validate_info_paths(local_info)
_schema_from_info(local_info, include_task=False)
LeRobotDataset = _import_lerobot_dataset()
dataset = _open_resolved_dataset(

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.

[P2] Handle an empty local dataset before opening LeRobot. A valid zero-frame directory created by LeRobotDataset.create() has meta/info.json but no data/episode/task Parquet files. Opening it here makes LeRobot treat the missing files as a cache miss and attempt a Hub download using the synthetic local/pypaimon-import repo ID, so the documented None no-op path at line 108 is never reached. Please create or validate the table from local_info, then return None before constructing LeRobotDataset when total_frames == 0.

else:
values = [_normalize_value(value, feature, name)
for value in values]
arrays.append(pa.array(values, type=field.type))

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.

[P2] Validate source values before coercing them to the target Arrow type. pa.array(values, type=field.type) is not a safe cast: for example, a Parquet float value 1.5 under int32 metadata becomes 1, and an oversized float becomes inf under float32. Because strict_arrow_table runs afterwards, it sees only the already-coerced target type and cannot reject the mismatch. Please validate the raw Arrow element type/range or use an explicitly safe cast before rebuilding the column.

@JingsongLi

Copy link
Copy Markdown
Contributor

+1

@JingsongLi
JingsongLi merged commit d557d05 into apache:master Aug 31, 2026
11 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