Add community_dataset_v3 -> v3.0 SO-arm migration scripts

2026-07-23 09:46:00 +00:00 · 2026-07-16 14:44:05 +02:00
parent 279c6c7af3
commit 323febcede
4 changed files with 482 additions and 0 deletions
@@ -0,0 +1,84 @@
+"""Classify each sub-dataset: is it SO-100/101, and what joint encoding is it in?
+
+Detection = robot_type string (recording-time signal) cross-checked against the
+per-episode stats min/max (magnitude + exact-boundary saturation). Mismatches are
+flagged as `ambiguous` for manual review rather than silently converted.
+"""
+import json
+from pathlib import Path
+import numpy as np
+
+SO_PREFIXES = ("so100", "so101")   # so100, so101, so100_follower, so101_follower, so100_bimanual, ...
+SO_EXACT = {"so_follower"}
+# Robots that superficially look SO-like but are NOT in scope for the joint fix:
+NEVER_FIX = {"koch", "koch_follower", "koch_bimanual", "moss", "moss_follower"}
+
+RAD_MAX = 3.5     # |val| below this => radians
+DEG_MIN = 105.0   # |val| above this => old-convention degrees
+SAT_ATOL = 0.5    # closeness to +/-100 / 0 / 100 counted as normalization saturation
+
+
+def load_info(root: Path) -> dict:
+    return json.loads((Path(root) / "meta" / "info.json").read_text())
+
+
+def _global_bounds(root: Path):
+    """Per-joint global min/max over action (fallback observation.state), across episodes."""
+    lo = hi = None
+    key_used = None
+    with open(Path(root) / "meta" / "episodes_stats.jsonl") as f:
+        for line in f:
+            s = json.loads(line)["stats"]
+            key = "action" if "action" in s else ("observation.state" if "observation.state" in s else None)
+            if key is None:
+                continue
+            key_used = key
+            mn = np.asarray(s[key]["min"], dtype=float)
+            mx = np.asarray(s[key]["max"], dtype=float)
+            lo = mn if lo is None else np.minimum(lo, mn)
+            hi = mx if hi is None else np.maximum(hi, mx)
+    return lo, hi, key_used
+
+
+def classify(root) -> dict:
+    root = Path(root)
+    info = load_info(root)
+    rt = info.get("robot_type", "") or ""
+    dim = (info.get("features", {}).get("action", {}).get("shape") or [None])[0]
+    out = {"root": str(root), "robot_type": rt, "action_dim": dim,
+           "codebase_version": info.get("codebase_version"), "ambiguous": False}
+
+    is_so = (rt.startswith(SO_PREFIXES) or rt in SO_EXACT) and rt not in NEVER_FIX
+    if not is_so:
+        return {**out, "is_so": False, "encoding": "non_so"}
+
+    lo, hi, key_used = _global_bounds(root)
+    if lo is None:
+        return {**out, "is_so": True, "encoding": "unknown", "ambiguous": True,
+                "note": "no action/state stats found"}
+
+    maxabs = float(np.nanmax(np.abs(np.concatenate([lo, hi]))))
+    # saturation on any arm joint (index != gripper) at +/-100, or gripper at 0/100
+    n = 6
+    sat = False
+    for a in range(len(hi) // n):
+        arm_hi, arm_lo = hi[a * n:a * n + n], lo[a * n:a * n + n]
+        joints_hi, joints_lo = arm_hi[:5], arm_lo[:5]
+        grip_hi, grip_lo = arm_hi[5], arm_lo[5]
+        sat |= bool(np.any(np.isclose(joints_hi, 100, atol=SAT_ATOL)) or
+                    np.any(np.isclose(joints_lo, -100, atol=SAT_ATOL)) or
+                    np.isclose(grip_hi, 100, atol=SAT_ATOL) or np.isclose(grip_lo, 0, atol=SAT_ATOL))
+
+    if maxabs <= RAD_MAX:
+        enc = "radians"
+    elif maxabs > DEG_MIN:
+        enc = "degrees_old"
+    elif sat:
+        enc = "normalized"
+    else:
+        enc = "degrees_new"
+
+    name_says_new = rt.endswith(("_follower", "_bimanual"))
+    ambiguous = (enc == "degrees_old" and name_says_new) or (enc in ("normalized", "degrees_new") and not name_says_new)
+    return {**out, "is_so": True, "encoding": enc, "maxabs": round(maxabs, 2),
+            "saturates": sat, "stats_key": key_used, "ambiguous": ambiguous}