From 944e74d46b48e04fea715d55cabb359042c0d77a Mon Sep 17 00:00:00 2001 From: Yann Date: Mon, 7 Sep 2026 10:31:29 +0800 Subject: [PATCH 1/4] feat(python): add direct LeRobot capture writer Add a LeRobot-compatible writer that batches completed episodes into multimodal Paimon commits, resumes compatible tables, and stores image frames as PNG BLOBs. Co-Authored-By: Codex AI-Model: gpt-5 AI-Contributed/Feature: 424/424 AI-Contributed/UT: 366/366 --- docs/docs/pypaimon/multimodal-api.mdx | 54 +++ .../pypaimon/multimodal/arrow_utils.py | 5 + .../pypaimon/multimodal/lerobot/__init__.py | 4 +- .../pypaimon/multimodal/lerobot/loader.py | 7 +- .../pypaimon/multimodal/lerobot/writer.py | 354 +++++++++++++++++ .../tests/multimodal_lerobot_writer_test.py | 366 ++++++++++++++++++ 6 files changed, 787 insertions(+), 3 deletions(-) create mode 100644 paimon-python/pypaimon/multimodal/lerobot/writer.py create mode 100644 paimon-python/pypaimon/tests/multimodal_lerobot_writer_test.py diff --git a/docs/docs/pypaimon/multimodal-api.mdx b/docs/docs/pypaimon/multimodal-api.mdx index ddef5d1475e6..c1647a165741 100644 --- a/docs/docs/pypaimon/multimodal-api.mdx +++ b/docs/docs/pypaimon/multimodal-api.mdx @@ -694,6 +694,60 @@ together, and keep writers paused until the call returns. Scalars map to scalar types, vectors to `VECTOR`, higher-rank tensors to nested `ARRAY`, and images to `BLOB`. Images keep their compressed bytes. +## Capture LeRobot frames directly into Paimon + +`PaimonLeRobotWriter` implements the write-side surface used by LeRobot's +recording loop without first creating a LeRobot Parquet/image dataset. Pass the +same user feature mapping that would be passed to `LeRobotDataset.create`. +The writer adds the standard `timestamp`, `frame_index`, `episode_index`, +`index`, and `task_index` features itself. + +```python +from pypaimon.multimodal.lerobot import PaimonLeRobotWriter + +writer = PaimonLeRobotWriter( + conn, + "robot_data", + fps=30, + features=dataset_features, +) + +# LeRobot's record_loop only needs writer.fps, writer.features, and +# writer.add_frame(frame), so the writer can be passed as its dataset argument. +record_loop(..., fps=30, dataset=writer) +writer.save_episode() + +# Optional durability/visibility boundary before finalize. +writer.flush() +writer.finalize() +``` + +Like native LeRobot, `add_frame` requires every declared user feature plus a +string `task`, and rejects caller-provided generated fields. Numeric features +must be NumPy arrays (Torch tensors are converted) with the declared dtype and +shape. Image metadata uses `(channels, height, width)`; image values may be CHW, +HWC, or PIL. Images are encoded as PNG bytes and stored in Paimon `BLOB` +columns; no LeRobot data directory or MP4 is created. `video` features remain +unsupported. + +`save_episode` accepts the current episode and writes it to a long-lived Paimon +batch writer. The default `episodes_per_commit=-1` keeps all completed episodes +in that batch until `finalize()`. Set a positive threshold for periodic commits, +or call `flush()` at an operational boundary. `save_episode`, `flush`, and +`finalize` always return `None`. Calling `save_episode()` without any buffered +frames raises `ValueError`, matching native LeRobot. + +Call `clear_episode_buffer()` before `save_episode()` to discard a re-recorded +episode without advancing frame or episode indices. `finalize()` rejects an +unfinished episode instead of silently dropping its frames. + +The writer creates a missing table and appends to an existing compatible table. +Before writing, it requires the table columns, order, Arrow types, nullability, +and LeRobot feature metadata to match. On resume, new `index` and +`episode_index` values continue after the existing maxima, existing task +mappings are retained, and the episode-local `frame_index` starts again at zero. +A commit exception has an unknown result and is not automatically retried. + ## Overwrite `overwrite` accepts the same input formats as `add` and replaces existing data diff --git a/paimon-python/pypaimon/multimodal/arrow_utils.py b/paimon-python/pypaimon/multimodal/arrow_utils.py index 76d056ed9211..47f92452220f 100644 --- a/paimon-python/pypaimon/multimodal/arrow_utils.py +++ b/paimon-python/pypaimon/multimodal/arrow_utils.py @@ -26,6 +26,11 @@ def strict_arrow_table( source_path, batch_index, format_name): + """Validate one Arrow batch against the exact target table schema. + + Reject missing, extra, or reordered columns, validate nested nullability, + and apply only Arrow safe casts before returning a ``pyarrow.Table``. + """ if isinstance(data, pa.RecordBatch): table = pa.Table.from_batches([data]) elif isinstance(data, pa.Table): diff --git a/paimon-python/pypaimon/multimodal/lerobot/__init__.py b/paimon-python/pypaimon/multimodal/lerobot/__init__.py index a40f2a8ccef0..5f598d299160 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/__init__.py +++ b/paimon-python/pypaimon/multimodal/lerobot/__init__.py @@ -14,11 +14,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""One-time LeRobot Dataset v3 import into a multimodal Paimon table.""" +"""LeRobot Dataset v3 import and direct Paimon capture.""" from pypaimon.multimodal.lerobot.api import load_from_lerobot +from pypaimon.multimodal.lerobot.writer import PaimonLeRobotWriter __all__ = [ + "PaimonLeRobotWriter", "load_from_lerobot", ] diff --git a/paimon-python/pypaimon/multimodal/lerobot/loader.py b/paimon-python/pypaimon/multimodal/lerobot/loader.py index 295eee99c363..71719869c39e 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/loader.py +++ b/paimon-python/pypaimon/multimodal/lerobot/loader.py @@ -450,7 +450,7 @@ def _image_bytes(value, root): return _encode_media_frame(value) -def _encode_media_frame(value): +def _encode_media_frame(value, channel_first=None): try: import numpy as np from PIL import Image @@ -466,7 +466,10 @@ def _encode_media_frame(value): if callable(detach): value = detach().cpu().numpy() array = np.asarray(value) - if array.ndim == 3 and array.shape[0] in (1, 3, 4): + if channel_first is True or ( + channel_first is None + and array.ndim == 3 + and array.shape[0] in (1, 3, 4)): array = np.transpose(array, (1, 2, 0)) if np.issubdtype(array.dtype, np.floating): array = np.rint(np.clip(array, 0.0, 1.0) * 255.0).astype(np.uint8) diff --git a/paimon-python/pypaimon/multimodal/lerobot/writer.py b/paimon-python/pypaimon/multimodal/lerobot/writer.py new file mode 100644 index 000000000000..77a320696b93 --- /dev/null +++ b/paimon-python/pypaimon/multimodal/lerobot/writer.py @@ -0,0 +1,354 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""LeRobot-compatible capture writer for multimodal Paimon tables.""" + +import copy +from typing import Mapping, Optional + +import numpy as np +import pyarrow as pa + +from pypaimon.multimodal.arrow_utils import strict_arrow_table +from pypaimon.multimodal.hdf5 import _SnapshotRecorder +from pypaimon.multimodal.lerobot.loader import ( + _encode_media_frame, + _normalize_value, + _safe_array, + _value_shape, +) +from pypaimon.multimodal.lerobot.schema import ( + _feature_shape, + _schema_from_info, + _validate_lerobot_schema, +) +from pypaimon.multimodal.table import _target_schema + + +_DEFAULT_FEATURES = { + "timestamp": {"dtype": "float32", "shape": (1,), "names": None}, + "frame_index": {"dtype": "int64", "shape": (1,), "names": None}, + "episode_index": {"dtype": "int64", "shape": (1,), "names": None}, + "index": {"dtype": "int64", "shape": (1,), "names": None}, + "task_index": {"dtype": "int64", "shape": (1,), "names": None}, +} +_TASK_FEATURE = {"dtype": "string", "shape": (1,), "names": None} + + +class PaimonLeRobotWriter: + """Collect LeRobot frames and commit completed episodes to Paimon.""" + + def __init__( + self, + connection, + table_name: str, + *, + fps: int, + features: Mapping[str, Mapping[str, object]], + episodes_per_commit: int = -1, + options: Optional[Mapping[str, object]] = None): + if isinstance(fps, bool) or not isinstance(fps, int) or fps <= 0: + raise ValueError("fps must be a positive integer.") + if isinstance(episodes_per_commit, bool) \ + or not isinstance(episodes_per_commit, int) \ + or episodes_per_commit == 0 \ + or episodes_per_commit < -1: + raise ValueError( + "episodes_per_commit must be -1 or a positive integer.") + if not isinstance(features, Mapping) or not features: + raise ValueError("features must be a non-empty mapping.") + if "task" in features: + raise ValueError("task is managed by PaimonLeRobotWriter.") + + self.fps = fps + self.episodes_per_commit = episodes_per_commit + self.features = copy.deepcopy(dict(features)) + self._user_features = copy.deepcopy(dict(features)) + self.features.update(copy.deepcopy(_DEFAULT_FEATURES)) + schema_features = dict(self.features) + schema_features["task"] = _TASK_FEATURE + self._source_schema = _schema_from_info({"features": schema_features}) + self._table = connection.create_table( + table_name, + schema=self._source_schema, + options=options, + ignore_if_exists=True, + ) + self._target_schema = _target_schema(self._table.raw_table) + _validate_lerobot_schema( + self._source_schema, self._target_schema, table_name) + strict_arrow_table( + pa.Table.from_batches([], schema=self._source_schema), + self._target_schema, + table_name, + 0, + "LeRobot", + ) + + self.num_frames, self.num_episodes, self._task_indices = \ + self._load_existing_state() + self._next_task_index = ( + max(self._task_indices.values()) + 1 + if self._task_indices else 0 + ) + self.pending_episodes = 0 + self._episode_frames = [] + self._table_write = None + self._table_commit = None + self._snapshot_recorder = None + self._finalized = False + self._failed = False + + def _load_existing_state(self): + # ponytail: one resume scan; persist counters if startup cost matters. + rows = self._table.scan().select([ + "index", "episode_index", "task_index", "task" + ]).to_arrow().to_pylist() + if not rows: + return 0, 0, {} + + task_indices = {} + index_tasks = {} + for row in rows: + task = row["task"] + task_index = row["task_index"] + if ((task in task_indices + and task_indices[task] != task_index) + or (task_index in index_tasks + and index_tasks[task_index] != task)): + raise ValueError( + "Existing LeRobot task and task_index values conflict.") + task_indices[task] = task_index + index_tasks[task_index] = task + return ( + max(row["index"] for row in rows) + 1, + max(row["episode_index"] for row in rows) + 1, + task_indices, + ) + + def add_frame(self, frame): + self._require_open("add_frame") + if not isinstance(frame, Mapping): + raise ValueError("frame must be a mapping.") + expected = set(self._user_features) + actual = set(frame) - {"task"} + if actual != expected: + missing = sorted(expected - actual) + extra = sorted(actual - expected) + raise ValueError( + "LeRobot frame fields do not match features; missing=%s, " + "extra=%s." % (missing, extra)) + task = frame.get("task") + if not isinstance(task, str): + raise ValueError("LeRobot frame task must be a string.") + + values = {"task": task} + for name, feature in self._user_features.items(): + if feature.get("dtype") == "image": + value = self._image_bytes(frame[name], feature, name) + else: + value = self._normalize_frame_value( + frame[name], feature, name) + _safe_array( + [value], + self._source_schema.field(name), + name, + str(feature.get("dtype", "")), + ) + values[name] = value + self._episode_frames.append(values) + + @staticmethod + def _normalize_frame_value(value, feature, name): + numpy = getattr(value, "numpy", None) + if callable(numpy): + value = numpy() + dtype = str(feature.get("dtype", "")) + if dtype == "string": + if not isinstance(value, str): + raise ValueError( + "LeRobot feature %s must be a string." % name) + else: + if not isinstance(value, np.ndarray): + raise ValueError( + "LeRobot feature %s must be a NumPy array." % name) + if value.dtype != np.dtype(dtype): + raise ValueError( + "LeRobot feature %s expected dtype %s, got %s." + % (name, dtype, value.dtype)) + expected_shape = _feature_shape(feature, name) + if value.shape != expected_shape: + raise ValueError( + "LeRobot feature %s expected shape %s, got %s." + % (name, expected_shape, value.shape)) + return _normalize_value(value, feature, name) + + @staticmethod + def _image_bytes(value, feature, name): + actual_shape = _value_shape(value) + getbands = getattr(value, "getbands", None) + image_size = getattr(value, "size", None) + if not actual_shape and callable(getbands) \ + and isinstance(image_size, tuple) and len(image_size) == 2: + actual_shape = (image_size[1], image_size[0], len(getbands())) + expected_shape = _feature_shape(feature, name) + channel_last_shape = ( + (expected_shape[1], expected_shape[2], expected_shape[0]) + if len(expected_shape) == 3 else () + ) + if actual_shape and actual_shape != expected_shape \ + and actual_shape != channel_last_shape: + raise ValueError( + "LeRobot feature %s expected shape %s, got %s." + % (name, expected_shape, actual_shape)) + return _encode_media_frame( + value, + channel_first=actual_shape == expected_shape, + ) + + def save_episode(self): + self._require_open("save_episode") + if not self._episode_frames: + raise ValueError("Cannot save an empty LeRobot episode.") + + episode = self._episode_table() + try: + self._ensure_batch() + self._table_write.write_arrow(episode) + except BaseException: + self._fail_batch(abort=True) + raise + + self.num_frames += episode.num_rows + self.num_episodes += 1 + self.pending_episodes += 1 + self._episode_frames = [] + if self.episodes_per_commit != -1 \ + and self.pending_episodes >= self.episodes_per_commit: + self.flush() + return None + + def clear_episode_buffer(self, delete_images=True): + self._require_open("clear_episode_buffer") + self._episode_frames = [] + + def has_pending_frames(self): + return bool(self._episode_frames) + + def flush(self): + self._require_open("flush") + if self.pending_episodes == 0: + return None + commit_started = False + try: + messages = self._table_write.prepare_commit() + commit_started = True + self._table_commit.commit(messages) + snapshot_id = self._snapshot_recorder.snapshot_id + if snapshot_id is None: + raise RuntimeError( + "LeRobot batch committed without reporting a snapshot id.") + except BaseException: + self._fail_batch(abort=not commit_started) + raise + self._close_batch() + self.pending_episodes = 0 + return None + + def finalize(self): + if self._finalized: + return None + self._require_open("finalize") + if self._episode_frames: + raise RuntimeError( + "Cannot finalize with unsaved LeRobot frames; call " + "save_episode() or clear_episode_buffer() first.") + self.flush() + self._finalized = True + return None + + def _episode_table(self): + """Build one episode as a target-schema ``pyarrow.Table``.""" + episode_index = self.num_episodes + first_index = self.num_frames + size = len(self._episode_frames) + tasks = [frame["task"] for frame in self._episode_frames] + task_indices = [] + for task in tasks: + if task not in self._task_indices: + self._task_indices[task] = self._next_task_index + self._next_task_index += 1 + task_indices.append(self._task_indices[task]) + + generated = { + "timestamp": [index / self.fps for index in range(size)], + "frame_index": list(range(size)), + "episode_index": [episode_index] * size, + "index": list(range(first_index, first_index + size)), + "task_index": task_indices, + } + arrays = [] + for name, feature in self.features.items(): + values = generated.get(name) + if values is None: + values = [frame[name] for frame in self._episode_frames] + field = self._source_schema.field(name) + arrays.append(_safe_array( + values, field, name, str(feature.get("dtype", "")))) + arrays.append(pa.array(tasks, type=pa.string())) + source = pa.Table.from_arrays(arrays, schema=self._source_schema) + return strict_arrow_table( + source, + self._target_schema, + self._table.identifier, + self.num_episodes, + "LeRobot", + ) + + def _ensure_batch(self): + if self._table_write is not None: + return + builder = self._table.raw_table.new_batch_write_builder() + self._table_write = builder.new_write() + self._table_commit = builder.new_commit() + self._snapshot_recorder = _SnapshotRecorder() + self._table_commit.add_commit_callback(self._snapshot_recorder) + + def _fail_batch(self, abort): + self._failed = True + if abort and self._table_write is not None: + self._table_write.abort() + self._close_batch() + + def _close_batch(self): + try: + if self._table_write is not None: + self._table_write.close() + finally: + if self._table_commit is not None: + self._table_commit.close() + self._table_write = None + self._table_commit = None + self._snapshot_recorder = None + + def _require_open(self, method): + if self._failed: + raise RuntimeError( + "Cannot call %s() after a Paimon write failure." % method) + if self._finalized: + raise RuntimeError( + "Cannot call %s() after finalize()." % method) diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_writer_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_writer_test.py new file mode 100644 index 000000000000..1bc117772f31 --- /dev/null +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_writer_test.py @@ -0,0 +1,366 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import io +import shutil +import tempfile +import unittest +from pathlib import Path + +import numpy as np +from PIL import Image + +import pypaimon.multimodal as pmm +from pypaimon.multimodal.lerobot import PaimonLeRobotWriter + + +class PaimonLeRobotWriterTest(unittest.TestCase): + + def setUp(self): + self.temp_dir = Path(tempfile.mkdtemp( + prefix="pypaimon_lerobot_writer_")) + self.connection = pmm.connect(options={ + "warehouse": str(self.temp_dir / "warehouse"), + }) + + def tearDown(self): + shutil.rmtree(self.temp_dir, ignore_errors=True) + + def test_default_commits_only_on_finalize_and_returns_none(self): + writer = PaimonLeRobotWriter( + self.connection, + "finalize_only", + fps=10, + features={ + "action": { + "dtype": "float32", + "shape": (1,), + "names": None, + }, + }, + ) + + self.assertEqual(-1, writer.episodes_per_commit) + for value in (1.0, 2.0): + writer.add_frame({ + "action": np.array([value], dtype=np.float32), + "task": "pick", + }) + self.assertIsNone(writer.save_episode()) + table = self.connection.get_table("finalize_only") + self.assertIsNone( + table.raw_table.snapshot_manager().get_latest_snapshot()) + + self.assertIsNone(writer.finalize()) + self.assertEqual( + 1, table.raw_table.snapshot_manager().get_latest_snapshot().id) + + def test_existing_table_resumes_global_and_task_indices(self): + features = { + "action": { + "dtype": "float32", + "shape": (1,), + "names": None, + }, + } + first = PaimonLeRobotWriter( + self.connection, "resume", fps=10, features=features) + for value in (1.0, 2.0): + first.add_frame({ + "action": np.array([value], dtype=np.float32), + "task": "pick", + }) + first.save_episode() + first.finalize() + + resumed = PaimonLeRobotWriter( + self.connection, "resume", fps=10, features=features) + self.assertEqual(2, resumed.num_frames) + self.assertEqual(1, resumed.num_episodes) + resumed.add_frame({ + "action": np.array([3.0], dtype=np.float32), + "task": "pick", + }) + resumed.add_frame({ + "action": np.array([4.0], dtype=np.float32), + "task": "place", + }) + resumed.save_episode() + resumed.finalize() + + rows = self.connection.get_table("resume").scan().select([ + "episode_index", "frame_index", "index", "task_index", "task" + ]).to_arrow().sort_by("index").to_pylist() + self.assertEqual([0, 0, 1, 1], [r["episode_index"] for r in rows]) + self.assertEqual([0, 1, 0, 1], [r["frame_index"] for r in rows]) + self.assertEqual([0, 1, 2, 3], [r["index"] for r in rows]) + self.assertEqual([0, 0, 0, 1], [r["task_index"] for r in rows]) + + def test_existing_table_requires_matching_feature_schema(self): + PaimonLeRobotWriter( + self.connection, + "schema_mismatch", + fps=10, + features={ + "action": { + "dtype": "float32", + "shape": (1,), + "names": None, + }, + }, + ) + + with self.assertRaisesRegex(ValueError, "LeRobot feature action"): + PaimonLeRobotWriter( + self.connection, + "schema_mismatch", + fps=10, + features={ + "action": { + "dtype": "int64", + "shape": (1,), + "names": None, + }, + }, + ) + + def test_commits_multiple_completed_episodes_as_one_batch(self): + writer = PaimonLeRobotWriter( + self.connection, + "robot_data", + fps=10, + episodes_per_commit=2, + features={ + "observation.state": { + "dtype": "float32", + "shape": (2,), + "names": ["x", "y"], + }, + "action": { + "dtype": "float32", + "shape": (1,), + "names": ["gripper"], + }, + }, + ) + + writer.add_frame({ + "observation.state": np.array([1.0, 2.0], dtype=np.float32), + "action": np.array([0.5], dtype=np.float32), + "task": "pick", + }) + self.assertIsNone(writer.save_episode()) + table = self.connection.get_table("robot_data") + self.assertIsNone( + table.raw_table.snapshot_manager().get_latest_snapshot()) + + writer.add_frame({ + "observation.state": np.array([3.0, 4.0], dtype=np.float32), + "action": np.array([0.0], dtype=np.float32), + "task": "place", + }) + self.assertIsNone(writer.save_episode()) + self.assertEqual( + 1, table.raw_table.snapshot_manager().get_latest_snapshot().id) + + rows = table.scan().select([ + "episode_index", + "frame_index", + "timestamp", + "index", + "task_index", + "task", + "observation.state", + "action", + ]).to_arrow().sort_by("index").to_pylist() + self.assertEqual([0, 1], [row["episode_index"] for row in rows]) + self.assertEqual([0, 0], [row["frame_index"] for row in rows]) + self.assertEqual([0, 1], [row["index"] for row in rows]) + self.assertEqual([0, 1], [row["task_index"] for row in rows]) + self.assertEqual(["pick", "place"], [row["task"] for row in rows]) + self.assertEqual([0.0, 0.0], [row["timestamp"] for row in rows]) + self.assertEqual([0.5, 0.0], [row["action"] for row in rows]) + + def test_writes_raw_image_frame_as_png_blob(self): + writer = PaimonLeRobotWriter( + self.connection, + "images", + fps=30, + features={ + "observation.image": { + "dtype": "image", + "shape": (3, 4, 5), + "names": ["channels", "height", "width"], + }, + }, + ) + writer.add_frame({ + "observation.image": np.full( + (4, 5, 3), 73, dtype=np.uint8), + "task": "inspect", + }) + writer.save_episode() + self.assertIsNone(writer.flush()) + + table = self.connection.get_table("images") + scalar, blobs = table.scan().select([ + "index", "observation.image" + ]).read_blobs() + self.assertEqual([0], scalar.column("index").to_pylist()) + image = Image.open(io.BytesIO(blobs["observation.image"][0])) + self.assertEqual((5, 4), image.size) + self.assertEqual((73, 73, 73), image.getpixel((0, 0))) + + def test_discards_rerecorded_episode_and_finalizes_tail_batch(self): + writer = PaimonLeRobotWriter( + self.connection, + "rerecord", + fps=10, + features={ + "action": { + "dtype": "float32", + "shape": (1,), + "names": None, + }, + }, + ) + writer.add_frame({ + "action": np.array([9.0], dtype=np.float32), + "task": "discard", + }) + self.assertTrue(writer.has_pending_frames()) + writer.clear_episode_buffer() + self.assertFalse(writer.has_pending_frames()) + + writer.add_frame({ + "action": np.array([1.0], dtype=np.float32), + "task": "keep", + }) + writer.save_episode() + self.assertEqual(1, writer.num_episodes) + self.assertEqual(1, writer.pending_episodes) + writer.finalize() + writer.finalize() + + rows = self.connection.get_table("rerecord").scan().select([ + "episode_index", "index", "task", "action" + ]).to_arrow().to_pylist() + self.assertEqual([{ + "episode_index": 0, + "index": 0, + "task": "keep", + "action": 1.0, + }], rows) + with self.assertRaisesRegex(RuntimeError, "after finalize"): + writer.add_frame({ + "action": np.array([2.0], dtype=np.float32), + "task": "late", + }) + + def test_finalize_rejects_unsaved_frames(self): + writer = PaimonLeRobotWriter( + self.connection, + "unsaved", + fps=10, + features={ + "action": { + "dtype": "float32", + "shape": (1,), + "names": None, + }, + }, + ) + writer.add_frame({ + "action": np.array([1.0], dtype=np.float32), + "task": "pick", + }) + + with self.assertRaisesRegex(RuntimeError, "unsaved"): + writer.finalize() + self.assertTrue(writer.has_pending_frames()) + self.assertIsNone( + self.connection.get_table("unsaved").raw_table + .snapshot_manager().get_latest_snapshot()) + + def test_add_frame_rejects_invalid_value_without_buffering_it(self): + writer = PaimonLeRobotWriter( + self.connection, + "invalid_frame", + fps=10, + features={ + "sensor": { + "dtype": "uint8", + "shape": (1,), + "names": None, + }, + }, + ) + + with self.assertRaisesRegex(ValueError, "expected dtype"): + writer.add_frame({ + "sensor": np.array([1], dtype=np.int16), + "task": "measure", + }) + self.assertFalse(writer.has_pending_frames()) + + def test_add_frame_requires_native_numeric_dtype_and_shape(self): + writer = PaimonLeRobotWriter( + self.connection, + "native_frame_contract", + fps=10, + features={ + "action": { + "dtype": "float32", + "shape": (1,), + "names": None, + }, + }, + ) + + with self.assertRaisesRegex(ValueError, "NumPy array"): + writer.add_frame({"action": [1.0], "task": "pick"}) + with self.assertRaisesRegex(ValueError, "expected dtype"): + writer.add_frame({ + "action": np.array([1.0], dtype=np.float64), + "task": "pick", + }) + self.assertFalse(writer.has_pending_frames()) + + def test_add_frame_validates_pil_image_shape(self): + writer = PaimonLeRobotWriter( + self.connection, + "pil_shape", + fps=10, + features={ + "observation.image": { + "dtype": "image", + "shape": (3, 4, 5), + "names": ["channels", "height", "width"], + }, + }, + ) + + with self.assertRaisesRegex(ValueError, "expected shape"): + writer.add_frame({ + "observation.image": Image.new("RGB", (6, 4)), + "task": "inspect", + }) + self.assertFalse(writer.has_pending_frames()) + + +if __name__ == "__main__": + unittest.main() From 0c3ebdd1157181c04563276f0dc00f0fa1a83092 Mon Sep 17 00:00:00 2001 From: Yann Date: Tue, 8 Sep 2026 21:39:32 +0800 Subject: [PATCH 2/4] feat(python): persist LeRobot writer state in snapshots Allow table commits to attach snapshot properties and use them to resume LeRobot frame, episode, and task indices without scanning table data. Co-Authored-By: Codex AI-Model: gpt-5 Co-Authored-By: Codex Co-Authored-By: Codex AI-Contributed/Feature: 71/71 AI-Contributed/UT: 40/40 --- docs/docs/pypaimon/multimodal-api.mdx | 12 +++- .../pypaimon/multimodal/lerobot/writer.py | 59 ++++++++++++++++++- .../tests/multimodal_lerobot_writer_test.py | 40 ++++++++++--- 3 files changed, 99 insertions(+), 12 deletions(-) diff --git a/docs/docs/pypaimon/multimodal-api.mdx b/docs/docs/pypaimon/multimodal-api.mdx index c1647a165741..ef0a38e28c87 100644 --- a/docs/docs/pypaimon/multimodal-api.mdx +++ b/docs/docs/pypaimon/multimodal-api.mdx @@ -738,15 +738,21 @@ or call `flush()` at an operational boundary. `save_episode`, `flush`, and frames raises `ValueError`, matching native LeRobot. Call `clear_episode_buffer()` before `save_episode()` to discard a re-recorded -episode without advancing frame or episode indices. `finalize()` rejects an -unfinished episode instead of silently dropping its frames. +episode without advancing frame or episode indices. Once `save_episode()` +accepts an episode, `clear_episode_buffer()` no longer affects it, even when the +batch has not yet been committed. `finalize()` rejects an unfinished episode +instead of silently dropping its frames. The writer creates a missing table and appends to an existing compatible table. Before writing, it requires the table columns, order, Arrow types, nullability, and LeRobot feature metadata to match. On resume, new `index` and `episode_index` values continue after the existing maxima, existing task mappings are retained, and the episode-local `frame_index` starts again at zero. -A commit exception has an unknown result and is not automatically retried. +Each commit records the next global indices and task mapping in Snapshot +properties. Resume normally reads only that metadata; a non-empty table created +before these properties existed is scanned once and upgraded by its next +commit. A commit exception has an unknown result and is not automatically +retried. ## Overwrite diff --git a/paimon-python/pypaimon/multimodal/lerobot/writer.py b/paimon-python/pypaimon/multimodal/lerobot/writer.py index 77a320696b93..7190d32190d3 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/writer.py +++ b/paimon-python/pypaimon/multimodal/lerobot/writer.py @@ -18,6 +18,7 @@ """LeRobot-compatible capture writer for multimodal Paimon tables.""" import copy +import json from typing import Mapping, Optional import numpy as np @@ -47,6 +48,11 @@ "task_index": {"dtype": "int64", "shape": (1,), "names": None}, } _TASK_FEATURE = {"dtype": "string", "shape": (1,), "names": None} +_STATE_PREFIX = "pypaimon.lerobot." +_STATE_VERSION = _STATE_PREFIX + "state-version" +_NEXT_INDEX = _STATE_PREFIX + "next-index" +_NEXT_EPISODE_INDEX = _STATE_PREFIX + "next-episode-index" +_TASK_INDICES = _STATE_PREFIX + "task-indices" class PaimonLeRobotWriter: @@ -114,6 +120,16 @@ def __init__( self._failed = False def _load_existing_state(self): + snapshot = self._table.raw_table.snapshot_manager() \ + .get_latest_snapshot() + if snapshot is None: + return 0, 0, {} + properties = snapshot.properties or {} + if _STATE_VERSION in properties: + return self._state_from_snapshot_properties(properties) + if any(key.startswith(_STATE_PREFIX) for key in properties): + raise ValueError("Existing LeRobot snapshot state is incomplete.") + # ponytail: one resume scan; persist counters if startup cost matters. rows = self._table.scan().select([ "index", "episode_index", "task_index", "task" @@ -140,6 +156,45 @@ def _load_existing_state(self): task_indices, ) + @staticmethod + def _state_from_snapshot_properties(properties): + if properties[_STATE_VERSION] != "1": + raise ValueError( + "Unsupported LeRobot snapshot state version %r." + % properties[_STATE_VERSION]) + try: + next_index = int(properties[_NEXT_INDEX]) + next_episode_index = int(properties[_NEXT_EPISODE_INDEX]) + task_indices = json.loads(properties[_TASK_INDICES]) + except (KeyError, TypeError, ValueError) as error: + raise ValueError( + "Existing LeRobot snapshot state is invalid.") from error + if next_index < 0 or next_episode_index < 0 \ + or not isinstance(task_indices, dict): + raise ValueError("Existing LeRobot snapshot state is invalid.") + indices = list(task_indices.values()) + if any(not isinstance(task, str) + or isinstance(index, bool) + or not isinstance(index, int) + or index < 0 + for task, index in task_indices.items()) \ + or len(set(indices)) != len(indices): + raise ValueError("Existing LeRobot snapshot state is invalid.") + return next_index, next_episode_index, task_indices + + def _snapshot_properties(self): + return { + _STATE_VERSION: "1", + _NEXT_INDEX: str(self.num_frames), + _NEXT_EPISODE_INDEX: str(self.num_episodes), + _TASK_INDICES: json.dumps( + self._task_indices, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ), + } + def add_frame(self, frame): self._require_open("add_frame") if not isinstance(frame, Mapping): @@ -257,7 +312,9 @@ def flush(self): try: messages = self._table_write.prepare_commit() commit_started = True - self._table_commit.commit(messages) + self._table_commit.commit( + messages, + snapshot_properties=self._snapshot_properties()) snapshot_id = self._snapshot_recorder.snapshot_id if snapshot_id is None: raise RuntimeError( diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_writer_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_writer_test.py index 1bc117772f31..0e931e41f059 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_writer_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_writer_test.py @@ -20,6 +20,7 @@ import tempfile import unittest from pathlib import Path +from unittest.mock import patch import numpy as np from PIL import Image @@ -87,8 +88,15 @@ def test_existing_table_resumes_global_and_task_indices(self): first.save_episode() first.finalize() - resumed = PaimonLeRobotWriter( - self.connection, "resume", fps=10, features=features) + snapshot = self.connection.get_table( + "resume").raw_table.snapshot_manager().get_latest_snapshot() + self.assertEqual("1", snapshot.properties[ + "pypaimon.lerobot.state-version"]) + with patch( + "pypaimon.multimodal.table.MultimodalTable.scan", + side_effect=AssertionError("resume must not scan table data")): + resumed = PaimonLeRobotWriter( + self.connection, "resume", fps=10, features=features) self.assertEqual(2, resumed.num_frames) self.assertEqual(1, resumed.num_episodes) resumed.add_frame({ @@ -253,18 +261,34 @@ def test_discards_rerecorded_episode_and_finalizes_tail_batch(self): writer.save_episode() self.assertEqual(1, writer.num_episodes) self.assertEqual(1, writer.pending_episodes) + writer.clear_episode_buffer() + self.assertEqual(1, writer.pending_episodes) + + writer.add_frame({ + "action": np.array([2.0], dtype=np.float32), + "task": "also keep", + }) + writer.save_episode() writer.finalize() writer.finalize() rows = self.connection.get_table("rerecord").scan().select([ "episode_index", "index", "task", "action" ]).to_arrow().to_pylist() - self.assertEqual([{ - "episode_index": 0, - "index": 0, - "task": "keep", - "action": 1.0, - }], rows) + self.assertEqual([ + { + "episode_index": 0, + "index": 0, + "task": "keep", + "action": 1.0, + }, + { + "episode_index": 1, + "index": 1, + "task": "also keep", + "action": 2.0, + }, + ], rows) with self.assertRaisesRegex(RuntimeError, "after finalize"): writer.add_frame({ "action": np.array([2.0], dtype=np.float32), From f6d7e458ef3120324dad5db38cb6f30af1899836 Mon Sep 17 00:00:00 2001 From: Yann Date: Wed, 9 Sep 2026 11:16:58 +0800 Subject: [PATCH 3/4] fix(python): accept native LeRobot feature layouts Preserve HWC camera layouts from native feature metadata and normalize validated NumPy arrays before Arrow conversion. Co-Authored-By: Codex AI-Model: gpt-5 Co-Authored-By: Codex AI-Contributed/Feature: 22/22 AI-Contributed/UT: 63/63 --- .../pypaimon/multimodal/lerobot/writer.py | 22 ++++--- .../tests/multimodal_lerobot_writer_test.py | 63 +++++++++++++++++++ 2 files changed, 78 insertions(+), 7 deletions(-) diff --git a/paimon-python/pypaimon/multimodal/lerobot/writer.py b/paimon-python/pypaimon/multimodal/lerobot/writer.py index 7190d32190d3..faa3c6b5bdb4 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/writer.py +++ b/paimon-python/pypaimon/multimodal/lerobot/writer.py @@ -250,7 +250,8 @@ def _normalize_frame_value(value, feature, name): raise ValueError( "LeRobot feature %s expected shape %s, got %s." % (name, expected_shape, value.shape)) - return _normalize_value(value, feature, name) + value = _normalize_value(value, feature, name) + return value.tolist() if isinstance(value, np.ndarray) else value @staticmethod def _image_bytes(value, feature, name): @@ -261,18 +262,25 @@ def _image_bytes(value, feature, name): and isinstance(image_size, tuple) and len(image_size) == 2: actual_shape = (image_size[1], image_size[0], len(getbands())) expected_shape = _feature_shape(feature, name) - channel_last_shape = ( - (expected_shape[1], expected_shape[2], expected_shape[0]) - if len(expected_shape) == 3 else () - ) - if actual_shape and actual_shape != expected_shape \ + names = tuple(feature.get("names") or ()) + if names == ("height", "width", "channels"): + channel_first_shape = ( + expected_shape[2], expected_shape[0], expected_shape[1]) + channel_last_shape = expected_shape + else: + channel_first_shape = expected_shape + channel_last_shape = ( + (expected_shape[1], expected_shape[2], expected_shape[0]) + if len(expected_shape) == 3 else () + ) + if actual_shape and actual_shape != channel_first_shape \ and actual_shape != channel_last_shape: raise ValueError( "LeRobot feature %s expected shape %s, got %s." % (name, expected_shape, actual_shape)) return _encode_media_frame( value, - channel_first=actual_shape == expected_shape, + channel_first=actual_shape == channel_first_shape, ) def save_episode(self): diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_writer_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_writer_test.py index 0e931e41f059..f5654dd941fc 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_writer_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_writer_test.py @@ -233,6 +233,69 @@ def test_writes_raw_image_frame_as_png_blob(self): self.assertEqual((5, 4), image.size) self.assertEqual((73, 73, 73), image.getpixel((0, 0))) + def test_writes_native_hwc_image_frame_as_png_blob(self): + writer = PaimonLeRobotWriter( + self.connection, + "native_hwc_images", + fps=30, + features={ + "observation.image": { + "dtype": "image", + "shape": (4, 5, 3), + "names": ["height", "width", "channels"], + }, + }, + ) + writer.add_frame({ + "observation.image": np.full( + (4, 5, 3), 73, dtype=np.uint8), + "task": "inspect", + }) + writer.save_episode() + writer.finalize() + + _, blobs = self.connection.get_table( + "native_hwc_images").scan().select([ + "observation.image" + ]).read_blobs() + image = Image.open(io.BytesIO(blobs["observation.image"][0])) + self.assertEqual((5, 4), image.size) + self.assertEqual((73, 73, 73), image.getpixel((0, 0))) + + def test_writes_multidimensional_and_boolean_numpy_features(self): + writer = PaimonLeRobotWriter( + self.connection, + "numpy_features", + fps=10, + features={ + "observation.matrix": { + "dtype": "float32", + "shape": (2, 2), + "names": None, + }, + "observation.flags": { + "dtype": "bool", + "shape": (2,), + "names": None, + }, + }, + ) + writer.add_frame({ + "observation.matrix": np.array( + [[1.0, 2.0], [3.0, 4.0]], dtype=np.float32), + "observation.flags": np.array([True, False], dtype=np.bool_), + "task": "inspect", + }) + writer.save_episode() + writer.finalize() + + rows = self.connection.get_table("numpy_features").scan().select([ + "observation.matrix", "observation.flags" + ]).to_arrow().to_pylist() + self.assertEqual([[1.0, 2.0], [3.0, 4.0]], + rows[0]["observation.matrix"]) + self.assertEqual([True, False], rows[0]["observation.flags"]) + def test_discards_rerecorded_episode_and_finalizes_tail_batch(self): writer = PaimonLeRobotWriter( self.connection, From 5306e253d25f3ef7a1f1357126dc53a0cf961857 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Wed, 9 Sep 2026 13:22:20 +0800 Subject: [PATCH 4/4] [python] Guard optional Pillow dependency in LeRobot writer tests --- .../pypaimon/tests/multimodal_lerobot_writer_test.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_writer_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_writer_test.py index f5654dd941fc..f446de75818f 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_writer_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_writer_test.py @@ -23,11 +23,15 @@ from unittest.mock import patch import numpy as np -from PIL import Image import pypaimon.multimodal as pmm from pypaimon.multimodal.lerobot import PaimonLeRobotWriter +try: + from PIL import Image +except ImportError: + Image = None + class PaimonLeRobotWriterTest(unittest.TestCase): @@ -203,6 +207,7 @@ def test_commits_multiple_completed_episodes_as_one_batch(self): self.assertEqual([0.0, 0.0], [row["timestamp"] for row in rows]) self.assertEqual([0.5, 0.0], [row["action"] for row in rows]) + @unittest.skipUnless(Image is not None, "Pillow is required for image tests") def test_writes_raw_image_frame_as_png_blob(self): writer = PaimonLeRobotWriter( self.connection, @@ -233,6 +238,7 @@ def test_writes_raw_image_frame_as_png_blob(self): self.assertEqual((5, 4), image.size) self.assertEqual((73, 73, 73), image.getpixel((0, 0))) + @unittest.skipUnless(Image is not None, "Pillow is required for image tests") def test_writes_native_hwc_image_frame_as_png_blob(self): writer = PaimonLeRobotWriter( self.connection, @@ -427,6 +433,7 @@ def test_add_frame_requires_native_numeric_dtype_and_shape(self): }) self.assertFalse(writer.has_pending_frames()) + @unittest.skipUnless(Image is not None, "Pillow is required for image tests") def test_add_frame_validates_pil_image_shape(self): writer = PaimonLeRobotWriter( self.connection,