diff --git a/.gitignore b/.gitignore index 15e2979..bcfb859 100644 --- a/.gitignore +++ b/.gitignore @@ -22,6 +22,7 @@ packages/hflow-server/src/hflow_server/static/ # Local maintainer tooling (agent skills, settings); not part of the public repo. .claude/ .agents/ +.zcode/ # Transient: `pnpm gen:api` dumps the schema here on its way to src/apiSchema.ts. .openapi.json diff --git a/src/hflow/importers/lerobot.py b/src/hflow/importers/lerobot.py index 7c6bfc8..e2c82e6 100644 --- a/src/hflow/importers/lerobot.py +++ b/src/hflow/importers/lerobot.py @@ -54,9 +54,13 @@ # can prove a landing file belongs to this exact selection (#303). Those # fields change the canonical bytes that content_episode_id hashes, so v5 # and v6 outputs must not share a converter identity. -CONVERTER_VERSION = "lerobot-converter-v6" +CONVERTER_VERSION = "lerobot-converter-v7" # Canonical transform knobs that affect published bytes for this importer. IMPORT_GOP_SECONDS = 1.0 +# The v3 per-episode aggregate of the collector's frame-level next.success +# label. Optional: a corpus that declares no outcome feature has no such column. +_OUTCOME_AGGREGATE_COLUMN = "stats/next.success/max" +_SUCCESS_DERIVATION = f"max({_OUTCOME_AGGREGATE_COLUMN.removesuffix('/max')})" PRESENTATION_TIMESTAMP_EPSILON_S = 0.050 EPISODE_METADATA_TREE_PREFIX = PurePosixPath("meta/episodes") @@ -75,6 +79,9 @@ class _EpisodeRow(TypedDict): data_from: int data_to: int video_windows: NotRequired[dict[str, "_VideoWindow"]] + # MAX of the episode's collector-labeled next.success frames, present only + # when the source declares that outcome feature at all. + success_outcome: NotRequired[bool] class _VideoWindow(TypedDict): @@ -584,12 +591,31 @@ def _ensure_source_archive(dataset_source: DatasetSource, cache_dir: Path) -> _S ) connection = duckdb.connect() try: + # Column discovery first: the outcome aggregate is optional in v3 + # (not every corpus declares a collector-labeled next.success), and + # naming a missing column would fail the read. A corpus without one + # is normal, not malformed. + episodes_columns = [ + column_description[0] + for column_description in connection.execute( + f"SELECT * FROM {episode_metadata_relation} LIMIT 1" + ).description + ] + # Built outside the f-string below: a quoted identifier cannot be + # nested in a same-quoted f-string on Python 3.11, which this repo + # still supports. + has_outcome_aggregate = _OUTCOME_AGGREGATE_COLUMN in episodes_columns + outcome_aggregate_selector = ( + f', "{_OUTCOME_AGGREGATE_COLUMN}"' if has_outcome_aggregate else "" + ) + episode_rows: list[_EpisodeRow] = [] parquet_episode_rows = connection.execute( f""" SELECT "episode_index", "tasks", "length", "data/chunk_index", "data/file_index", "dataset_from_index", "dataset_to_index" + {outcome_aggregate_selector} FROM {episode_metadata_relation} ORDER BY "episode_index" """ @@ -597,17 +623,25 @@ def _ensure_source_archive(dataset_source: DatasetSource, cache_dir: Path) -> _S for parquet_episode_row in parquet_episode_rows: tasks = parquet_episode_row[1] task = (str(tasks[0]) if tasks else "") if isinstance(tasks, list) else 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]), - } - ) + episode_row_value: _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]), + } + if has_outcome_aggregate: + # Episode outcome = MAX over the episode's collector-labeled + # next.success frames: any success frame makes the episode a + # success. An empty aggregate carries no label either way. + outcome_frames = parquet_episode_row[7] + if outcome_frames is not None and len(outcome_frames) > 0: + episode_row_value["success_outcome"] = any( + bool(value) for value in outcome_frames + ) + episode_rows.append(episode_row_value) # Video window columns: videos//{chunk_index,file_index,from_timestamp,to_timestamp} flattened_column_names = [ @@ -1193,20 +1227,30 @@ def _feature_rows(feature_name: str) -> list | None: sequence=frame_index, ) + episode_record: dict[str, str] = { + "task": str(episode_row["task"] or ""), + "operator": "lerobot_converter", + "embodiment": str(source_archive["info"].get("robot_type") or "unknown"), + "source_dataset": dataset_source.repo_id, + "source_revision": dataset_source.revision, + "source_episode_index": str(episode_index), + "converter_version": CONVERTER_VERSION, + "camera_keys": _encode_camera_keys(camera_keys), + "gop_seconds": f"{IMPORT_GOP_SECONDS:g}", + } + # success is the collector's label, never ours: when the source + # declares the outcome feature, report MAX over the episode's + # frames and name the derivation so the methodology travels with + # the data; when it does not, the key is omitted rather than + # invented (FORMAT.md: every episode/v1 key is optional and the + # record is copied/merged from the source recording). + success_outcome = episode_row.get("success_outcome") + if success_outcome is not None: + episode_record["success"] = "true" if success_outcome else "false" + episode_record["success_derivation"] = _SUCCESS_DERIVATION mcap_writer.add_metadata( name="episode/v1", - data={ - "task": str(episode_row["task"] or ""), - "operator": "lerobot_converter", - "success": "true", - "embodiment": str(source_archive["info"].get("robot_type") or "unknown"), - "source_dataset": dataset_source.repo_id, - "source_revision": dataset_source.revision, - "source_episode_index": str(episode_index), - "converter_version": CONVERTER_VERSION, - "camera_keys": _encode_camera_keys(camera_keys), - "gop_seconds": f"{IMPORT_GOP_SECONDS:g}", - }, + data=episode_record, ) mcap_writer.add_metadata( name="source-provenance/v1", diff --git a/tests/test_lerobot_converter.py b/tests/test_lerobot_converter.py index 99c6406..5b145c6 100755 --- a/tests/test_lerobot_converter.py +++ b/tests/test_lerobot_converter.py @@ -1866,3 +1866,247 @@ def should_not_convert(**_kwargs: object) -> prep._PublishedEpisode: manifest = json.loads((output_dir / "prepared-manifest.json").read_text()) assert manifest["episodes_converted"] == 0 assert len(manifest["episodes"]) == 2 + + +# --- success label: read the collector's outcome, never invent it (#395) ----- + + +def _build_success_label_corpus(root: Path, outcome_mode: str) -> dict: + """One two-frame episode. + + outcome_mode: 'transition', 'all-false', 'empty-aggregate', or 'none'. + """ + has_outcome = outcome_mode != "none" + info = { + "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", + "features": { + "action": {"dtype": "float32", "shape": [1]}, + "observation.state": {"dtype": "float32", "shape": [1]}, + "observation.images.up": {"dtype": "video", "shape": [480, 640, 3]}, + "timestamp": {"dtype": "float32", "shape": [1]}, + }, + "robot_type": "so101", + } + if has_outcome: + info["features"]["next.success"] = {"dtype": "bool", "shape": [1]} + (root / "meta").mkdir(parents=True, exist_ok=True) + (root / "meta" / "info.json").write_text(json.dumps(info)) + + import duckdb + + conn = duckdb.connect() + ep_cols = [ + "episode_index", + "length", + "data/chunk_index", + "data/file_index", + "dataset_from_index", + "dataset_to_index", + "videos/observation.images.up/chunk_index", + "videos/observation.images.up/file_index", + "videos/observation.images.up/from_timestamp", + "videos/observation.images.up/to_timestamp", + "tasks", + ] + row: list[object] = [0, 2, "000", "000", 0, 2, "000", "000", 0.0, 0.0, ["push the block"]] + if has_outcome: + if outcome_mode == "transition": + stats_min, stats_max = [False], [True] + elif outcome_mode == "empty-aggregate": + # The column exists but carries no value for this episode, which + # is a declared feature with nothing recorded rather than a label. + stats_min, stats_max = [], [] + else: + stats_min, stats_max = [False], [False] + ep_cols += ["stats/next.success/min", "stats/next.success/max"] + row += [stats_min, stats_max] + ep_path = root / "meta" / "episodes" / "chunk-000" / "file-000.parquet" + ep_path.parent.mkdir(parents=True, exist_ok=True) + vals = ( + "(" + + ",".join( + "[" + ",".join(str(bool(item)) for item in value) + "]" + if isinstance(value, list) and value and all(isinstance(item, bool) for item in value) + else "[" + ",".join(f"'{item}'" for item in value) + "]" + if isinstance(value, list) + else f"'{value}'" + if isinstance(value, str) + else str(value) + for value in row + ) + + ")" + ) + conn.execute( + f"COPY (SELECT * FROM (VALUES {vals}) AS t({','.join(chr(34) + c + chr(34) for c in ep_cols)})) " + f"TO '{str(ep_path).replace(chr(39), chr(39) * 2)}' (FORMAT parquet)" + ) + + frame_outcomes = [False, True] if outcome_mode == "transition" else [False, False] + data_cols = 'index, episode_index, frame_index, timestamp, "observation.state", action' + data_rows = [ + f"({index}, 0, {frame_index}, 0.0, [0.0], [0.5]" + for index, frame_index in enumerate(range(2)) + ] + if has_outcome: + data_cols += ', "next.success"' + data_rows = [ + data_row + f", {str(frame_outcomes[frame_index]).lower()})" + for frame_index, data_row in enumerate(data_rows) + ] + else: + data_rows = [data_row + ")" for data_row in data_rows] + data_path = root / "data" / "chunk-000" / "file-000.parquet" + data_path.parent.mkdir(parents=True, exist_ok=True) + conn.execute( + f"COPY (SELECT * FROM (VALUES {','.join(data_rows)}) AS t({data_cols})) " + f"TO '{str(data_path).replace(chr(39), chr(39) * 2)}' (FORMAT parquet)" + ) + conn.close() + return {"info": info} + + +def _import_success_label_corpus( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, outcome_mode: str +) -> Path: + root = tmp_path / "corpus" + corpus = _build_success_label_corpus(root, outcome_mode) + output_dir = tmp_path / "out" + + monkeypatch.setattr( + prep, "_hf_repo_info", lambda repo, revision: {"sha": "abc1234", "license": "apache-2.0"} + ) + monkeypatch.setattr(prep, "_fetch_info_json", lambda repo, rev, cache: corpus["info"]) + monkeypatch.setattr( + prep, + "_hf_tree", + lambda repo, rev, path: ( + [{"path": "meta/episodes/chunk-000/file-000.parquet", "type": "file"}] + if "episodes" in path + else [{"path": "meta/info.json", "type": "file"}] + ), + ) + + def fake_download(url: str, dest: Path, **_kwargs: object) -> None: + dest.parent.mkdir(parents=True, exist_ok=True) + if "meta/episodes" in url: + shutil.copy(root / "meta" / "episodes" / "chunk-000" / "file-000.parquet", dest) + elif url.endswith("info.json"): + shutil.copy(root / "meta" / "info.json", dest) + else: + shutil.copy(root / "data" / "chunk-000" / "file-000.parquet", dest) + + monkeypatch.setattr(prep, "_download_file", fake_download) + monkeypatch.setattr( + prep, + "_transcode_mp4_to_h264", + lambda mp4_path, gop, fps: ( + [ + b"\x00\x00\x00\x01\x09\x10\x00\x00\x00\x01\x67\x42\x00" + b"\x00\x00\x00\x01\x68\x88\x80\x00\x00\x00\x01\x65\x88" + ] + * 2 + ), + ) + monkeypatch.setattr(prep, "_get_video_pts_times", lambda path: [0, 0]) + monkeypatch.setattr(prep, "ffmpeg_version", lambda: "test-ffmpeg") + + prep.import_lerobot_dataset( + dataset_repo="fake/repo", + output_dir=output_dir, + camera_keys=("observation.images.up",), + ) + return output_dir + + +def test_success_label_reports_max_over_episode_frames( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """MAX over the collector's next.success frames: a False frame followed + by a True frame makes the episode a success, even though the LAST frame + is False. The derivation is stamped so the methodology travels.""" + from hflow.episode import Episode + + output_dir = _import_success_label_corpus(tmp_path, monkeypatch, "transition") + landing = sorted((output_dir / "landing").glob("*.mcap")) + with Episode(landing[0]) as episode: + record = episode.metadata_records["episode/v1"] + assert record["success"] == "true" + assert record["success_derivation"] == "max(stats/next.success)" + + +def test_success_label_reports_false_when_source_is_all_false( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """An all-false source label ships as 'false', never as an invented + 'true': the collector's judgment, reported verbatim.""" + from hflow.episode import Episode + + output_dir = _import_success_label_corpus(tmp_path, monkeypatch, "all-false") + landing = sorted((output_dir / "landing").glob("*.mcap")) + with Episode(landing[0]) as episode: + record = episode.metadata_records["episode/v1"] + assert record["success"] == "false" + assert record["success_derivation"] == "max(stats/next.success)" + + +def test_success_label_omitted_when_the_outcome_aggregate_is_empty( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """A declared outcome feature with nothing recorded is not a label. + + Dropping the length check stamps ``success: "false"`` here, because + ``any([])`` is False. That is the same invention the hardcoded ``"true"`` + was, one value over, so the empty aggregate needs its own case rather than + riding on the no-feature one. + """ + from hflow.episode import Episode + + output_dir = _import_success_label_corpus(tmp_path, monkeypatch, "empty-aggregate") + landing = sorted((output_dir / "landing").glob("*.mcap")) + with Episode(landing[0]) as episode: + record = episode.metadata_records["episode/v1"] + assert "success" not in record + assert "success_derivation" not in record + + +def test_success_label_omitted_when_source_has_no_outcome_feature( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """A corpus without the outcome feature is normal, not malformed: the key + is omitted (never substituted), the import succeeds, and the catalog + promotion renders the omitted key as SQL NULL (catalog.py:886).""" + import duckdb + + from hflow.catalog import Catalog + from hflow.episode import Episode + from hflow.transform import stamps_from_provenance + + output_dir = _import_success_label_corpus(tmp_path, monkeypatch, "none") + landing = sorted((output_dir / "landing").glob("*.mcap")) + with Episode(landing[0]) as episode: + record = episode.metadata_records["episode/v1"] + assert "success" not in record + assert "success_derivation" not in record + + catalog_root = tmp_path / "catalog" + catalog = Catalog(catalog_root) + catalog.append_episode( + canonical_path=landing[0], + stamps=stamps_from_provenance(episode.metadata), + episode_metadata=dict(episode.metadata), + check_rows=[], + ) + + rows = duckdb.sql( + f"SELECT success FROM read_parquet('{catalog_root / 'episodes' / '*.parquet'}')" + ).fetchall() + assert len(rows) == 1 + assert rows[0][0] is None + + +def test_converter_version_bumped_with_the_label_support() -> None: + """The label changes episode/v1 bytes, which content_episode_id hashes: + the converter version moves with the change, not after it.""" + assert prep.CONVERTER_VERSION == "lerobot-converter-v7"