From 5898bc9cfa0907e2eaf0a8efd0f0de11606577e0 Mon Sep 17 00:00:00 2001 From: victorwon2001 <192616110+victorwon2001@users.noreply.github.com> Date: Fri, 4 Sep 2026 00:09:13 +0900 Subject: [PATCH] refactor(lerobot): type metadata at import boundary --- src/hflow/importers/lerobot.py | 347 +++++++++++++++++--------------- tests/test_lerobot_converter.py | 321 +++++++++++++++++++++-------- 2 files changed, 421 insertions(+), 247 deletions(-) diff --git a/src/hflow/importers/lerobot.py b/src/hflow/importers/lerobot.py index 309d440..53b04aa 100644 --- a/src/hflow/importers/lerobot.py +++ b/src/hflow/importers/lerobot.py @@ -24,9 +24,9 @@ import tempfile import urllib.request from collections.abc import Sequence -from dataclasses import dataclass +from dataclasses import dataclass, replace from pathlib import Path -from typing import NotRequired, TypedDict +from typing import TypedDict from urllib.parse import urlsplit from mcap.writer import Writer as McapWriter @@ -50,7 +50,17 @@ HUGGING_FACE_TOKEN_ENVIRONMENT_VARIABLES = ("HF_TOKEN", "HUGGING_FACE_HUB_TOKEN") -class _EpisodeRow(TypedDict): +@dataclass(frozen=True) +class _VideoWindow: + camera_key: str + chunk_index: str + file_index: str + from_timestamp: float + to_timestamp: float + + +@dataclass(frozen=True) +class _EpisodeRow: episode_index: int task: str length: int @@ -58,14 +68,7 @@ class _EpisodeRow(TypedDict): data_file: str data_from: int data_to: int - video_windows: NotRequired[dict[str, "_VideoWindow"]] - - -class _VideoWindow(TypedDict): - chunk_index: str - file_index: str - from_timestamp: float - to_timestamp: float + video_windows: tuple[_VideoWindow, ...] = () @dataclass(frozen=True) @@ -75,19 +78,31 @@ class DatasetSource: license: str +@dataclass(frozen=True) +class _NumericFeatureSpecification: + feature_name: str + dimension: int + + +@dataclass(frozen=True) +class _DatasetInformation: + fps: int | float + data_path_template: str + video_path_template: str + robot_type: str | None + video_feature_names: tuple[str, ...] + numeric_features: tuple[_NumericFeatureSpecification, ...] + + class _DatasetRepositoryInformation(TypedDict): sha: str license: str -class _SourceArchive(TypedDict): - info: dict - fps: int | float - data_path: str - video_path: str - episodes: list[_EpisodeRow] - video_keys: list[str] - numeric_features: dict[str, dict] +@dataclass(frozen=True) +class _SourceArchive: + dataset_information: _DatasetInformation + episodes: tuple[_EpisodeRow, ...] cache_dir: Path dataset: DatasetSource @@ -349,6 +364,81 @@ def _fetch_info_json(repo_id: str, revision: str, cache_dir: Path) -> dict: return dataset_information +def _parse_dataset_information(dataset_information: dict) -> _DatasetInformation: + frames_per_second = dataset_information.get("fps") + if ( + isinstance(frames_per_second, bool) + or not isinstance(frames_per_second, int | float) + or not math.isfinite(frames_per_second) + or frames_per_second <= 0 + ): + raise ValueError( + f"LeRobot meta/info.json has invalid fps={frames_per_second!r}; " + "FPS must be finite and positive" + ) + + data_path_template = dataset_information.get("data_path") + if not isinstance(data_path_template, str) or not data_path_template.strip(): + raise ValueError("LeRobot meta/info.json must define a non-empty data_path template") + + video_path_template = dataset_information.get( + "video_path", "videos/{camera_key}/{chunk_index:06d}/{file_index:06d}.mp4" + ) + if not isinstance(video_path_template, str) or not video_path_template.strip(): + raise ValueError("LeRobot meta/info.json must define a non-empty video_path template") + + robot_type = dataset_information.get("robot_type") + if robot_type is not None and not isinstance(robot_type, str): + raise ValueError("LeRobot meta/info.json robot_type must be a string or null") + + dataset_features = dataset_information.get("features") + if not isinstance(dataset_features, dict): + raise ValueError("LeRobot meta/info.json must define a features object") + + video_feature_names: list[str] = [] + numeric_features: list[_NumericFeatureSpecification] = [] + for feature_name, feature_specification in dataset_features.items(): + if not isinstance(feature_specification, dict): + continue + declared_dtype = feature_specification.get("dtype") + if declared_dtype == "video": + video_feature_names.append(feature_name) + continue + if declared_dtype != "float32" or not ( + feature_name == "action" or feature_name.startswith("observation.") + ): + continue + + declared_shape = feature_specification.get("shape") or [] + if ( + not isinstance(declared_shape, list) + or len(declared_shape) != 1 + or isinstance(declared_shape[0], bool) + or not isinstance(declared_shape[0], int) + or declared_shape[0] < 1 + ): + raise ValueError( + f"unsupported feature {feature_name}: " + f"dtype={declared_dtype}, shape={declared_shape} " + "(only 1-D fixed-width numeric vectors are supported)" + ) + numeric_features.append( + _NumericFeatureSpecification( + feature_name=feature_name, + dimension=int(declared_shape[0]), + ) + ) + + return _DatasetInformation( + fps=frames_per_second, + data_path_template=data_path_template, + video_path_template=video_path_template, + robot_type=robot_type, + video_feature_names=tuple(sorted(video_feature_names)), + numeric_features=tuple(sorted(numeric_features, key=lambda feature: feature.feature_name)), + ) + + def _download_file(url: str, destination_path: Path, chunk_size: int = 1 << 20) -> None: destination_path.parent.mkdir(parents=True, exist_ok=True) with tempfile.NamedTemporaryFile( @@ -379,28 +469,9 @@ def _ensure_source_archive(dataset_source: DatasetSource, cache_dir: Path) -> _S f"{dataset_source.repo_id}/resolve/{dataset_source.revision}" ) metadata_directory = cache_dir / "meta" - dataset_information = _fetch_info_json( - dataset_source.repo_id, dataset_source.revision, cache_dir - ) - frames_per_second = dataset_information.get("fps") - if ( - isinstance(frames_per_second, bool) - or not isinstance(frames_per_second, int | float) - or not math.isfinite(frames_per_second) - or frames_per_second <= 0 - ): - raise ValueError( - f"LeRobot meta/info.json has invalid fps={frames_per_second!r}; " - "FPS must be finite and positive" - ) - data_path_template = dataset_information.get("data_path") - if not isinstance(data_path_template, str) or not data_path_template.strip(): - raise ValueError("LeRobot meta/info.json must define a non-empty data_path template") - video_path_template = dataset_information.get( - "video_path", "videos/{camera_key}/{chunk_index:06d}/{file_index:06d}.mp4" + dataset_information = _parse_dataset_information( + _fetch_info_json(dataset_source.repo_id, dataset_source.revision, cache_dir) ) - if not isinstance(video_path_template, str) or not video_path_template.strip(): - raise ValueError("LeRobot meta/info.json must define a non-empty video_path template") # Determine the episodes parquet location (v3 uses meta/episodes/*.parquet) episodes_metadata_directory = metadata_directory / "episodes" @@ -438,17 +509,17 @@ def _ensure_source_archive(dataset_source: DatasetSource, cache_dir: Path) -> _S else: task = str(tasks or "") episode_rows.append( - { - "episode_index": int(parquet_episode_row[0]), - "task": task, - "length": int(parquet_episode_row[2]), - "data_chunk": str(parquet_episode_row[3]).split("/")[-1], - "data_file": str(parquet_episode_row[4]).split("/")[-1], - "data_from": int(parquet_episode_row[5]), - "data_to": int(parquet_episode_row[6]), - } + _EpisodeRow( + episode_index=int(parquet_episode_row[0]), + task=task, + length=int(parquet_episode_row[2]), + data_chunk=str(parquet_episode_row[3]).split("/")[-1], + data_file=str(parquet_episode_row[4]).split("/")[-1], + data_from=int(parquet_episode_row[5]), + data_to=int(parquet_episode_row[6]), + ) ) - episode_rows.sort(key=lambda episode: episode["episode_index"]) + episode_rows.sort(key=lambda episode: episode.episode_index) # Video window columns: videos//{chunk_index,file_index,from_timestamp,to_timestamp} flattened_column_names = [ @@ -468,7 +539,7 @@ def _ensure_source_archive(dataset_source: DatasetSource, cache_dir: Path) -> _S ) # Gather video windows per episode+camera - video_windows_by_episode: dict[int, dict[str, _VideoWindow]] = {} + video_windows_by_episode: dict[int, list[_VideoWindow]] = {} video_window_selectors: list[str] = [] for camera_key in video_keys: video_window_selectors += [ @@ -490,86 +561,43 @@ def _ensure_source_archive(dataset_source: DatasetSource, cache_dir: Path) -> _S zip(video_window_column_names, video_window_row, strict=True) ) episode_index = int(video_window_by_column["episode_index"]) - video_windows_by_episode[episode_index] = {} + video_windows_by_episode[episode_index] = [] for camera_key in video_keys: video_chunk_index = video_window_by_column.get(f"vc_{camera_key}") video_file_index = video_window_by_column.get(f"vf_{camera_key}") - video_windows_by_episode[episode_index][camera_key] = { - "chunk_index": ( - "" if video_chunk_index is None else str(video_chunk_index).split("/")[-1] - ), - "file_index": ( - "" if video_file_index is None else str(video_file_index).split("/")[-1] - ), - "from_timestamp": float( - video_window_by_column.get(f"vfrom_{camera_key}") or 0.0 - ), - "to_timestamp": float(video_window_by_column.get(f"vto_{camera_key}") or 0.0), - } - for episode_row in episode_rows: - episode_row["video_windows"] = dict( - video_windows_by_episode.get(episode_row["episode_index"], {}) + video_windows_by_episode[episode_index].append( + _VideoWindow( + camera_key=camera_key, + chunk_index=( + "" + if video_chunk_index is None + else str(video_chunk_index).split("/")[-1] + ), + file_index=( + "" if video_file_index is None else str(video_file_index).split("/")[-1] + ), + from_timestamp=float( + video_window_by_column.get(f"vfrom_{camera_key}") or 0.0 + ), + to_timestamp=float(video_window_by_column.get(f"vto_{camera_key}") or 0.0), + ) + ) + episode_rows = [ + replace( + episode_row, + video_windows=tuple(video_windows_by_episode.get(episode_row.episode_index, [])), ) + for episode_row in episode_rows + ] finally: connection.close() - dataset_features = dataset_information.get("features") - if not isinstance(dataset_features, dict): - raise ValueError("LeRobot meta/info.json must define a features object") - video_keys_from_schema = sorted( - feature_name - for feature_name, feature_specification in dataset_features.items() - if isinstance(feature_specification, dict) and feature_specification.get("dtype") == "video" + return _SourceArchive( + dataset_information=dataset_information, + episodes=tuple(episode_rows), + cache_dir=cache_dir, + dataset=dataset_source, ) - if not video_keys: - video_keys = video_keys_from_schema - numeric_features = { - feature_name: feature_specification - for feature_name, feature_specification in dataset_features.items() - if isinstance(feature_specification, dict) - and feature_specification.get("dtype") == "float32" - } - - return { - "info": dataset_information, - "fps": frames_per_second, - "data_path": data_path_template, - "video_path": video_path_template, - "episodes": episode_rows, - "video_keys": video_keys, - "numeric_features": numeric_features, - "cache_dir": cache_dir, - "dataset": dataset_source, - } - - -@dataclass -class _NumericSchema: - name: str - dim: int - - -def _derive_numeric_schema(feature_name: str, feature_specification: dict) -> _NumericSchema: - declared_dtype = feature_specification.get("dtype") - declared_shape = feature_specification.get("shape") or [] - if declared_dtype != "float32": - raise ValueError( - f"unsupported feature {feature_name}: " - f"dtype={declared_dtype}, shape={declared_shape} " - "(only float32 fixed-width numeric vectors are supported)" - ) - if ( - len(declared_shape) != 1 - or isinstance(declared_shape[0], bool) - or not isinstance(declared_shape[0], int) - or declared_shape[0] < 1 - ): - raise ValueError( - f"unsupported feature {feature_name}: " - f"dtype={declared_dtype}, shape={declared_shape} " - "(only 1-D fixed-width numeric vectors are supported)" - ) - return _NumericSchema(name=feature_name, dim=int(declared_shape[0])) def import_lerobot_dataset( @@ -629,28 +657,26 @@ def import_lerobot_dataset( cache_directory, ) + available_video_feature_names = source_archive.dataset_information.video_feature_names for camera_key in resolved_camera_keys: - if camera_key not in source_archive["video_keys"]: + if camera_key not in available_video_feature_names: raise ValueError( f"camera key '{camera_key}' not found in dataset. " - f"Available: {source_archive['video_keys']}" + f"Available: {list(available_video_feature_names)}" ) - # Numeric schemas derived from metadata (fail before any conversion) - numeric_schemas = { - feature_name: _derive_numeric_schema(feature_name, feature_specification) - for feature_name, feature_specification in source_archive["numeric_features"].items() - if feature_name in ("observation.state", "action") - or feature_name.startswith("observation.") + numeric_features = { + feature.feature_name: feature + for feature in source_archive.dataset_information.numeric_features } - missing_required_features = {"observation.state", "action"} - numeric_schemas.keys() + missing_required_features = {"observation.state", "action"} - numeric_features.keys() if missing_required_features: raise ValueError( "LeRobot dataset is missing supported required features: " + ", ".join(sorted(missing_required_features)) ) - episode_rows = source_archive["episodes"] + episode_rows = source_archive.episodes if episode_index is not None and episode_index >= len(episode_rows): raise ValueError( f"episode_index {episode_index} is out of range for {len(episode_rows)} episode(s)" @@ -660,7 +686,7 @@ def import_lerobot_dataset( ) canonical_episode_paths: list[Path] = [] - dataset_source = source_archive["dataset"] + dataset_source = source_archive.dataset for selected_episode_index in selected_episode_indexes: canonical_episode_path = _convert_single_episode( source_archive=source_archive, @@ -668,8 +694,8 @@ def import_lerobot_dataset( output_dir=output_dir, episode_index=selected_episode_index, camera_keys=resolved_camera_keys, - numeric_schemas=numeric_schemas, - frames_per_second=int(source_archive["fps"]), + numeric_features=numeric_features, + frames_per_second=int(source_archive.dataset_information.fps), ) canonical_episode_paths.append(canonical_episode_path) @@ -716,26 +742,26 @@ def _convert_single_episode( output_dir: Path, episode_index: int, camera_keys: tuple[str, ...], - numeric_schemas: dict[str, _NumericSchema], + numeric_features: dict[str, _NumericFeatureSpecification], frames_per_second: int, ) -> Path: """Convert a single episode to canonical MCAP. Returns output path.""" import duckdb - episode_row = source_archive["episodes"][episode_index] - if episode_row["length"] is None or episode_row["length"] < 1: + episode_row = source_archive.episodes[episode_index] + if episode_row.length < 1: raise ValueError(f"episode {episode_index} has no frames") dataset_base_url = ( "https://huggingface.co/datasets/" f"{dataset_source.repo_id}/resolve/{dataset_source.revision}" ) - cache_directory = source_archive["cache_dir"] + cache_directory = source_archive.cache_dir # Locate the data parquet for this episode - data_chunk_index = episode_row["data_chunk"] - data_file_index = episode_row["data_file"] - data_relative_path = source_archive["data_path"].format( + data_chunk_index = episode_row.data_chunk + data_file_index = episode_row.data_file + data_relative_path = source_archive.dataset_information.data_path_template.format( chunk_index=int(data_chunk_index), file_index=int(data_file_index) ) local_data_path = ( @@ -749,12 +775,15 @@ def _convert_single_episode( # Episode video windows per camera (v3 flat columns: videos//from_timestamp etc.) video_time_window_by_camera: dict[str, tuple[float, float]] = {} video_metadata_by_camera: dict[str, _VideoWindow] = {} + episode_video_windows_by_camera = { + video_window.camera_key: video_window for video_window in episode_row.video_windows + } for camera_key in camera_keys: - video_window = episode_row.get("video_windows", {}).get(camera_key) - if video_window: + video_window = episode_video_windows_by_camera.get(camera_key) + if video_window is not None: video_time_window_by_camera[camera_key] = ( - video_window["from_timestamp"], - video_window["to_timestamp"], + video_window.from_timestamp, + video_window.to_timestamp, ) video_metadata_by_camera[camera_key] = video_window @@ -774,8 +803,8 @@ def _convert_single_episode( ] else "frame_index" ) - data_start_index = int(episode_row["data_from"]) - data_end_index = int(episode_row["data_to"]) + data_start_index = episode_row.data_from + data_end_index = episode_row.data_to episode_data_rows = connection.execute( f"SELECT * FROM read_parquet('{escaped_data_path}') " f"WHERE {index_column_name} >= {data_start_index} " @@ -809,24 +838,24 @@ def _feature_rows(feature_name: str) -> list | None: # Per-camera video: download chunk video, slice to episode window, transcode video_data_by_camera: dict[str, tuple[list[bytes], list[float]]] = {} for camera_key in camera_keys: - camera_video_metadata = video_metadata_by_camera.get(camera_key, {}) + camera_video_metadata = video_metadata_by_camera.get(camera_key) video_time_window = video_time_window_by_camera.get(camera_key) video_start_seconds = video_time_window[0] if video_time_window is not None else 0.0 video_end_seconds = video_time_window[1] if video_time_window is not None else 0.0 video_chunk_index = ( - str(camera_video_metadata.get("chunk_index")) - if camera_video_metadata.get("chunk_index") is not None + camera_video_metadata.chunk_index + if camera_video_metadata is not None else (data_chunk_index or "0") ) video_chunk_index = video_chunk_index.split("/")[-1] video_file_index = ( - str(camera_video_metadata.get("file_index")) - if camera_video_metadata.get("file_index") is not None + camera_video_metadata.file_index + if camera_video_metadata is not None else (data_file_index or "0") ) video_file_index = video_file_index.split("/")[-1] - video_relative_path = source_archive["video_path"].format( + video_relative_path = source_archive.dataset_information.video_path_template.format( chunk_index=int(video_chunk_index or 0), file_index=int(video_file_index or 0), video_key=camera_key, @@ -897,8 +926,8 @@ def _feature_rows(feature_name: str) -> list | None: video_schema_data = build_file_descriptor_set(CompressedVideo).SerializeToString() state_schema_name = "lerobot_msgs/msg/State" action_schema_name = "lerobot_msgs/msg/Action" - state_schema_text = f"float32[{numeric_schemas['observation.state'].dim}] position" - action_schema_text = f"float32[{numeric_schemas['action'].dim}] action" + state_schema_text = f"float32[{numeric_features['observation.state'].dimension}] position" + action_schema_text = f"float32[{numeric_features['action'].dimension}] action" state_schema_data = state_schema_text.encode("utf-8") action_schema_data = action_schema_text.encode("utf-8") @@ -970,10 +999,10 @@ def _feature_rows(feature_name: str) -> list | None: mcap_writer.add_metadata( name="episode/v1", data={ - "task": str(episode_row["task"] or ""), + "task": episode_row.task or "", "operator": "lerobot_converter", "success": "true", - "embodiment": str(source_archive["info"].get("robot_type") or "unknown"), + "embodiment": source_archive.dataset_information.robot_type or "unknown", "source_dataset": dataset_source.repo_id, "source_revision": dataset_source.revision, "source_episode_index": str(episode_index), diff --git a/tests/test_lerobot_converter.py b/tests/test_lerobot_converter.py index 39a24a1..03dc3d7 100755 --- a/tests/test_lerobot_converter.py +++ b/tests/test_lerobot_converter.py @@ -7,18 +7,19 @@ import io import json +import re import shutil import subprocess import urllib.request from pathlib import Path -from typing import cast import pytest +from mcap.reader import make_reader import hflow.importers.lerobot as prep from hflow.cli import main as cli_main -_DERIVE = prep._derive_numeric_schema +_PARSE = prep._parse_dataset_information _ENCODE = prep._encode_cdr_float32_array _FFMPEG = shutil.which("ffmpeg") @@ -32,9 +33,8 @@ ) -def _build_fake_corpus(tmp_path: Path) -> dict: - """Synthetic v3 metadata: 4 episodes, 2 cameras, 6-dim state/action.""" - info = { +def _valid_info() -> dict: + return { "fps": 30, "data_path": "data/chunk-{chunk_index:03d}/file-{file_index:03d}.parquet", "video_path": "videos/{video_key}/chunk-{chunk_index:03d}/file-{file_index:03d}.mp4", @@ -47,6 +47,11 @@ def _build_fake_corpus(tmp_path: Path) -> dict: }, "robot_type": "so101", } + + +def _build_fake_corpus(tmp_path: Path) -> dict: + """Synthetic v3 metadata: 4 episodes, 2 cameras, 6-dim state/action.""" + info = _valid_info() (tmp_path / "meta").mkdir(parents=True, exist_ok=True) (tmp_path / "meta" / "info.json").write_text(json.dumps(info)) @@ -145,32 +150,80 @@ def _build_fake_corpus(tmp_path: Path) -> dict: } -def test_derive_numeric_schema_float32_vector() -> None: - schema = _DERIVE("observation.state", {"dtype": "float32", "shape": [6]}) - assert schema.name == "observation.state" - assert schema.dim == 6 +def test_parse_dataset_information_returns_typed_supported_metadata() -> None: + info = _valid_info() + info["unknown_top_level"] = {"future": True} + info["features"]["action"]["unused_upstream_field"] = "ignored" + + parsed = _PARSE(info) + + assert parsed.fps == 30 + assert parsed.data_path_template == info["data_path"] + assert parsed.video_path_template == info["video_path"] + assert parsed.robot_type == "so101" + assert parsed.video_feature_names == ( + "observation.images.side", + "observation.images.up", + ) + assert {feature.feature_name: feature.dimension for feature in parsed.numeric_features} == { + "action": 6, + "observation.state": 6, + } + assert isinstance(parsed.numeric_features, tuple) + +@pytest.mark.parametrize("shape", [[2, 3], [], [True], [False]]) +def test_parse_dataset_information_rejects_unsupported_numeric_shape(shape: list[object]) -> None: + info = _valid_info() + info["features"]["action"]["shape"] = shape -def test_derive_numeric_schema_rejects_unsupported() -> None: - with pytest.raises(ValueError, match="unsupported feature"): - _DERIVE("action", {"dtype": "float64", "shape": [6]}) - with pytest.raises(ValueError, match="unsupported feature"): - _DERIVE("observation.state", {"dtype": "float32", "shape": [2, 3]}) - with pytest.raises(ValueError, match="unsupported feature"): - _DERIVE("observation.state", {"dtype": "float32", "shape": []}) - for shape in ([True], [False]): - with pytest.raises( - ValueError, - match=rf"unsupported feature action: dtype=float32, shape=\[{shape[0]}\]", - ): - _DERIVE("action", {"dtype": "float32", "shape": shape}) + with pytest.raises( + ValueError, + match=rf"unsupported feature action: dtype=float32, shape={re.escape(str(shape))}", + ): + _PARSE(info) -def test_import_rejects_required_boolean_dimension_without_dataset_output( +@pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("data_path", "", "non-empty data_path"), + ("data_path", None, "non-empty data_path"), + ("video_path", "", "non-empty video_path"), + ("video_path", 3, "non-empty video_path"), + ("robot_type", 3, "robot_type must be a string or null"), + ], +) +def test_parse_dataset_information_rejects_invalid_supported_fields( + field: str, value: object, message: str +) -> None: + info = _valid_info() + info[field] = value + + with pytest.raises(ValueError, match=message): + _PARSE(info) + + +def test_parse_dataset_information_accepts_null_robot_type_and_default_video_path() -> None: + info = _valid_info() + info["robot_type"] = None + info.pop("video_path") + + parsed = _PARSE(info) + + assert parsed.robot_type is None + assert parsed.video_path_template == ( + "videos/{camera_key}/{chunk_index:06d}/{file_index:06d}.mp4" + ) + + +def test_import_rejects_required_unsupported_dtype_without_dataset_output( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: output_dir = tmp_path / "out" dataset_source = prep.DatasetSource(repo_id="fake/repo", revision="abc", license="apache-2.0") + info = _valid_info() + info["features"]["action"]["dtype"] = "float64" monkeypatch.setattr( prep, "_hf_repo_info", @@ -179,22 +232,23 @@ def test_import_rejects_required_boolean_dimension_without_dataset_output( monkeypatch.setattr( prep, "_ensure_source_archive", - lambda source, cache_dir: { - "numeric_features": { - "action": {"dtype": "float32", "shape": [True]}, - "observation.state": {"dtype": "float32", "shape": [6]}, - }, - "video_keys": [prep.DEFAULT_CAMERA_KEY], - "episodes": [], - "dataset": dataset_source, - }, + lambda source, cache_dir: prep._SourceArchive( + dataset_information=_PARSE(info), + episodes=(), + cache_dir=cache_dir, + dataset=dataset_source, + ), ) with pytest.raises( ValueError, - match=r"unsupported feature action: dtype=float32, shape=\[True\]", + match=r"missing supported required features: action", ): - prep.import_lerobot_dataset(dataset_repo="fake/repo", output_dir=output_dir) + prep.import_lerobot_dataset( + dataset_repo="fake/repo", + output_dir=output_dir, + camera_keys="observation.images.up", + ) assert not (output_dir / "landing").exists() assert not (output_dir / "prepared-manifest.json").exists() @@ -247,54 +301,66 @@ def fake_dl(url: str, dest: Path, **kw: object) -> None: ds = prep.DatasetSource(repo_id="fake/repo", revision="abc", license="apache-2.0") found = prep._ensure_source_archive(ds, tmp_path) - assert len(found["episodes"]) == 4 - assert found["episodes"][0]["length"] == 60 - assert found["episodes"][1]["length"] == 65 - assert found["episodes"][0]["data_from"] == 0 - assert found["episodes"][0]["data_to"] == 60 - assert set(found["video_keys"]) == {"observation.images.up", "observation.images.side"} - assert found["episodes"][0]["video_windows"]["observation.images.up"][ - "to_timestamp" + assert len(found.episodes) == 4 + assert found.episodes[0].length == 60 + assert found.episodes[1].length == 65 + assert found.episodes[0].data_from == 0 + assert found.episodes[0].data_to == 60 + assert set(found.dataset_information.video_feature_names) == { + "observation.images.up", + "observation.images.side", + } + assert isinstance(found.episodes, tuple) + assert isinstance(found.episodes[0].video_windows, tuple) + assert {window.camera_key: window.to_timestamp for window in found.episodes[0].video_windows}[ + "observation.images.up" ] == pytest.approx(2.0) +@pytest.mark.parametrize( + ("robot_type", "expected_embodiment"), + [("so101", "so101"), (None, "unknown")], +) def test_video_cache_distinguishes_file_indices_and_reuses_same_source( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + robot_type: str | None, + expected_embodiment: str, ) -> None: corpus = _build_fake_corpus(tmp_path) + corpus["info"]["robot_type"] = robot_type camera_key = "observation.images.up" dataset_source = prep.DatasetSource(repo_id="fake/repo", revision="abc", license="apache-2.0") - def episode_row(episode_index: int, video_file_index: int) -> dict: - return { - "episode_index": episode_index, - "task": f"task-{episode_index}", - "length": 1, - "data_chunk": "000", - "data_file": "000", - "data_from": 0, - "data_to": 1, - "video_windows": { - camera_key: { - "chunk_index": "000", - "file_index": f"{video_file_index:03d}", - "from_timestamp": 0.0, - "to_timestamp": 0.0, - } - }, - } + def episode_row(episode_index: int, video_file_index: int) -> prep._EpisodeRow: + return prep._EpisodeRow( + episode_index=episode_index, + task=f"task-{episode_index}", + length=1, + data_chunk="000", + data_file="000", + data_from=0, + data_to=1, + video_windows=( + prep._VideoWindow( + camera_key=camera_key, + chunk_index="000", + file_index=f"{video_file_index:03d}", + from_timestamp=0.0, + to_timestamp=0.0, + ), + ), + ) - source_archive = cast( - prep._SourceArchive, - { - **corpus, - "episodes": [episode_row(0, 0), episode_row(1, 1)], - "video_keys": [camera_key], - }, + source_archive = prep._SourceArchive( + dataset_information=_PARSE(corpus["info"]), + episodes=(episode_row(0, 0), episode_row(1, 1)), + cache_dir=tmp_path, + dataset=dataset_source, ) - numeric_schemas = { - "observation.state": prep._NumericSchema(name="observation.state", dim=6), - "action": prep._NumericSchema(name="action", dim=6), + numeric_features = { + feature.feature_name: feature + for feature in source_archive.dataset_information.numeric_features } video_downloaded_urls: set[str] = set() cache_path_by_url: dict[str, Path] = {} @@ -332,7 +398,7 @@ def fake_transcode(mp4_path: Path, *_args: object, **_kwargs: object) -> list[by output_dir=tmp_path / "output", episode_index=episode_index, camera_keys=(camera_key,), - numeric_schemas=numeric_schemas, + numeric_features=numeric_features, frames_per_second=30, ) @@ -349,6 +415,26 @@ def fake_transcode(mp4_path: Path, *_args: object, **_kwargs: object) -> list[by ] assert cache_path_by_url[video_urls[0]] != cache_path_by_url[video_urls[1]] + output_path = tmp_path / "output" / "landing" / "lerobot_episode_0001.mcap" + with output_path.open("rb") as stream: + reader = make_reader(stream, validate_crcs=True) + summary = reader.get_summary() + metadata_records = {record.name: record.metadata for record in reader.iter_metadata()} + assert summary is not None + schemas_by_name = {schema.name: schema.data for schema in summary.schemas.values()} + assert schemas_by_name["lerobot_msgs/msg/State"] == b"float32[6] position" + assert schemas_by_name["lerobot_msgs/msg/Action"] == b"float32[6] action" + assert metadata_records["episode/v1"] == { + "task": "task-0", + "operator": "lerobot_converter", + "success": "true", + "embodiment": expected_embodiment, + "source_dataset": "fake/repo", + "source_revision": "abc", + "source_episode_index": "0", + "converter_version": prep.CONVERTER_VERSION, + } + def test_camera_selection_validates_keys(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: corpus = _build_fake_corpus(tmp_path) @@ -413,26 +499,32 @@ def test_import_namespaces_source_cache_by_resolved_revision( lambda repo, revision: {"sha": resolved_shas[revision], "license": "apache-2.0"}, ) - def fake_ensure_source_archive(dataset_source: prep.DatasetSource, cache_dir: Path) -> dict: + def fake_ensure_source_archive( + dataset_source: prep.DatasetSource, cache_dir: Path + ) -> prep._SourceArchive: cache_dir.mkdir(parents=True, exist_ok=True) source_marker = cache_dir / "source-marker.txt" if not source_marker.exists(): source_marker.write_text(dataset_source.revision) cache_observations.append((dataset_source.revision, cache_dir, source_marker.read_text())) - return { - "info": {}, - "fps": 30, - "data_path": "data/{chunk_index}/{file_index}.parquet", - "video_path": "videos/{camera_key}/{chunk_index}/{file_index}.mp4", - "episodes": [], - "video_keys": [prep.DEFAULT_CAMERA_KEY], - "numeric_features": { - "action": {"dtype": "float32", "shape": [1]}, - "observation.state": {"dtype": "float32", "shape": [1]}, - }, - "cache_dir": cache_dir, - "dataset": dataset_source, - } + return prep._SourceArchive( + dataset_information=_PARSE( + { + "fps": 30, + "data_path": "data/{chunk_index}/{file_index}.parquet", + "video_path": "videos/{camera_key}/{chunk_index}/{file_index}.mp4", + "features": { + "action": {"dtype": "float32", "shape": [1]}, + "observation.state": {"dtype": "float32", "shape": [1]}, + prep.DEFAULT_CAMERA_KEY: {"dtype": "video"}, + }, + "robot_type": "so101", + } + ), + episodes=(), + cache_dir=cache_dir, + dataset=dataset_source, + ) monkeypatch.setattr(prep, "_ensure_source_archive", fake_ensure_source_archive) @@ -450,6 +542,17 @@ def fake_ensure_source_archive(dataset_source: prep.DatasetSource, cache_dir: Pa sha_a, sha_b, ] + assert json.loads((tmp_path / "prepared-manifest.json").read_text()) == { + "schema_version": 2, + "dataset": { + "repo_id": "fake/repo", + "revision": sha_a, + "license": "apache-2.0", + }, + "camera_keys": [prep.DEFAULT_CAMERA_KEY], + "episodes_converted": 0, + "converter_version": prep.CONVERTER_VERSION, + } @pytest.mark.parametrize("resolved_sha", ["../../evil", "/tmp/probe-328-absolute"]) @@ -711,6 +814,26 @@ def test_fetch_info_json_malformed_json_raises_contextual_value_error( assert isinstance(excinfo.value.__cause__, json.JSONDecodeError) +def test_fetch_info_json_caches_original_json_document( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + info = _valid_info() + info["unknown_top_level"] = {"preserve": [1, True, None]} + tree_body = json.dumps([{"path": "meta/info.json", "type": "file"}]).encode() + _stub_urlopen( + monkeypatch, + { + "recursive=true": tree_body, + "meta/info.json": json.dumps(info).encode(), + }, + ) + + fetched = prep._fetch_info_json("lerobot/pusht", "main", tmp_path) + + assert fetched == info + assert json.loads((tmp_path / "meta" / "info.json").read_text()) == info + + def test_hf_repo_info_valid_json_still_resolves(monkeypatch: pytest.MonkeyPatch) -> None: # A real resolved commit sha: at least the 7 hex characters the sha # validation requires, since a cache directory is named after it. @@ -889,6 +1012,28 @@ def fail_tree(repo: str, rev: str, path: str) -> list[dict]: assert not cache_dir.exists() or not (cache_dir / "meta" / "episodes").exists() +def test_info_json_refuses_unsupported_numeric_shape_before_episode_discovery( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + info = _valid_info() + info["features"]["action"]["shape"] = [True] + monkeypatch.setattr(prep, "_fetch_info_json", lambda repo, rev, cache: info) + + def fail_tree(repo: str, rev: str, path: str) -> list[dict]: + raise AssertionError("episode metadata discovery must not run after invalid shape") + + monkeypatch.setattr(prep, "_hf_tree", fail_tree) + + dataset_source = prep.DatasetSource(repo_id="fake/repo", revision="abc", license="apache-2.0") + cache_dir = tmp_path / "cache" + with pytest.raises( + ValueError, match=r"unsupported feature action: dtype=float32, shape=\[True\]" + ): + prep._ensure_source_archive(dataset_source, cache_dir) + + assert not cache_dir.exists() or not (cache_dir / "meta" / "episodes").exists() + + def test_info_json_accepts_normal_positive_fps( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: @@ -921,4 +1066,4 @@ def fake_dl(url: str, dest: Path, **kw: object) -> None: dataset_source = prep.DatasetSource(repo_id="fake/repo", revision="abc", license="apache-2.0") source_archive = prep._ensure_source_archive(dataset_source, tmp_path / "cache") - assert source_archive["fps"] == 30 + assert source_archive.dataset_information.fps == 30