This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new d95ebff28f [python] Add direct LeRobot capture writer (#9665)
d95ebff28f is described below
commit d95ebff28fafd0443b6a2b0801f61a1ec027193e
Author: Yann Byron <[email protected]>
AuthorDate: Wed Sep 9 14:04:57 2026 +0800
[python] Add direct LeRobot capture writer (#9665)
---
docs/docs/pypaimon/multimodal-api.mdx | 60 +++
paimon-python/pypaimon/multimodal/arrow_utils.py | 5 +
.../pypaimon/multimodal/lerobot/__init__.py | 4 +-
.../pypaimon/multimodal/lerobot/loader.py | 7 +-
.../pypaimon/multimodal/lerobot/writer.py | 419 +++++++++++++++++++
.../tests/multimodal_lerobot_writer_test.py | 460 +++++++++++++++++++++
6 files changed, 952 insertions(+), 3 deletions(-)
diff --git a/docs/docs/pypaimon/multimodal-api.mdx
b/docs/docs/pypaimon/multimodal-api.mdx
index ddef5d1475..ef0a38e28c 100644
--- a/docs/docs/pypaimon/multimodal-api.mdx
+++ b/docs/docs/pypaimon/multimodal-api.mdx
@@ -694,6 +694,66 @@ 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. 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.
+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
`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 76d056ed92..47f9245222 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 a40f2a8cce..5f598d2991 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 295eee99c3..71719869c3 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 0000000000..faa3c6b5bd
--- /dev/null
+++ b/paimon-python/pypaimon/multimodal/lerobot/writer.py
@@ -0,0 +1,419 @@
+# 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
+import json
+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}
+_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:
+ """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):
+ 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"
+ ]).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,
+ )
+
+ @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):
+ 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))
+ value = _normalize_value(value, feature, name)
+ return value.tolist() if isinstance(value, np.ndarray) else value
+
+ @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)
+ 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 == channel_first_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_properties=self._snapshot_properties())
+ 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 0000000000..f446de7581
--- /dev/null
+++ b/paimon-python/pypaimon/tests/multimodal_lerobot_writer_test.py
@@ -0,0 +1,460 @@
+# 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
+from unittest.mock import patch
+
+import numpy as np
+
+import pypaimon.multimodal as pmm
+from pypaimon.multimodal.lerobot import PaimonLeRobotWriter
+
+try:
+ from PIL import Image
+except ImportError:
+ Image = None
+
+
+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()
+
+ 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({
+ "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])
+
+ @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,
+ "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)))
+
+ @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,
+ "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,
+ "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.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,
+ },
+ {
+ "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),
+ "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())
+
+ @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,
+ "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()
