From 5cd6b516c0820da922ca30f9fae2bf41bb2f78ea Mon Sep 17 00:00:00 2001 From: CarolinePascal Date: Fri, 31 Jul 2026 16:51:15 +0200 Subject: [PATCH] fix(aggregate stats): fix episode statistics aggregation. Frames and episodes index statistics are offseted and tastk index statistics recomputed from the new labels. --- src/lerobot/datasets/aggregate.py | 41 +++++++++++++++++++ tests/datasets/test_aggregate.py | 65 +++++++++++++++++++++++++++++++ 2 files changed, 106 insertions(+) diff --git a/src/lerobot/datasets/aggregate.py b/src/lerobot/datasets/aggregate.py index 45d8d82e1..f08fc6670 100644 --- a/src/lerobot/datasets/aggregate.py +++ b/src/lerobot/datasets/aggregate.py @@ -22,6 +22,7 @@ from pathlib import Path from typing import Any, NotRequired, TypedDict import datasets +import numpy as np import pandas as pd import tqdm @@ -303,6 +304,46 @@ def update_meta_data( df["dataset_to_index"] = df["dataset_to_index"] + dst_meta.info.total_frames df["episode_index"] = df["episode_index"] + dst_meta.info.total_episodes + # Per-episode stats still describe the pre-merge values of the bookkeeping columns + # reindexed above. index/episode_index shift by a constant; task_index is relabeled, + # so recompute it from the episode's (stable) task strings via the unified tasks table. + shift_stat_keys = ("min", "max", "mean", "q01", "q10", "q50", "q90", "q99") + for name, offset in ( + ("episode_index", dst_meta.info.total_episodes), + ("index", dst_meta.info.total_frames), + ): + for stat in shift_stat_keys: + col = f"stats/{name}/{stat}" + if col in df.columns: + df[col] = df[col] + offset + + if any(c.startswith("stats/task_index/") for c in df.columns): + quantiles = {"q01": 0.01, "q10": 0.10, "q50": 0.50, "q90": 0.90, "q99": 0.99} + ids_per_row = [ + np.array([dst_meta.tasks.loc[t, "task_index"] for t in tasks], dtype=np.float64) + for tasks in df["tasks"] + ] + + def _task_stat(ids, stat): + if stat == "min": + return ids.min() + if stat == "max": + return ids.max() + if stat == "std": + return ids.std() + if stat in quantiles: + return np.quantile(ids, quantiles[stat]) + return ids.mean() + + for stat in ("min", "max", "mean", "std", *quantiles): + col = f"stats/task_index/{stat}" + if col in df.columns: + # np.full_like preserves each cell container and dtype so the parquet schema is unchanged. + df[col] = [ + np.full_like(orig, _task_stat(ids, stat)) + for orig, ids in zip(df[col], ids_per_row, strict=True) + ] + return df diff --git a/tests/datasets/test_aggregate.py b/tests/datasets/test_aggregate.py index 2fafd2777..78c622db7 100644 --- a/tests/datasets/test_aggregate.py +++ b/tests/datasets/test_aggregate.py @@ -23,12 +23,15 @@ import pytest pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])") import datasets # noqa: E402 +import numpy as np +import pandas as pd import torch from lerobot.configs import VIDEO_ENCODER_INFO_KEYS from lerobot.datasets.aggregate import aggregate_datasets from lerobot.datasets.feature_utils import features_equal_for_merge from lerobot.datasets.lerobot_dataset import LeRobotDataset +from lerobot.datasets.utils import EPISODES_DIR from tests.fixtures.constants import ( DUMMY_CAMERA_FEATURES_WITH_DEPTH, DUMMY_REPO_ID, @@ -857,3 +860,65 @@ def test_aggregate_already_merged_dataset(tmp_path, lerobot_dataset_factory): # This would raise FileNotFoundError before the fix assert_dataset_iteration_works(ds_abc) + + +def test_aggregate_updates_per_episode_stats(tmp_path): + """episode_index/index/task_index per-episode stats follow the merge; all other stats are copied verbatim.""" + features = {"observation.state": {"dtype": "float32", "shape": (2,), "names": None}} + + def _make_dataset(suffix, tasks): + ds = LeRobotDataset.create( + f"{DUMMY_REPO_ID}_{suffix}", fps=10, features=features, root=tmp_path / suffix + ) + for task in tasks: + for _ in range(4): + ds.add_frame({"observation.state": torch.randn(2), "task": task}) + ds.save_episode() + ds.finalize() + return ds + + # Overlapping tasks so relabeling collapses shared "b" and introduces new "c". + sources = [_make_dataset("s0", ["a", "b"]), _make_dataset("s1", ["b", "c"])] + aggr_root = tmp_path / "aggr" + aggregate_datasets( + repo_ids=[d.repo_id for d in sources], + roots=[d.root for d in sources], + aggr_repo_id=f"{DUMMY_REPO_ID}_aggr", + aggr_root=aggr_root, + ) + with ( + patch("lerobot.datasets.dataset_metadata.get_safe_version", return_value="v3.0"), + patch("lerobot.datasets.dataset_metadata.snapshot_download", return_value=str(aggr_root)), + ): + aggr = LeRobotDataset(f"{DUMMY_REPO_ID}_aggr", root=aggr_root) + assert aggr.meta.total_tasks == 3 # "b" deduped, "c" added + + def _load_stats(root): + # load_episodes drops stats/* columns, so read the episodes parquet shards directly. + shards = sorted((root / EPISODES_DIR).rglob("*.parquet")) + return pd.concat([pd.read_parquet(s) for s in shards]).set_index("episode_index") + + merged = _load_stats(aggr_root) + src_rows, ep_off, fr_off = [], 0, 0 + for d in sources: + src = _load_stats(d.root) + src_rows += [(src.loc[ep], ep_off, fr_off) for ep in range(d.num_episodes)] + ep_off, fr_off = ep_off + d.num_episodes, fr_off + d.num_frames + + shift = {"min", "max", "mean", "q01", "q10", "q50", "q90", "q99"} + for ep, (src_row, e_off, f_off) in enumerate(src_rows): + row = merged.loc[ep] + new_id = float(aggr.meta.tasks.loc[row["tasks"][0], "task_index"]) + for col in (c for c in merged.columns if c.startswith("stats/")): + feat, key = col[len("stats/") :].rsplit("/", 1) + got = np.asarray(row[col], dtype=np.float64).reshape(-1) + base = np.asarray(src_row[col], dtype=np.float64).reshape(-1) + if feat == "episode_index": + expected = base + e_off if key in shift else base + elif feat == "index": + expected = base + f_off if key in shift else base + elif feat == "task_index": + expected = base if key == "count" else np.full_like(base, 0.0 if key == "std" else new_id) + else: + expected = base + assert np.allclose(got, expected), f"ep{ep} {col}: {got} != {expected}"