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 3988f2ac3a [python] Preserve LeRobot recording metadata in direct
writer (#9708)
3988f2ac3a is described below
commit 3988f2ac3aaf09e2dbeb3e1f15145ab5cace181c
Author: Yann Byron <[email protected]>
AuthorDate: Thu Sep 10 17:15:19 2026 +0800
[python] Preserve LeRobot recording metadata in direct writer (#9708)
---
docs/docs/pypaimon/multimodal-api.mdx | 61 ++-
.../pypaimon/multimodal/lerobot/metadata.py | 30 ++
.../pypaimon/multimodal/lerobot/writer.py | 543 ++++++++++++++++++---
.../tests/multimodal_lerobot_writer_test.py | 513 ++++++++++++++++++-
4 files changed, 1068 insertions(+), 79 deletions(-)
diff --git a/docs/docs/pypaimon/multimodal-api.mdx
b/docs/docs/pypaimon/multimodal-api.mdx
index 6edbeeb798..a39ecaac27 100644
--- a/docs/docs/pypaimon/multimodal-api.mdx
+++ b/docs/docs/pypaimon/multimodal-api.mdx
@@ -702,6 +702,14 @@ 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.
+The writer preserves native LeRobot recording metadata in a managed Paimon
+table group. The named table stores frames, while `<table>__episodes`,
+`<table>__tasks`, `<table>__info`, and `<table>__stats` store episode
+boundaries, task labels, JSON-encoded dataset information, and global
+statistics. When `subtask_index` is declared, `<table>__subtasks` stores its
+ordered text vocabulary. The root table's managed options identify these
+companions.
+
```python
from pypaimon.multimodal.lerobot import PaimonLeRobotWriter
@@ -710,6 +718,8 @@ writer = PaimonLeRobotWriter(
"recorded_frames",
fps=30,
features=dataset_features,
+ # Required for a new table when features includes subtask_index:
+ # subtasks=["approach object", "grasp object"],
)
# LeRobot's record_loop only needs writer.fps, writer.features, and
@@ -725,10 +735,23 @@ 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.
+shape. Image dimension names may declare CHW or native HWC layout; 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. Task text is stored once in the tasks table; frame rows
+retain only `task_index`. Optional subtask text is likewise stored once in the
+subtasks table, while each frame supplies its declared NumPy `subtask_index`.
+Subtask labels must be non-empty and unique, and frame indices must reference
+that ordered vocabulary. On resume, the writer restores the vocabulary from
+the existing table; an explicitly supplied vocabulary must match it exactly.
+
+When `save_episode()` accepts an episode, the writer uses LeRobot's native
+statistics implementation to calculate `min`, `max`, `mean`, `std`, `count`,
+`q01`, `q10`, `q50`, `q90`, and `q99` for every non-string feature. Flattened
+episode statistics are appended to the episodes table. Image statistics use
+the encoded PNG frames and LeRobot's sampling, downsampling, CHW, and `[0,1]`
+normalization rules. Each flush aggregates the accepted episode statistics and
+replaces the global stats table; it does not rescan frame data.
`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
@@ -743,16 +766,26 @@ 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.
+The writer creates a missing table group and appends to an existing compatible
+group. Before writing, it requires the root columns, order, Arrow types,
+nullability, LeRobot feature metadata, managed options, companion schemas, and
+component counts to match. A legacy frame-only writer table is not implicitly
+migrated. An existing group must have the stats companion and must not have a
+subtasks companion unless its frame schema declares `subtask_index`. On resume,
+new `index` and `episode_index` values continue after the published metadata,
+existing task and subtask mappings come from their companion tables, and the
+episode-local `frame_index` starts again at zero. Snapshot properties retain
+only the next global frame and episode indices; resume restores global
+statistics from the stats table and does not scan frame data.
+
+When configured, the first `flush()` writes the immutable subtask vocabulary.
+Each flush appends new tasks and episodes, replaces global stats and info, and
+commits frames last. Paimon
+does not provide a transaction across these tables. A component commit
+exception leaves the group result unknown, makes the writer terminal, and is
+not automatically retried. Reopening validates the component state and rejects
+a partial batch. Pause writes and use `create_lerobot_tag` before training to
+pin one named snapshot on every component.
## Train with Paimon LeRobot data
diff --git a/paimon-python/pypaimon/multimodal/lerobot/metadata.py
b/paimon-python/pypaimon/multimodal/lerobot/metadata.py
index b1130bbbb1..e9daa02066 100644
--- a/paimon-python/pypaimon/multimodal/lerobot/metadata.py
+++ b/paimon-python/pypaimon/multimodal/lerobot/metadata.py
@@ -51,6 +51,10 @@ _EMPTY_TASKS_SCHEMA = pa.schema([
pa.field("task_index", pa.int64(), nullable=False),
pa.field("task", pa.string(), nullable=False),
])
+_EMPTY_SUBTASKS_SCHEMA = pa.schema([
+ pa.field("subtask_index", pa.int64(), nullable=False),
+ pa.field("subtask", pa.string(), nullable=False),
+])
_EMPTY_EPISODES_SCHEMA = pa.schema([
pa.field("episode_index", pa.int64(), nullable=False),
pa.field("dataset_from_index", pa.int64(), nullable=False),
@@ -328,6 +332,32 @@ def _append_arrow(table, data):
return _append_arrow_tables(table, [data])
+def _overwrite_arrow(table, data):
+ target_schema = _target_schema(table)
+ if not data.schema.equals(target_schema, check_metadata=False):
+ raise ValueError(
+ "LeRobot component schema %s does not match target %s."
+ % (data.schema, target_schema))
+ builder = table.new_batch_write_builder().overwrite()
+ table_write = builder.new_write()
+ table_commit = builder.new_commit()
+ commit_started = False
+ try:
+ table_write.write_arrow(data)
+ messages = table_write.prepare_commit()
+ commit_started = True
+ table_commit.commit(messages)
+ except BaseException:
+ if not commit_started:
+ table_write.abort()
+ raise
+ finally:
+ try:
+ table_write.close()
+ finally:
+ table_commit.close()
+
+
def _append_arrow_tables(table, tables):
builder = table.new_batch_write_builder()
table_write = None
diff --git a/paimon-python/pypaimon/multimodal/lerobot/writer.py
b/paimon-python/pypaimon/multimodal/lerobot/writer.py
index faa3c6b5bd..f73fadf110 100644
--- a/paimon-python/pypaimon/multimodal/lerobot/writer.py
+++ b/paimon-python/pypaimon/multimodal/lerobot/writer.py
@@ -18,14 +18,29 @@
"""LeRobot-compatible capture writer for multimodal Paimon tables."""
import copy
+import io
import json
-from typing import Mapping, Optional
+from typing import Mapping, Optional, Sequence
import numpy as np
import pyarrow as pa
+from pypaimon.catalog.catalog_exception import (
+ DatabaseNotExistException,
+ TableNotExistException,
+)
from pypaimon.multimodal.arrow_utils import strict_arrow_table
from pypaimon.multimodal.hdf5 import _SnapshotRecorder
+from pypaimon.multimodal.lerobot.metadata import (
+ _COMPANION_OPTION_KEYS,
+ _EMPTY_EPISODES_SCHEMA,
+ _append_arrow,
+ _companion_table_identifiers,
+ _managed_table_options,
+ _metadata_table,
+ _overwrite_arrow,
+ _prepare_metadata_tables,
+)
from pypaimon.multimodal.lerobot.loader import (
_encode_media_frame,
_normalize_value,
@@ -36,6 +51,7 @@ from pypaimon.multimodal.lerobot.schema import (
_feature_shape,
_schema_from_info,
_validate_lerobot_schema,
+ _validate_v3_required_features,
)
from pypaimon.multimodal.table import _target_schema
@@ -47,12 +63,204 @@ _DEFAULT_FEATURES = {
"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"
+_STATE_VERSION_VALUE = "1"
_NEXT_INDEX = _STATE_PREFIX + "next-index"
_NEXT_EPISODE_INDEX = _STATE_PREFIX + "next-episode-index"
-_TASK_INDICES = _STATE_PREFIX + "task-indices"
+_STAT_NAMES = (
+ "min", "max", "mean", "std", "count",
+ "q01", "q10", "q50", "q90", "q99",
+)
+_INTEGER_DTYPES = {
+ "int8", "int16", "int32", "int64", "uint8", "uint16", "uint32",
+}
+_EPISODE_STATE_COLUMNS = list(_EMPTY_EPISODES_SCHEMA.names) + [
+ "stats/index/count",
+]
+
+
+def _read_arrow(table, projection=None):
+ builder = table.new_read_builder()
+ if projection is not None:
+ builder = builder.with_projection(projection)
+ return builder.new_read().to_arrow(builder.new_scan().plan().splits())
+
+
+def _lerobot_stats_functions():
+ try:
+ from lerobot.datasets.compute_stats import (
+ aggregate_stats,
+ auto_downsample_height_width,
+ compute_episode_stats,
+ get_feature_stats,
+ sample_indices,
+ )
+ except ImportError as error:
+ raise ImportError(
+ "PaimonLeRobotWriter statistics require LeRobot; install "
+ "'pypaimon[lerobot]'.") from error
+ return (aggregate_stats, auto_downsample_height_width,
+ compute_episode_stats, get_feature_stats, sample_indices)
+
+
+def _nested_list_type(value_type, depth):
+ for _ in range(depth):
+ value_type = pa.list_(value_type)
+ return value_type
+
+
+def _episode_schema(features):
+ fields = list(_EMPTY_EPISODES_SCHEMA)
+ for name, feature in features.items():
+ dtype = str(feature.get("dtype", ""))
+ if dtype == "string":
+ continue
+ feature_shape = _feature_shape(feature, name)
+ for stat in _STAT_NAMES:
+ if stat == "count":
+ value_type = pa.int64()
+ depth = 1
+ elif dtype == "image":
+ value_type = pa.float64()
+ depth = 3
+ elif stat in ("min", "max") and dtype in _INTEGER_DTYPES:
+ value_type = pa.int64()
+ elif stat in ("min", "max") and dtype in ("bool", "boolean"):
+ value_type = pa.bool_()
+ else:
+ value_type = pa.float64()
+ if stat != "count" and dtype != "image":
+ depth = max(1, len(feature_shape))
+ fields.append(pa.field(
+ "stats/%s/%s" % (name, stat),
+ _nested_list_type(value_type, depth),
+ nullable=False,
+ ))
+ return pa.schema(fields)
+
+
+def _metadata_values(table, component):
+ result = {}
+ for row in _read_arrow(table).to_pylist():
+ if row["key"] in result:
+ raise ValueError(
+ "Existing LeRobot %s metadata repeats key %r."
+ % (component, row["key"]))
+ try:
+ result[row["key"]] = json.loads(row["value"])
+ except (TypeError, ValueError) as error:
+ raise ValueError(
+ "Existing LeRobot %s metadata is invalid."
+ % component) from error
+ return result
+
+
+def _image_stats(values):
+ try:
+ from PIL import Image
+ except ImportError as error:
+ raise ImportError(
+ "PaimonLeRobotWriter image statistics require Pillow from "
+ "'pypaimon[lerobot]'.") from error
+ (_, downsample, _, get_feature_stats,
+ sample_indices) = _lerobot_stats_functions()
+ images = []
+ for index in sample_indices(len(values)):
+ with Image.open(io.BytesIO(values[index])) as image:
+ array = np.asarray(image.convert("RGB"), dtype=np.uint8)
+ images.append(downsample(np.transpose(array, (2, 0, 1))))
+ stats = get_feature_stats(
+ np.stack(images), axis=(0, 2, 3), keepdims=True)
+ return {
+ name: value if name == "count" else np.squeeze(
+ value / 255.0, axis=0)
+ for name, value in stats.items()
+ }
+
+
+def _compute_stats(episode, features):
+ _, _, compute_episode_stats, _, _ = \
+ _lerobot_stats_functions()
+ data = {}
+ numeric_features = {}
+ reshaped_features = {}
+ result = {}
+ for name, feature in features.items():
+ dtype = str(feature.get("dtype", ""))
+ values = episode.column(name).to_pylist()
+ if dtype == "image":
+ result[name] = _image_stats(values)
+ else:
+ numeric_features[name] = feature
+ array = (
+ values if dtype == "string" else np.asarray(
+ values,
+ dtype=np.dtype("bool" if dtype == "boolean" else dtype),
+ )
+ )
+ if dtype != "string" and array.ndim > 2:
+ # Keep higher-rank stats stable across one- and multi-frame
+ # episodes while delegating the calculation to LeRobot.
+ reshaped_features[name] = array.shape[1:]
+ array = array.reshape(array.shape[0], -1)
+ data[name] = array
+ numeric_stats = compute_episode_stats(data, numeric_features)
+ for name, shape in reshaped_features.items():
+ numeric_stats[name] = {
+ stat: value if stat == "count" else value.reshape(shape)
+ for stat, value in numeric_stats[name].items()
+ }
+ result.update(numeric_stats)
+ return result
+
+
+def _aggregate_stats(stats_list, features):
+ aggregate_stats = _lerobot_stats_functions()[0]
+ result = {}
+ for name, feature in features.items():
+ if feature.get("dtype") == "string":
+ continue
+ key = "image" if feature.get("dtype") == "image" else "feature"
+ result[name] = aggregate_stats([
+ {key: stats[name]} for stats in stats_list
+ ])[key]
+ return result
+
+
+def _indexed_metadata_table(component, entries):
+ import pandas as pd
+
+ entries = list(entries)
+ indices, labels = zip(*entries) if entries else ((), ())
+ return pa.Table.from_pandas(pd.DataFrame(
+ {component + "_index": np.asarray(indices, dtype=np.int64)},
+ index=pd.Index(
+ labels, dtype="string", name=component,
+ ),
+ ))
+
+
+def _subtasks_table(subtasks):
+ return _indexed_metadata_table("subtask", enumerate(subtasks))
+
+
+def _validate_subtasks(subtasks, has_feature):
+ if subtasks is None:
+ return None
+ if not has_feature:
+ raise ValueError(
+ "subtasks require a subtask_index feature.")
+ if isinstance(subtasks, (str, bytes)) \
+ or not isinstance(subtasks, Sequence):
+ raise ValueError("subtasks must be a sequence of strings.")
+ result = tuple(subtasks)
+ if not result or any(not isinstance(value, str) or not value
+ for value in result):
+ raise ValueError("subtasks must contain non-empty strings.")
+ if len(set(result)) != len(result):
+ raise ValueError("subtasks must not contain duplicates.")
+ return result
class PaimonLeRobotWriter:
@@ -65,6 +273,7 @@ class PaimonLeRobotWriter:
*,
fps: int,
features: Mapping[str, Mapping[str, object]],
+ subtasks: Optional[Sequence[str]] = None,
episodes_per_commit: int = -1,
options: Optional[Mapping[str, object]] = None):
if isinstance(fps, bool) or not isinstance(fps, int) or fps <= 0:
@@ -79,21 +288,44 @@ class PaimonLeRobotWriter:
raise ValueError("features must be a non-empty mapping.")
if "task" in features:
raise ValueError("task is managed by PaimonLeRobotWriter.")
+ requested_subtasks = _validate_subtasks(
+ subtasks, "subtask_index" in features)
+ _lerobot_stats_functions()
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,
- )
+ _validate_v3_required_features({"features": self.features})
+ self._source_schema = _schema_from_info({"features": self.features})
+ metadata = self._writer_metadata(
+ fps, self.features, requested_subtasks)
+ self._episodes_schema = metadata["episodes_schema"]
+ create_options = dict(options or {})
+ reserved_options = set(_COMPANION_OPTION_KEYS.values()).intersection(
+ create_options)
+ if reserved_options:
+ raise ValueError(
+ "%s are managed by PaimonLeRobotWriter."
+ % sorted(reserved_options))
+ create_options.update(_managed_table_options(
+ connection._identifier(table_name), metadata))
+ try:
+ self._table = connection.get_table(table_name)
+ created = False
+ except (DatabaseNotExistException, TableNotExistException):
+ if "subtask_index" in self.features \
+ and requested_subtasks is None:
+ raise ValueError(
+ "subtasks are required when creating a table with "
+ "subtask_index.")
+ self._table = connection.create_table(
+ table_name,
+ schema=self._source_schema,
+ options=create_options,
+ )
+ created = True
self._target_schema = _target_schema(self._table.raw_table)
_validate_lerobot_schema(
self._source_schema, self._target_schema, table_name)
@@ -104,95 +336,221 @@ class PaimonLeRobotWriter:
0,
"LeRobot",
)
+ self._metadata_tables = (
+ _prepare_metadata_tables(
+ connection, self._table.raw_table, metadata)
+ if created else self._open_metadata_tables(
+ connection, self._table.raw_table, metadata)
+ )
- self.num_frames, self.num_episodes, self._task_indices = \
- self._load_existing_state()
+ (self.num_frames, self.num_episodes, self._task_indices,
+ self._stats, stored_subtasks) = self._load_existing_state()
+ if stored_subtasks is not None:
+ if requested_subtasks is not None \
+ and requested_subtasks != stored_subtasks:
+ raise ValueError(
+ "subtasks do not match the existing LeRobot table.")
+ self.subtasks = stored_subtasks
+ else:
+ self.subtasks = requested_subtasks
+ if "subtask_index" in self.features and self.subtasks is None:
+ raise ValueError(
+ "subtasks are required when creating a table with "
+ "subtask_index.")
self._next_task_index = (
max(self._task_indices.values()) + 1
if self._task_indices else 0
)
self.pending_episodes = 0
self._episode_frames = []
+ self._pending_episode_rows = []
+ self._committed_task_count = len(self._task_indices)
+ self._subtasks_committed = self.num_frames > 0
self._table_write = None
self._table_commit = None
self._snapshot_recorder = None
self._finalized = False
self._failed = False
+ @staticmethod
+ def _writer_metadata(fps, features, subtasks):
+ info = {
+ "codebase_version": "v3.0",
+ "fps": fps,
+ "features": features,
+ "total_frames": 0,
+ "total_episodes": 0,
+ "total_tasks": 0,
+ "splits": {},
+ }
+ return {
+ "info_table": _metadata_table(info),
+ "episodes_schema": _episode_schema(features),
+ "tasks_table": _indexed_metadata_table("task", ()),
+ "stats_table": _metadata_table({}),
+ "subtasks_table": (
+ _subtasks_table(subtasks or ())
+ if "subtask_index" in features else None
+ ),
+ }
+
+ @staticmethod
+ def _open_metadata_tables(connection, frames_table, metadata):
+ identifiers = _companion_table_identifiers(frames_table)
+ expected = {
+ "info": metadata["info_table"].schema,
+ "episodes": metadata["episodes_schema"],
+ "tasks": metadata["tasks_table"].schema,
+ "stats": metadata["stats_table"].schema,
+ }
+ if metadata["subtasks_table"] is not None:
+ expected["subtasks"] = metadata["subtasks_table"].schema
+ if set(identifiers) != set(expected):
+ raise ValueError(
+ "PaimonLeRobotWriter companion tables do not match "
+ "the declared features.")
+ tables = {
+ name: connection.catalog.get_table(identifier)
+ for name, identifier in identifiers.items()
+ }
+ for name, table in tables.items():
+ if not _target_schema(table).equals(
+ expected[name], check_metadata=False):
+ raise ValueError(
+ "LeRobot %s companion schema does not match "
+ "PaimonLeRobotWriter." % name)
+ return tables
+
def _load_existing_state(self):
snapshot = self._table.raw_table.snapshot_manager() \
.get_latest_snapshot()
if snapshot is None:
- return 0, 0, {}
+ if any(table.snapshot_manager().get_latest_snapshot() is not None
+ for table in self._metadata_tables.values()):
+ raise ValueError(
+ "Existing LeRobot table group state is inconsistent.")
+ return 0, 0, {}, None, None
properties = snapshot.properties or {}
if _STATE_VERSION in properties:
- return self._state_from_snapshot_properties(properties)
+ state = self._state_from_snapshot_properties(properties)
+ metadata_state = self._state_from_companion_tables()
+ if state != metadata_state[:2]:
+ raise ValueError(
+ "Existing LeRobot table group state is inconsistent.")
+ return metadata_state
if any(key.startswith(_STATE_PREFIX) for key in properties):
raise ValueError("Existing LeRobot snapshot state is incomplete.")
+ return self._state_from_companion_tables()
- # 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, {}
-
+ def _state_from_companion_tables(self):
+ task_rows = _read_arrow(self._metadata_tables["tasks"]).to_pylist()
+ task_rows.sort(key=lambda row: row["task_index"])
task_indices = {}
- index_tasks = {}
- for row in rows:
+ for expected, row in enumerate(task_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)):
+ if row["task_index"] != expected or not isinstance(task, str) \
+ or task in task_indices:
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,
- )
+ "Existing LeRobot task metadata is invalid.")
+ task_indices[task] = expected
+
+ subtasks = None
+ if "subtasks" in self._metadata_tables:
+ subtask_rows = _read_arrow(
+ self._metadata_tables["subtasks"]).to_pylist()
+ subtask_rows.sort(key=lambda row: row["subtask_index"])
+ labels = []
+ for expected, row in enumerate(subtask_rows):
+ label = row["subtask"]
+ if row["subtask_index"] != expected \
+ or not isinstance(label, str) or not label \
+ or label in labels:
+ raise ValueError(
+ "Existing LeRobot subtask metadata is invalid.")
+ labels.append(label)
+ if not labels:
+ raise ValueError(
+ "Existing LeRobot subtask metadata is empty.")
+ subtasks = tuple(labels)
+
+ episode_rows = _read_arrow(
+ self._metadata_tables["episodes"],
+ _EPISODE_STATE_COLUMNS,
+ ).to_pylist()
+ episode_rows.sort(key=lambda row: row["episode_index"])
+ next_index = 0
+ for expected, row in enumerate(episode_rows):
+ if row["episode_index"] != expected \
+ or row["dataset_from_index"] != next_index \
+ or row["dataset_to_index"] <= next_index \
+ or row["length"] != (
+ row["dataset_to_index"] - next_index) \
+ or row["stats/index/count"] != [row["length"]] \
+ or any(task not in task_indices
+ for task in row["tasks"]):
+ raise ValueError(
+ "Existing LeRobot episode metadata is invalid.")
+ next_index = row["dataset_to_index"]
+
+ info = _metadata_values(self._metadata_tables["info"], "info")
+ if int(info.get("fps", -1)) != self.fps \
+ or info.get("features") != json.loads(json.dumps(
+ self.features, ensure_ascii=False)) \
+ or info.get("total_frames") != next_index \
+ or info.get("total_episodes") != len(episode_rows) \
+ or info.get("total_tasks") != len(task_indices):
+ raise ValueError(
+ "Existing LeRobot info metadata is inconsistent.")
+ stats = _metadata_values(self._metadata_tables["stats"], "stats")
+ expected_stats = {
+ name for name, feature in self.features.items()
+ if feature.get("dtype") != "string"
+ }
+ if set(stats) != expected_stats or any(
+ set(feature_stats) != set(_STAT_NAMES)
+ for feature_stats in stats.values()):
+ raise ValueError(
+ "Existing LeRobot stats metadata is inconsistent.")
+ numpy_stats = {
+ name: {
+ stat: np.asarray(value)
+ for stat, value in feature_stats.items()
+ }
+ for name, feature_stats in stats.items()
+ }
+ try:
+ _aggregate_stats([numpy_stats], self.features)
+ except (TypeError, ValueError) as error:
+ raise ValueError(
+ "Existing LeRobot stats metadata is invalid.") from error
+ if set(numpy_stats["index"]) != set(_STAT_NAMES) \
+ or numpy_stats["index"]["count"].tolist() != [next_index]:
+ raise ValueError(
+ "Existing LeRobot stats metadata is inconsistent.")
+ return (next_index, len(episode_rows), task_indices, numpy_stats,
+ subtasks)
@staticmethod
def _state_from_snapshot_properties(properties):
- if properties[_STATE_VERSION] != "1":
+ if properties[_STATE_VERSION] != _STATE_VERSION_VALUE:
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):
+ if next_index < 0 or next_episode_index < 0:
raise ValueError("Existing LeRobot snapshot state is invalid.")
- return next_index, next_episode_index, task_indices
+ return next_index, next_episode_index
def _snapshot_properties(self):
return {
- _STATE_VERSION: "1",
+ _STATE_VERSION: _STATE_VERSION_VALUE,
_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):
@@ -218,6 +576,11 @@ class PaimonLeRobotWriter:
else:
value = self._normalize_frame_value(
frame[name], feature, name)
+ if name == "subtask_index" \
+ and (value < 0 or value >= len(self.subtasks)):
+ raise ValueError(
+ "LeRobot frame subtask_index %d outside [0, %d)."
+ % (value, len(self.subtasks)))
_safe_array(
[value],
self._source_schema.field(name),
@@ -289,6 +652,11 @@ class PaimonLeRobotWriter:
raise ValueError("Cannot save an empty LeRobot episode.")
episode = self._episode_table()
+ episode_stats = _compute_stats(episode, self.features)
+ stats = (
+ _aggregate_stats([self._stats, episode_stats], self.features)
+ if self._stats is not None else episode_stats
+ )
try:
self._ensure_batch()
self._table_write.write_arrow(episode)
@@ -299,6 +667,20 @@ class PaimonLeRobotWriter:
self.num_frames += episode.num_rows
self.num_episodes += 1
self.pending_episodes += 1
+ self._pending_episode_rows.append({
+ "episode_index": self.num_episodes - 1,
+ "dataset_from_index": self.num_frames - episode.num_rows,
+ "dataset_to_index": self.num_frames,
+ "tasks": list(dict.fromkeys(
+ frame["task"] for frame in self._episode_frames)),
+ "length": episode.num_rows,
+ **{
+ "stats/%s/%s" % (feature, stat): value.tolist()
+ for feature, feature_stats in episode_stats.items()
+ for stat, value in feature_stats.items()
+ },
+ })
+ self._stats = stats
self._episode_frames = []
if self.episodes_per_commit != -1 \
and self.pending_episodes >= self.episodes_per_commit:
@@ -319,6 +701,47 @@ class PaimonLeRobotWriter:
commit_started = False
try:
messages = self._table_write.prepare_commit()
+ if "subtasks" in self._metadata_tables \
+ and not self._subtasks_committed:
+ _append_arrow(
+ self._metadata_tables["subtasks"],
+ _subtasks_table(self.subtasks),
+ )
+ task_rows = [
+ (index, task)
+ for task, index in sorted(
+ self._task_indices.items(), key=lambda item: item[1])
+ if index >= self._committed_task_count
+ ]
+ _append_arrow(
+ self._metadata_tables["tasks"],
+ _indexed_metadata_table("task", task_rows),
+ )
+ _append_arrow(
+ self._metadata_tables["episodes"],
+ pa.Table.from_pylist(
+ self._pending_episode_rows,
+ schema=self._episodes_schema,
+ ),
+ )
+ _overwrite_arrow(
+ self._metadata_tables["stats"],
+ _metadata_table(self._stats),
+ )
+ _overwrite_arrow(
+ self._metadata_tables["info"],
+ _metadata_table({
+ "codebase_version": "v3.0",
+ "fps": self.fps,
+ "features": self.features,
+ "total_frames": self.num_frames,
+ "total_episodes": self.num_episodes,
+ "total_tasks": len(self._task_indices),
+ "splits": {
+ "train": "0:%d" % self.num_episodes,
+ },
+ }),
+ )
commit_started = True
self._table_commit.commit(
messages,
@@ -332,6 +755,9 @@ class PaimonLeRobotWriter:
raise
self._close_batch()
self.pending_episodes = 0
+ self._pending_episode_rows = []
+ self._committed_task_count = len(self._task_indices)
+ self._subtasks_committed = True
return None
def finalize(self):
@@ -374,7 +800,6 @@ class PaimonLeRobotWriter:
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,
diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_writer_test.py
b/paimon-python/pypaimon/tests/multimodal_lerobot_writer_test.py
index f446de7581..26fb37a248 100644
--- a/paimon-python/pypaimon/tests/multimodal_lerobot_writer_test.py
+++ b/paimon-python/pypaimon/tests/multimodal_lerobot_writer_test.py
@@ -16,6 +16,8 @@
# under the License.
import io
+import importlib
+import json
import shutil
import tempfile
import unittest
@@ -23,19 +25,41 @@ from pathlib import Path
from unittest.mock import patch
import numpy as np
+import pyarrow as pa
import pypaimon.multimodal as pmm
from pypaimon.multimodal.lerobot import PaimonLeRobotWriter
+from pypaimon.multimodal.lerobot.metadata import (
+ _append_arrow,
+ _restore_pandas_metadata,
+)
+from pypaimon.multimodal.lerobot.writer import _read_arrow
try:
from PIL import Image
except ImportError:
Image = None
+try:
+ importlib.import_module("lerobot.datasets.compute_stats")
+ LEROBOT_AVAILABLE = True
+except ImportError:
+ LEROBOT_AVAILABLE = False
+
+
+def _catalog_rows(connection, name):
+ table = connection.catalog.get_table(connection._identifier(name))
+ builder = table.new_read_builder()
+ return builder.new_read().to_arrow(
+ builder.new_scan().plan().splits()).to_pylist()
+
class PaimonLeRobotWriterTest(unittest.TestCase):
def setUp(self):
+ if not LEROBOT_AVAILABLE and self._testMethodName != \
+
"test_missing_lerobot_stats_dependency_fails_before_table_creation":
+ self.skipTest("LeRobot is required for writer tests")
self.temp_dir = Path(tempfile.mkdtemp(
prefix="pypaimon_lerobot_writer_"))
self.connection = pmm.connect(options={
@@ -45,6 +69,205 @@ class PaimonLeRobotWriterTest(unittest.TestCase):
def tearDown(self):
shutil.rmtree(self.temp_dir, ignore_errors=True)
+ def test_creates_lerobot_table_group_on_finalize(self):
+ writer = PaimonLeRobotWriter(
+ self.connection,
+ "table_group",
+ fps=10,
+ features={
+ "action": {
+ "dtype": "float32",
+ "shape": (1,),
+ "names": None,
+ },
+ },
+ )
+ writer.add_frame({
+ "action": np.array([1.0], dtype=np.float32),
+ "task": "pick",
+ })
+ writer.save_episode()
+ writer.finalize()
+
+ frames = self.connection.get_table("table_group")
+ self.assertNotIn("task", [field.name for field in
frames.raw_table.fields])
+ options = frames.raw_table.table_schema.options
+ self.assertEqual("default.table_group__episodes", options[
+ "pypaimon.lerobot.episodes-table"])
+ self.assertEqual("default.table_group__tasks", options[
+ "pypaimon.lerobot.tasks-table"])
+ self.assertEqual("default.table_group__info", options[
+ "pypaimon.lerobot.info-table"])
+
+ episode = _catalog_rows(
+ self.connection, "table_group__episodes")[0]
+ self.assertEqual({
+ "episode_index": 0,
+ "dataset_from_index": 0,
+ "dataset_to_index": 1,
+ "tasks": ["pick"],
+ "length": 1,
+ }, {key: episode[key] for key in (
+ "episode_index", "dataset_from_index", "dataset_to_index",
+ "tasks", "length")})
+ self.assertEqual([{
+ "task_index": 0,
+ "task": "pick",
+ }], _catalog_rows(self.connection, "table_group__tasks"))
+ info = {
+ row["key"]: json.loads(row["value"])
+ for row in _catalog_rows(self.connection, "table_group__info")
+ }
+ self.assertEqual(10, info["fps"])
+ self.assertEqual(1, info["total_frames"])
+ self.assertEqual(1, info["total_episodes"])
+ self.assertEqual(1, info["total_tasks"])
+ self.assertEqual({"train": "0:1"}, info["splits"])
+ self.assertIn("action", info["features"])
+ self.assertEqual(
+ {"frames", "episodes", "tasks", "info", "stats"},
+ set(self.connection.create_lerobot_tag(
+ "table_group", "training")),
+ )
+
+ def test_writes_and_resumes_frame_subtasks(self):
+ features = {
+ "action": {
+ "dtype": "float32",
+ "shape": (1,),
+ "names": None,
+ },
+ "subtask_index": {
+ "dtype": "int64",
+ "shape": (1,),
+ "names": None,
+ },
+ }
+ writer = PaimonLeRobotWriter(
+ self.connection,
+ "with_subtasks",
+ fps=10,
+ features=features,
+ subtasks=["approach", "grasp"],
+ episodes_per_commit=1,
+ )
+ for index in (0, 1):
+ writer.add_frame({
+ "action": np.array([index], dtype=np.float32),
+ "subtask_index": np.array([index], dtype=np.int64),
+ "task": "pick",
+ })
+ writer.save_episode()
+ writer.add_frame({
+ "action": np.array([2], dtype=np.float32),
+ "subtask_index": np.array([0], dtype=np.int64),
+ "task": "pick",
+ })
+ writer.save_episode()
+ writer.finalize()
+
+ frames = self.connection.get_table("with_subtasks")
+ self.assertEqual([0, 1, 0], frames.scan().select([
+ "index", "subtask_index"
+ ]).to_arrow().sort_by("index").column(
+ "subtask_index").to_pylist())
+ self.assertEqual([
+ {"subtask_index": 0, "subtask": "approach"},
+ {"subtask_index": 1, "subtask": "grasp"},
+ ], _catalog_rows(self.connection, "with_subtasks__subtasks"))
+ for component, expected in (
+ ("tasks", ["pick"]),
+ ("subtasks", ["approach", "grasp"])):
+ table = self.connection.catalog.get_table(
+ self.connection._identifier(
+ "with_subtasks__%s" % component))
+ data = _restore_pandas_metadata(
+ table, _read_arrow(table)).to_pandas()
+ self.assertEqual(expected, data.index.tolist())
+ self.assertEqual(
+ "default.with_subtasks__subtasks",
+ frames.raw_table.table_schema.options[
+ "pypaimon.lerobot.subtasks-table"],
+ )
+ self.assertEqual(
+ {"frames", "episodes", "tasks", "info", "stats", "subtasks"},
+ set(self.connection.create_lerobot_tag(
+ "with_subtasks", "training")),
+ )
+
+ with self.assertRaisesRegex(ValueError, "do not match"):
+ PaimonLeRobotWriter(
+ self.connection,
+ "with_subtasks",
+ fps=10,
+ features=features,
+ subtasks=["approach", "release"],
+ )
+
+ resumed = PaimonLeRobotWriter(
+ self.connection,
+ "with_subtasks",
+ fps=10,
+ features=features,
+ )
+ self.assertEqual(("approach", "grasp"), resumed.subtasks)
+ self.assertEqual(2, resumed.num_episodes)
+ resumed.finalize()
+
+ def test_validates_subtask_contract_before_buffering(self):
+ action = {
+ "action": {
+ "dtype": "float32",
+ "shape": (1,),
+ "names": None,
+ },
+ }
+ with self.assertRaisesRegex(ValueError, "require a subtask_index"):
+ PaimonLeRobotWriter(
+ self.connection,
+ "subtasks_without_feature",
+ fps=10,
+ features=action,
+ subtasks=["approach"],
+ )
+
+ features = dict(action)
+ features["subtask_index"] = {
+ "dtype": "int64",
+ "shape": (1,),
+ "names": None,
+ }
+ with self.assertRaisesRegex(ValueError, "subtasks are required"):
+ PaimonLeRobotWriter(
+ self.connection,
+ "missing_subtasks",
+ fps=10,
+ features=features,
+ )
+ with self.assertRaisesRegex(ValueError, "duplicates"):
+ PaimonLeRobotWriter(
+ self.connection,
+ "duplicate_subtasks",
+ fps=10,
+ features=features,
+ subtasks=["approach", "approach"],
+ )
+
+ writer = PaimonLeRobotWriter(
+ self.connection,
+ "invalid_subtask_index",
+ fps=10,
+ features=features,
+ subtasks=["approach"],
+ )
+ with self.assertRaisesRegex(ValueError, "outside"):
+ writer.add_frame({
+ "action": np.array([1], dtype=np.float32),
+ "subtask_index": np.array([1], dtype=np.int64),
+ "task": "pick",
+ })
+ self.assertFalse(writer.has_pending_frames())
+
def test_default_commits_only_on_finalize_and_returns_none(self):
writer = PaimonLeRobotWriter(
self.connection,
@@ -115,13 +338,62 @@ class PaimonLeRobotWriterTest(unittest.TestCase):
resumed.finalize()
rows = self.connection.get_table("resume").scan().select([
- "episode_index", "frame_index", "index", "task_index", "task"
+ "episode_index", "frame_index", "index", "task_index"
]).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_snapshot_state_does_not_duplicate_task_metadata(self):
+ writer = PaimonLeRobotWriter(
+ self.connection,
+ "snapshot_state",
+ fps=10,
+ features={
+ "action": {
+ "dtype": "float32",
+ "shape": (1,),
+ "names": None,
+ },
+ },
+ )
+ writer.add_frame({
+ "action": np.array([1.0], dtype=np.float32),
+ "task": "pick",
+ })
+ writer.save_episode()
+ writer.finalize()
+
+ snapshot = self.connection.get_table(
+ "snapshot_state").raw_table.snapshot_manager() \
+ .get_latest_snapshot()
+ self.assertEqual("1", snapshot.properties[
+ "pypaimon.lerobot.state-version"])
+ self.assertNotIn(
+ "pypaimon.lerobot.task-indices", snapshot.properties)
+
+ def test_empty_frames_rejects_nonempty_companion_state(self):
+ features = {
+ "action": {
+ "dtype": "float32",
+ "shape": (1,),
+ "names": None,
+ },
+ }
+ PaimonLeRobotWriter(
+ self.connection, "inconsistent", fps=10, features=features)
+ tasks = self.connection.catalog.get_table(
+ self.connection._identifier("inconsistent__tasks"))
+ _append_arrow(tasks, pa.Table.from_pylist([{
+ "task_index": 0,
+ "task": "pick",
+ }], schema=_read_arrow(tasks).schema))
+
+ with self.assertRaisesRegex(ValueError, "inconsistent"):
+ PaimonLeRobotWriter(
+ self.connection, "inconsistent", fps=10, features=features)
+
def test_existing_table_requires_matching_feature_schema(self):
PaimonLeRobotWriter(
self.connection,
@@ -150,6 +422,51 @@ class PaimonLeRobotWriterTest(unittest.TestCase):
},
)
+ def test_rejects_subtask_feature_without_companion_support(self):
+ with self.assertRaisesRegex(ValueError, "subtask"):
+ PaimonLeRobotWriter(
+ self.connection,
+ "subtasks",
+ fps=10,
+ features={
+ "subtask_index": {
+ "dtype": "int64",
+ "shape": (1,),
+ "names": None,
+ },
+ },
+ )
+
+ def
test_missing_lerobot_stats_dependency_fails_before_table_creation(self):
+ self.connection.catalog.create_database(
+ "default", ignore_if_exists=True)
+ original_import = __import__
+
+ def reject_lerobot(name, *args, **kwargs):
+ if name.startswith("lerobot"):
+ raise ImportError("missing lerobot")
+ return original_import(name, *args, **kwargs)
+
+ with patch("builtins.__import__", side_effect=reject_lerobot):
+ with self.assertRaisesRegex(
+ ImportError, r"pypaimon\[lerobot\]"):
+ PaimonLeRobotWriter(
+ self.connection,
+ "missing_stats_dependency",
+ fps=10,
+ features={
+ "action": {
+ "dtype": "float32",
+ "shape": (1,),
+ "names": None,
+ },
+ },
+ )
+ self.assertNotIn(
+ "missing_stats_dependency",
+ self.connection.catalog.list_tables("default"),
+ )
+
def test_commits_multiple_completed_episodes_as_one_batch(self):
writer = PaimonLeRobotWriter(
self.connection,
@@ -195,7 +512,6 @@ class PaimonLeRobotWriterTest(unittest.TestCase):
"timestamp",
"index",
"task_index",
- "task",
"observation.state",
"action",
]).to_arrow().sort_by("index").to_pylist()
@@ -203,9 +519,178 @@ class PaimonLeRobotWriterTest(unittest.TestCase):
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])
+ self.assertEqual([
+ {"task_index": 0, "task": "pick"},
+ {"task_index": 1, "task": "place"},
+ ], _catalog_rows(self.connection, "robot_data__tasks"))
+
+ def test_multiple_flushes_append_rows_and_replace_info(self):
+ writer = PaimonLeRobotWriter(
+ self.connection,
+ "multiple_flushes",
+ fps=10,
+ episodes_per_commit=1,
+ features={
+ "action": {
+ "dtype": "float32",
+ "shape": (1,),
+ "names": None,
+ },
+ },
+ )
+ for value, task in ((1.0, "pick"), (2.0, "place")):
+ writer.add_frame({
+ "action": np.array([value], dtype=np.float32),
+ "task": task,
+ })
+ writer.save_episode()
+ writer.finalize()
+
+ self.assertEqual(2, len(_catalog_rows(
+ self.connection, "multiple_flushes__episodes")))
+ self.assertEqual(2, len(_catalog_rows(
+ self.connection, "multiple_flushes__tasks")))
+ info_rows = _catalog_rows(self.connection, "multiple_flushes__info")
+ self.assertEqual(len({row["key"] for row in info_rows}),
+ len(info_rows))
+ info = {
+ row["key"]: json.loads(row["value"])
+ for row in info_rows
+ }
+ self.assertEqual(2, info["total_frames"])
+ self.assertEqual(2, info["total_episodes"])
+ self.assertEqual(2, info["total_tasks"])
+ self.assertEqual({"train": "0:2"}, info["splits"])
+
+ def test_persists_native_episode_and_global_stats(self):
+ from lerobot.datasets.compute_stats import (
+ aggregate_stats,
+ compute_episode_stats,
+ )
+
+ writer = PaimonLeRobotWriter(
+ self.connection,
+ "stats",
+ fps=10,
+ features={
+ "action": {
+ "dtype": "float32",
+ "shape": (1,),
+ "names": None,
+ },
+ },
+ )
+ for values in ((1.0, 3.0), (5.0, 7.0)):
+ for value in values:
+ writer.add_frame({
+ "action": np.array([value], dtype=np.float32),
+ "task": "pick",
+ })
+ writer.save_episode()
+ writer.finalize()
+
+ options = self.connection.get_table(
+ "stats").raw_table.table_schema.options
+ self.assertEqual("default.stats__stats", options[
+ "pypaimon.lerobot.stats-table"])
+ episodes = _catalog_rows(self.connection, "stats__episodes")
+ native_episodes = [
+ compute_episode_stats(
+ {"action": np.asarray(values, dtype=np.float32)},
+ {"action": {"dtype": "float32", "shape": (1,)}},
+ )
+ for values in ((1.0, 3.0), (5.0, 7.0))
+ ]
+ for stat, expected in native_episodes[0]["action"].items():
+ np.testing.assert_allclose(
+ expected,
+ episodes[0]["stats/action/%s" % stat],
+ )
+ self.assertEqual(
+ {"min", "max", "mean", "std", "count",
+ "q01", "q10", "q50", "q90", "q99"},
+ {
+ name.removeprefix("stats/action/")
+ for name in episodes[0]
+ if name.startswith("stats/action/")
+ },
+ )
+
+ stats = {
+ row["key"]: json.loads(row["value"])
+ for row in _catalog_rows(self.connection, "stats__stats")
+ }
+ expected_stats = aggregate_stats(native_episodes)["action"]
+ for stat, expected in expected_stats.items():
+ np.testing.assert_allclose(expected, stats["action"][stat])
+ self.assertIn("timestamp", stats)
+ self.assertIn("index", stats)
+
+ def test_aggregates_numeric_feature_with_image_in_its_name(self):
+ features = {
+ "observation.image_embedding": {
+ "dtype": "float32",
+ "shape": (2,),
+ "names": None,
+ },
+ }
+ writer = PaimonLeRobotWriter(
+ self.connection,
+ "numeric_image_name",
+ fps=10,
+ features=features,
+ )
+ for offset in (0.0, 4.0):
+ for value in (1.0, 2.0):
+ writer.add_frame({
+ "observation.image_embedding": np.array(
+ [value + offset, value + offset + 1],
+ dtype=np.float32,
+ ),
+ "task": "inspect",
+ })
+ writer.save_episode()
+ writer.finalize()
+
+ resumed = PaimonLeRobotWriter(
+ self.connection,
+ "numeric_image_name",
+ fps=10,
+ features=features,
+ )
+ self.assertEqual(2, resumed.num_episodes)
+ resumed.finalize()
+
+ def test_metadata_read_projects_only_resume_columns(self):
+ writer = PaimonLeRobotWriter(
+ self.connection,
+ "projected_metadata",
+ fps=10,
+ features={
+ "action": {
+ "dtype": "float32",
+ "shape": (1,),
+ "names": None,
+ },
+ },
+ )
+ writer.add_frame({
+ "action": np.array([1.0], dtype=np.float32),
+ "task": "pick",
+ })
+ writer.save_episode()
+ writer.finalize()
+
+ episodes = self.connection.catalog.get_table(
+ self.connection._identifier("projected_metadata__episodes"))
+ projected = _read_arrow(
+ episodes, ["episode_index", "stats/index/count"])
+ self.assertEqual(
+ ["episode_index", "stats/index/count"],
+ projected.column_names,
+ )
@unittest.skipUnless(Image is not None, "Pillow is required for image
tests")
def test_writes_raw_image_frame_as_png_blob(self):
@@ -237,6 +722,20 @@ class PaimonLeRobotWriterTest(unittest.TestCase):
image = Image.open(io.BytesIO(blobs["observation.image"][0]))
self.assertEqual((5, 4), image.size)
self.assertEqual((73, 73, 73), image.getpixel((0, 0)))
+ episode = _catalog_rows(self.connection, "images__episodes")[0]
+ np.testing.assert_allclose(
+ np.full((3, 1, 1), 73 / 255.0),
+ episode["stats/observation.image/mean"],
+ )
+ stats = {
+ row["key"]: json.loads(row["value"])
+ for row in _catalog_rows(self.connection, "images__stats")
+ }
+ np.testing.assert_allclose(
+ np.full((3, 1, 1), 73 / 255.0),
+ stats["observation.image"]["mean"],
+ )
+ self.assertEqual([1], stats["observation.image"]["count"])
@unittest.skipUnless(Image is not None, "Pillow is required for image
tests")
def test_writes_native_hwc_image_frame_as_png_blob(self):
@@ -342,22 +841,24 @@ class PaimonLeRobotWriterTest(unittest.TestCase):
writer.finalize()
rows = self.connection.get_table("rerecord").scan().select([
- "episode_index", "index", "task", "action"
+ "episode_index", "index", "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)
+ self.assertEqual([
+ {"task_index": 0, "task": "keep"},
+ {"task_index": 1, "task": "also keep"},
+ ], _catalog_rows(self.connection, "rerecord__tasks"))
with self.assertRaisesRegex(RuntimeError, "after finalize"):
writer.add_frame({
"action": np.array([2.0], dtype=np.float32),
