Skip to content

[python] Add Paimon LeRobot map-style dataset - #9498

Draft
XiaoHongbo-Hope wants to merge 5 commits into
apache:masterfrom
XiaoHongbo-Hope:codex/paimon-lerobot-dataset
Draft

[python] Add Paimon LeRobot map-style dataset#9498
XiaoHongbo-Hope wants to merge 5 commits into
apache:masterfrom
XiaoHongbo-Hope:codex/paimon-lerobot-dataset

Conversation

@XiaoHongbo-Hope

Copy link
Copy Markdown
Contributor

Purpose

Add a map-style PaimonLeRobotDataset for training from LeRobot v3 image datasets that were imported into a multimodal Paimon table.

Changes

  • Accept matching LeRobotDatasetMetadata or a local v3 dataset/meta path and expose it as dataset.meta.
  • Reuse the lazy row-ID Torch reader from [python][torch] Make map-style reads lazy #9486, including batched __getitems__ and multi-worker DataLoader reads, without materializing the full table.
  • Read image BLOBs in coalesced batches and return Torch tensors.
  • Preserve episode order and support episode subsets, delta windows, padding masks, and image transforms.
  • Validate the table against the metadata dtype, shape, BLOB, VECTOR, and task contract.
  • Keep exact top-level LeRobot field names containing dots on the lazy reader path; only true nested projection still falls back.

Scope

This first version reuses external LeRobot v3 metadata; the imported table does not persist the complete meta/ directory and is not presented as a self-contained round-trip format. It supports image-backed datasets. Video-backed reads are an independent follow-up to #9494, and this PR does not depend on that PR.

PaimonLeRobotDataset is a minimal PyTorch Dataset-compatible wrapper rather than a subclass of the official LeRobotDataset.

Tests

  • 53 passed in pypaimon/tests/torch_read_test.py, including DataLoader(batch_size=2, num_workers=2).
  • 15 passed, 4 skipped in pypaimon/tests/multimodal_lerobot_test.py locally; LeRobot-dependent cases run in the optional-dependency CI lane.
  • Flake8, py_compile, and git diff --check pass.

@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.

Two remaining issues after re-checking the updated head.

"Paimon table is missing LeRobot fields: %s"
% sorted(missing))

self._dataset, splits, read_table, snapshot_id = _lazy_torch_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] Please reject unsupported lazy-reader fallbacks before loading payload columns. to_torch() silently calls _materialize() when the row-ID path is unavailable; a reachable example is query authorization masking _ROW_ID. In that case constructing this wrapper reads and retains the full projected table even though the documented contract says payload columns remain lazy, and enabling deltas materializes a second projection as well. This can turn dataset construction into O(table size) payload I/O and memory use on training-scale tables. Please preflight the lazy eligibility and fail with an actionable error, or provide a lazy routing path for these cases.

def _resolve_metadata(metadata):
if isinstance(metadata, (str, Path)):
root = Path(metadata)
if root.name == "meta":

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] Please disambiguate the dataset root and meta/ directory structurally instead of using the basename. A valid dataset rooted at /datasets/meta stores its metadata at /datasets/meta/meta/info.json, but this branch rewrites the root to /datasets and then rejects it. Check <input>/meta/info.json first as the dataset-root form; only if that is absent should <input>/info.json be treated as the metadata-directory form and resolved to its parent.

@XiaoHongbo-Hope
XiaoHongbo-Hope marked this pull request as ready for review August 31, 2026 18:07
@JingsongLi

Copy link
Copy Markdown
Contributor

I think the imported dataset should be self-contained in Paimon, but separate typed tables would be clearer than either duplicating metadata in every frame or storing all metadata in a generic kind + payload table.

A possible logical dataset bundle is:

  • frames: the existing per-frame data, with a stable dataset_id.
  • datasets: one row per dataset metadata version, containing the LeRobot format version, FPS, feature schema, global statistics, total counts, source information, and the associated frame-table snapshot/version.
  • episodes: one row per episode, keyed by dataset_id, metadata_version, and episode_index, with its frame range, length, and any episode-level statistics.
  • tasks: the task_index to task-text mapping, keyed by dataset_id and metadata_version.

features and global stats can remain JSON initially because their structures are format-dependent. Episodes and tasks have stable identities and relationships, so explicit columns and tables provide a much clearer contract and make them independently queryable. Video metadata could be added as another typed table when video-backed datasets are supported.

All component tables should use the same dataset_id and metadata_version. Since the tables cannot be committed atomically, datasets can act as the manifest and publication point: write the frame, episode, and task records first, then publish a completed dataset manifest which pins their corresponding versions. A reader must resolve exactly the versions referenced by that manifest and fail if any component is missing or mismatched. This avoids silently combining frame data with stale metadata.

With this contract, PaimonLeRobotDataset(table, dataset_id) becomes a thin adapter which reconstructs a LeRobot-compatible view from Paimon. The original LeRobot meta/ directory is only an import-time input and is no longer required during training.

The current PR leaves the source of truth split between Paimon frames and an external metadata directory. I suggest defining and implementing the self-contained multi-table import contract first, then keeping PaimonLeRobotDataset focused on reading that contract.

@JingsongLi

Copy link
Copy Markdown
Contributor

To clarify the business meaning of the three proposed metadata tables, they represent three different levels of the dataset rather than merely splitting one metadata blob by file type.

1. datasets: the published dataset manifest

One row represents one complete, reproducible dataset version. It should contain dataset-level information such as:

dataset_id
metadata_version
format / format_version
fps
features_json
global_stats_json
total_frames / total_episodes / total_tasks
frames_snapshot_id
episodes_snapshot_id
tasks_snapshot_id
source_uri
checksum
status

This table defines the global training contract. For example, fps determines how delta timestamps are converted to frame offsets, features_json defines the dtype and shape of every feature, and global statistics are used for state/action normalization.

More importantly, this row is the manifest for a reproducible release. A training job should be able to record something like aloha_pick_cube@3 and later reopen exactly the same frame, episode, and task versions. The importer can write the component tables first and publish the datasets row with status = READY only after validation succeeds. Readers should ignore incomplete versions.

2. episodes: the trajectory index and sampling boundaries

One row represents one robot rollout/trajectory:

dataset_id
metadata_version
episode_index
dataset_from_index
dataset_to_index
length
task_indices
split
episode_stats_json

This is not only redundant aggregation over the frame table. It is the authoritative trajectory directory used to:

  • select or shuffle complete episodes without scanning all frames;
  • prevent delta-window reads from crossing an episode boundary;
  • apply padding correctly at the beginning and end of a rollout;
  • create train/validation splits at episode granularity;
  • attach episode-level properties such as success, duration, quality, or statistics.

Although some values could be derived with GROUP BY episode_index, recomputing them whenever a dataset is opened would be expensive and would not provide a stable, versioned contract.

3. tasks: the task semantic dictionary

One row maps a stable task identifier to its human-readable instruction:

dataset_id
metadata_version
task_index
task
task_metadata_json

For example:

0 -> "Pick up the red cube"
1 -> "Place the cube in the drawer"

Frame rows can store only task_index, avoiding repeated task strings while still supporting filtering, balancing, and statistics by task. The task definition is also pinned to the same metadata version as the frames. If one episode contains multiple tasks, episodes.task_indices can preserve that relationship; a separate relation table is only needed if the model becomes more complex later.

The resulting mapping to the LeRobot API is straightforward:

meta.info      <- datasets
meta.stats     <- datasets
meta.episodes  <- episodes
meta.tasks     <- tasks
dataset[i]     <- frames

Therefore, the tables have distinct business granularity and lifecycle:

datasets: one row per published dataset version
episodes: one row per trajectory
tasks:    one row per task definition
frames:   one row per observation/frame

This makes the schema, keys, query behavior, and version relationships much clearer than a single generic metadata table, while keeping PaimonLeRobotDataset as a thin reconstruction layer.

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