From a6dcb185858546ecc67c633da1f4cb8fefff9f25 Mon Sep 17 00:00:00 2001
From: CarolinePascal <caroline8.pascal@gmail.com>
Date: Fri, 17 Jul 2026 16:59:10 +0200
Subject: [PATCH] migration: reconcile stale meta episode count to data+video
 files

When the data files and video files agree on episode count N but the metadata
lists a different count, rewrite meta/episodes.jsonl, meta/episodes_stats.jsonl
and info.json to N before the v2.1->v3.0 converter runs (which otherwise raises
'Number of episodes is not the same'). Only the safe direction is handled:
trimming metadata that lists MORE episodes than exist. Non-contiguous data,
data/video disagreement, or metadata missing episodes are left untouched.
---
 community_v3_migration/fix_dataset.py   | 55 +++++++++++++++++++++++++
 community_v3_migration/run_migration.py |  6 ++-
 2 files changed, 60 insertions(+), 1 deletion(-)

diff --git a/community_v3_migration/fix_dataset.py b/community_v3_migration/fix_dataset.py
index 9c18d003b..af8978f6a 100644
--- a/community_v3_migration/fix_dataset.py
+++ b/community_v3_migration/fix_dataset.py
@@ -67,6 +67,61 @@ def _regen_episode_stats(root: Path) -> None:
             f.write(json.dumps(orig[ep]) + "\n")
 
 
+def _read_jsonl(path: Path) -> list[dict]:
+    with open(path) as f:
+        return [json.loads(line) for line in f if line.strip()]
+
+
+def _write_jsonl(path: Path, rows: list[dict]) -> None:
+    with open(path, "w") as f:
+        for r in rows:
+            f.write(json.dumps(r) + "\n")
+
+
+def reconcile_episode_count(root) -> str | None:
+    """When the data files and video files agree on an episode count N but the metadata lists a
+    different count, rewrite the metadata (episodes.jsonl, episodes_stats.jsonl, info.json) to N.
+
+    Only the safe direction is handled: trimming metadata that lists MORE episodes than actually
+    exist. If the data itself is non-contiguous, the videos disagree with the data, or the metadata
+    lists FEWER episodes than the data (which would require fabricating per-episode stats), nothing
+    is changed and the stock converter's mismatch error is left to surface. Returns a note on fix."""
+    root = Path(root)
+    info = json.loads((root / "meta" / "info.json").read_text())
+
+    data_idx = sorted(int(p.stem.split("_")[-1]) for p in (root / "data").glob("*/episode_*.parquet"))
+    n = len(data_idx)
+    if n == 0 or data_idx != list(range(n)):
+        return None
+
+    vkeys = [k for k, f in info.get("features", {}).items() if f.get("dtype") == "video"]
+    for k in vkeys:
+        if len(list((root / "videos").glob(f"*/{k}/episode_*.mp4"))) != n:
+            return None  # data and videos disagree -> out of scope for this fix
+
+    eps_path = root / "meta" / "episodes.jsonl"
+    stats_path = root / "meta" / "episodes_stats.jsonl"
+    eps, stats = _read_jsonl(eps_path), _read_jsonl(stats_path)
+    if len(eps) == n and len(stats) == n:
+        return None
+
+    eps_keep = [e for e in eps if e.get("episode_index", -1) < n]
+    stats_keep = [s for s in stats if s.get("episode_index", -1) < n]
+    if len(eps_keep) != n or len(stats_keep) != n:
+        return None  # metadata is missing episodes present in the data -> can't safely fabricate
+
+    dropped = max(len(eps), len(stats)) - n
+    _write_jsonl(eps_path, eps_keep)
+    _write_jsonl(stats_path, stats_keep)
+    info["total_episodes"] = n
+    info["total_frames"] = int(sum(e.get("length", 0) for e in eps_keep))
+    if "total_videos" in info:
+        info["total_videos"] = n * len(vkeys)
+    info["splits"] = {"train": f"0:{n}"}
+    (root / "meta" / "info.json").write_text(json.dumps(info, indent=4))
+    return f"metadata episode count reconciled to {n} (data & videos agree; dropped {dropped} stale meta entries)"
+
+
 def fix_dataset_in_place(root) -> dict:
     """Returns the classification dict augmented with the action taken."""
     root = Path(root)
diff --git a/community_v3_migration/run_migration.py b/community_v3_migration/run_migration.py
index 86ee9bcf4..6a308ffb3 100644
--- a/community_v3_migration/run_migration.py
+++ b/community_v3_migration/run_migration.py
@@ -17,7 +17,7 @@ from huggingface_hub import HfApi
 
 import so_arm_frame
 from classify import classify, is_end_effector, load_info
-from fix_dataset import fix_dataset_in_place
+from fix_dataset import fix_dataset_in_place, reconcile_episode_count
 
 SRC_REPO = "HuggingFaceVLA/community_dataset_v3"
 
@@ -171,6 +171,10 @@ def migrate_one(api, dst_repo, sub, work_dir, no_upload) -> dict:
 
     result = fix_dataset_in_place(local)          # SO-arm value fix (or structural_only)
 
+    reconciled = reconcile_episode_count(local)   # align stale meta counts to data+video files
+    if reconciled:
+        result["action"] = f"{result['action']}; {reconciled}"
+
     from lerobot.scripts.convert_dataset_v21_to_v30 import convert_dataset
     convert_dataset(repo_id=sub, root=str(local), push_to_hub=False)  # v2.1 -> v3.0, in place