fix(aggregate stats): fix episode statistics aggregation. Frames and episodes index statistics are offseted and tastk index statistics recomputed from the new labels.

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