diff --git a/docs/docs/pypaimon/multimodal-api.mdx b/docs/docs/pypaimon/multimodal-api.mdx index ef0a38e28c87..8ba3644a29e7 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 `__episodes`, +`
__tasks`, `
__info`, and `
__stats` store episode +boundaries, task labels, JSON-encoded dataset information, and global +statistics. When `subtask_index` is declared, `
__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( "robot_data", 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. ## Overwrite diff --git a/paimon-python/pypaimon/multimodal/lerobot/metadata.py b/paimon-python/pypaimon/multimodal/lerobot/metadata.py index 0f568f47693a..a30459f7e082 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/metadata.py +++ b/paimon-python/pypaimon/multimodal/lerobot/metadata.py @@ -51,6 +51,10 @@ 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), @@ -310,6 +314,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 faa3c6b5bdb4..f73fadf1103e 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 @@ _feature_shape, _schema_from_info, _validate_lerobot_schema, + _validate_v3_required_features, ) from pypaimon.multimodal.table import _target_schema @@ -47,12 +63,204 @@ "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 @@ def __init__( *, 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 @@ def __init__( 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 @@ def __init__( 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 @@ def add_frame(self, frame): 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 @@ def save_episode(self): 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 @@ def save_episode(self): 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 @@ def flush(self): 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 @@ def flush(self): 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 @@ def _episode_table(self): 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 f446de75818f..26fb37a248fb 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 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 @@ def setUp(self): 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 @@ def test_existing_table_resumes_global_and_task_indices(self): 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 @@ def test_existing_table_requires_matching_feature_schema(self): }, ) + 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 @@ def test_commits_multiple_completed_episodes_as_one_batch(self): "timestamp", "index", "task_index", - "task", "observation.state", "action", ]).to_arrow().sort_by("index").to_pylist() @@ -203,9 +519,178 @@ def test_commits_multiple_completed_episodes_as_one_batch(self): 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 @@ def test_writes_raw_image_frame_as_png_blob(self): 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 @@ def test_discards_rerecorded_episode_and_finalizes_tail_batch(self): 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),