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

@@ -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
 
 

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}"