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

community_v3_migration/classify.py

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

community_v3_migration/fix_dataset.py

@@ -0,0 +1,80 @@
+"""Rewrite observation.state / action to degrees in a LOCAL v2.1 SO-arm dataset, then
+regenerate meta/episodes_stats.jsonl (action & state only; other features preserved).
+Run this BEFORE the stock v2.1->v3.0 converter so its stats aggregation stays correct.
+"""
+import json
+from pathlib import Path
+import numpy as np
+import pandas as pd
+
+from classify import classify
+from so_arm_frame import to_degrees
+
+VALUE_COLS = ("observation.state", "action")
+
+
+def _stack(col_values) -> np.ndarray:
+    return np.stack([np.asarray(v, dtype=np.float64) for v in col_values])  # (N, D)
+
+
+def _rewrite_parquet(root: Path, encoding: str) -> None:
+    for pq in sorted((root / "data").glob("*/*.parquet")):
+        df = pd.read_parquet(pq)
+        changed = False
+        for col in VALUE_COLS:
+            if col in df.columns:
+                conv = to_degrees(_stack(df[col].values), encoding, n_joints_per_arm=6)
+                df[col] = list(conv.astype(np.float32))
+                changed = True
+        if changed:
+            df.to_parquet(pq, index=False)
+
+
+def _regen_episode_stats(root: Path) -> None:
+    stats_path = root / "meta" / "episodes_stats.jsonl"
+    orig = {}
+    with open(stats_path) as f:
+        for line in f:
+            e = json.loads(line)
+            orig[e["episode_index"]] = e
+    for pq in sorted((root / "data").glob("*/*.parquet")):
+        df = pd.read_parquet(pq)
+        for ep in np.unique(df["episode_index"].values):
+            ep = int(ep)
+            sub = df[df["episode_index"] == ep]
+            entry = orig.get(ep)
+            if entry is None:
+                continue
+            for col in VALUE_COLS:
+                if col in sub.columns:
+                    a = _stack(sub[col].values)  # (n, D)
+                    entry["stats"][col] = {
+                        "min": a.min(0).tolist(), "max": a.max(0).tolist(),
+                        "mean": a.mean(0).tolist(), "std": a.std(0).tolist(),
+                        "count": [int(a.shape[0])],
+                    }
+    with open(stats_path, "w") as f:
+        for ep in sorted(orig):
+            f.write(json.dumps(orig[ep]) + "\n")
+
+
+def fix_dataset_in_place(root) -> dict:
+    """Returns the classification dict augmented with the action taken."""
+    root = Path(root)
+    cls = classify(root)
+    enc = cls.get("encoding")
+    if not cls.get("is_so") or enc in ("radians", "unknown", "non_so"):
+        reason = {
+            "non_so": "not an SO-100/101 dataset",
+            "radians": "SO-arm joints already in radians",
+            "unknown": "SO-arm but joint encoding could not be determined",
+        }.get(enc, "no joint conversion applicable")
+        return {**cls, "converted": False,
+                "action": f"structural v2.1->v3.0 only ({reason}); joint values left unchanged"}
+    # drop stray files that would otherwise be uploaded
+    for junk in (root / "meta").glob("info.json.bak"):
+        junk.unlink()
+    _rewrite_parquet(root, enc)
+    _regen_episode_stats(root)
+    return {**cls, "converted": True,
+            "action": f"structural v2.1->v3.0 + joint values converted ({enc} -> degrees)"}

community_v3_migration/run_migration.py

@@ -0,0 +1,250 @@
+"""End-to-end migration of the community_dataset_v3 monorepo to v3.0 + SO-arm degrees.
+
+For each `{user}/{dataset}` sub-dataset: stream-download it, fix SO-arm joint values
+(if applicable), run the stock v2.1->v3.0 structural converter locally, upload the v3.0
+result under the same path into a NEW repo, then delete the local copy. Resumable.
+
+  uv run python run_migration.py --dst-repo HuggingFaceVLA/community_dataset_v3_degrees \
+      --work-dir /big/disk/cdv3_work --manifest manifest.csv
+Flags: --only-classify (just write manifest), --no-push (fix+convert locally, keep output,
+       no upload), --folder-name A [B ...] (target specific dataset folders), --limit N,
+       --allow-uncalibrated (accept placeholder CANON ranges).
+"""
+import argparse, csv, json, shutil, sys, traceback
+from pathlib import Path
+from huggingface_hub import HfApi, snapshot_download
+
+import so_arm_frame
+from classify import classify, load_info
+from fix_dataset import fix_dataset_in_place
+
+SRC_REPO = "HuggingFaceVLA/community_dataset_v3"
+
+
+def list_datasets(api: HfApi, repo: str) -> list[str]:
+    files = api.list_repo_files(repo, repo_type="dataset")
+    roots = {p[: -len("/meta/info.json")] for p in files if p.endswith("/meta/info.json")}
+    return sorted(roots)
+
+
+def resolve_folders(api: HfApi, repo: str, names: list[str]) -> list[str]:
+    """Expand each --folder-name into concrete dataset roots (folders that contain
+    meta/info.json). A name may be a full dataset path (returned as-is) or a namespace/prefix
+    like 'Beegbrain' (expanded to every dataset beneath it). Unknown names pass through so
+    they surface as a clear per-item error instead of a confusing FileNotFoundError."""
+    out: list[str] = []
+    for name in names:
+        name = name.strip("/")
+        try:
+            paths = [e.path for e in api.list_repo_tree(
+                repo, path_in_repo=name, recursive=True, repo_type="dataset")]
+        except Exception:
+            out.append(name)  # let it fail loudly downstream
+            continue
+        roots = sorted({p[: -len("/meta/info.json")] for p in paths if p.endswith("/meta/info.json")})
+        out.extend(roots or [name])
+    seen: set[str] = set()
+    return [r for r in out if not (r in seen or seen.add(r))]
+
+
+def already_done(api: HfApi, dst: str, sub: str, dst_files: set[str]) -> bool:
+    return f"{sub}/meta/info.json" in dst_files  # present in target => skip (resume)
+
+
+def _write_dataset_card(local: Path, sub: str, result: dict) -> None:
+    """Regenerate the sub-dataset's card the way LeRobot does (create_lerobot_dataset_card
+    from meta/info.json), then append a migration section documenting provenance and the
+    joint-encoding fix."""
+    enc = result.get("encoding")
+    converted_degrees = bool(result.get("converted"))
+    approx = enc == "normalized" and not so_arm_frame.CANON_IS_CALIBRATED
+
+    enc_labels = {
+        "degrees_old": "legacy degrees (old community frame, pre-#777 convention)",
+        "degrees_new": "degrees (recorded with `use_degrees=True`)",
+        "normalized": "normalized units (joints -100..100, gripper 0..100)",
+        "radians": "radians",
+        "unknown": "undetermined",
+    }
+    joint_actions = {
+        "degrees_old": "per-joint offsets and axis directions corrected to the post-#777 frame (values stay in degrees)",
+        "degrees_new": "already in the post-#777 degrees frame; values unchanged",
+        "normalized": "un-normalized to physical degrees using canonical joint ranges",
+        "radians": "left unchanged (already in radians)",
+        "unknown": "left unchanged (encoding could not be determined)",
+    }
+
+    lines = [
+        "## Migration to LeRobotDataset v3.0",
+        "",
+        "Migrated to LeRobotDataset **v3.0**"
+        + (" with SO-100/101 joint state/action mapped to the post-#777 physical frame (in degrees)."
+           if converted_degrees else "."),
+        "",
+        f"- Source: [`{SRC_REPO}`](https://huggingface.co/datasets/{SRC_REPO}/tree/main/{sub}) (`{sub}`)",
+        "- Codebase version: v2.1 -> v3.0",
+    ]
+    if result.get("is_so"):
+        lines += [
+            f"- Original joint encoding: {enc_labels.get(enc, enc)}",
+            f"- Joint values: {joint_actions.get(enc, 'left unchanged')}",
+            f"- Robot type: `{result.get('robot_type')}`",
+            f"- Action dimension: {result.get('action_dim')}",
+        ]
+    else:
+        lines += ["- Joint values: not applicable (not an SO-100/101 dataset)"]
+    if approx:
+        lines += ["", "> **Note:** normalized (-100..100 / 0..100) joint values were un-normalized "
+                  "using *placeholder* canonical joint ranges (uncalibrated). These values are APPROXIMATE."]
+    if result.get("ambiguous"):
+        lines += ["", "> **Note:** joint-encoding detection was flagged ambiguous; conversion used the "
+                  "best-guess encoding above and may warrant manual review."]
+
+    section = "\n".join(lines) + "\n"
+
+    readme = local / "README.md"
+    try:
+        try:
+            from lerobot.datasets.utils import create_lerobot_dataset_card
+        except ImportError:
+            from lerobot.common.datasets.utils import create_lerobot_dataset_card
+
+        class _Info(dict):            # satisfies both the dict and .to_dict() card variants
+            def to_dict(self):
+                return dict(self)
+
+        rt = result.get("robot_type") or None
+        card = create_lerobot_dataset_card(
+            tags=[rt] if rt else None,
+            dataset_info=_Info(load_info(local)),
+            license="apache-2.0",
+            repo_id=sub,
+        )
+        card.text = card.text.rstrip() + "\n\n" + section
+        card.save(str(readme))
+    except Exception:
+        # LeRobot card generator unavailable at runtime: keep the standalone migration note.
+        if readme.exists():
+            readme.write_text(readme.read_text().rstrip() + "\n\n" + section)
+        else:
+            readme.write_text(f"# {sub}\n\n" + section)
+
+
+def migrate_one(api, dst_repo, sub, work_dir, no_upload) -> dict:
+    local = Path(work_dir) / sub
+    if local.parent.exists():
+        shutil.rmtree(local.parent, ignore_errors=True)  # clean any partial
+    snapshot_download(SRC_REPO, repo_type="dataset", revision="main",
+                      allow_patterns=[f"{sub}/*"], local_dir=work_dir)
+
+    info = load_info(local)
+    if info.get("codebase_version") != "v2.1":
+        return {"root": sub, "action": f"skipped: source codebase is {info.get('codebase_version')} (expected v2.1)"}
+
+    result = fix_dataset_in_place(local)          # SO-arm value fix (or structural_only)
+
+    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
+
+    _write_dataset_card(local, sub, result)       # document the conversion in the dataset card
+
+    base = {k: result.get(k) for k in
+            ("robot_type", "is_so", "encoding", "action_dim", "maxabs", "ambiguous", "action")}
+    if no_upload:
+        # keep the converted output on disk for inspection; do NOT delete or push
+        base["action"] = f"{base['action']}; not pushed (kept locally at {local})"
+        return {"root": sub, **base}
+    api.upload_folder(repo_id=dst_repo, repo_type="dataset", folder_path=str(local),
+                      path_in_repo=sub, commit_message=f"Add {sub} (v3.0, {result['action']})")
+    shutil.rmtree(Path(work_dir) / sub.split("/")[0], ignore_errors=True)  # drop after successful push
+    return {"root": sub, **base}
+
+
+def main():
+    ap = argparse.ArgumentParser(
+        description="Migrate the HuggingFaceVLA/community_dataset_v3 monorepo to LeRobotDataset "
+                    "v3.0, converting SO-100/101 joint state/action to physical degrees along "
+                    "the way. Processes one sub-dataset at a time (download -> fix -> convert -> "
+                    "upload -> cleanup) and is resumable.",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    ap.add_argument("--dst-repo", default=None, metavar="ORG/NAME",
+                    help="Destination HF dataset repo to push the converted v3.0 datasets to "
+                         "(created if missing). Required unless --no-push or --only-classify.")
+    ap.add_argument("--work-dir", default="./cdv3_work", metavar="DIR",
+                    help="Local scratch directory used to download, convert, and (unless pushing) "
+                         "retain each sub-dataset. Only one dataset lives here at a time on a push run.")
+    ap.add_argument("--manifest", default="manifest.csv", metavar="CSV",
+                    help="CSV log appended to as datasets are processed (robot_type, detected "
+                         "encoding, action taken, errors). Reused across resumed runs.")
+    ap.add_argument("--limit", type=int, default=None, metavar="N",
+                    help="Process only the first N sub-datasets (alphabetical). Ignored when "
+                         "--folder-name is given. Useful for a quick end-to-end smoke test.")
+    ap.add_argument("--folder-name", nargs="+", default=None, metavar="USER/DATASET",
+                    help=f"One or more folders WITHIN the {SRC_REPO} monorepo to process. Either a "
+                         "full dataset path ('Beegbrain/draw_pixel_art') or a whole namespace "
+                         "('Beegbrain'), which expands to every dataset under it. Skips the full "
+                         "791-dataset listing.")
+    ap.add_argument("--only-classify", action="store_true",
+                    help="Detect each dataset's robot type and joint encoding and write the "
+                         "manifest, without downloading data, converting, or pushing. Run this "
+                         "first to review scope (especially rows flagged ambiguous=True).")
+    ap.add_argument("--no-push", action="store_true",
+                    help="Fix + convert locally but do NOT upload; the converted v3.0 output is "
+                         "kept under --work-dir for inspection instead of being deleted.")
+    ap.add_argument("--allow-uncalibrated", action="store_true",
+                    help="Permit converting 'normalized' (-100..100) datasets using the PLACEHOLDER "
+                         "canonical joint ranges in so_arm_frame.py. Omit this to force running "
+                         "calibrate_canonical_ranges.py first (recommended for a real run).")
+    args = ap.parse_args()
+    no_upload = args.no_push
+    if not no_upload and not args.only_classify and not args.dst_repo:
+        ap.error("--dst-repo is required unless --no-push or --only-classify is set.")
+
+    if args.allow_uncalibrated:
+        so_arm_frame.CANON_IS_CALIBRATED = True
+    api = HfApi()
+    if args.folder_name:
+        subs = resolve_folders(api, SRC_REPO, args.folder_name)
+        print(f"targeting {len(subs)} sub-dataset(s): {', '.join(subs)}", file=sys.stderr)
+    else:
+        subs = list_datasets(api, SRC_REPO)
+        if args.limit:
+            subs = subs[: args.limit]
+        print(f"{len(subs)} sub-datasets found", file=sys.stderr)
+
+    if not args.only_classify and not no_upload:
+        api.create_repo(args.dst_repo, repo_type="dataset", exist_ok=True)
+    dst_files = set() if (args.only_classify or no_upload) else set(
+        api.list_repo_files(args.dst_repo, repo_type="dataset"))
+
+    first = not Path(args.manifest).exists()
+    with open(args.manifest, "a", newline="") as mf:
+        w = None
+        for i, sub in enumerate(subs):
+            try:
+                if args.only_classify:
+                    # classify without full download: fetch just the meta/ of this sub
+                    snapshot_download(SRC_REPO, repo_type="dataset",
+                                      allow_patterns=[f"{sub}/meta/*"], local_dir=args.work_dir)
+                    row = {"root": sub, **classify(Path(args.work_dir) / sub)}
+                    shutil.rmtree(Path(args.work_dir) / sub.split("/")[0], ignore_errors=True)
+                elif not no_upload and already_done(api, args.dst_repo, sub, dst_files):
+                    row = {"root": sub, "action": "skipped: already present in destination repo"}
+                else:
+                    row = migrate_one(api, args.dst_repo, sub, args.work_dir, no_upload)
+            except Exception as e:
+                row = {"root": sub, "action": f"ERROR: {e}"}
+                traceback.print_exc()
+            if w is None:
+                w = csv.DictWriter(mf, fieldnames=sorted(
+                    {"root", "robot_type", "is_so", "encoding", "action_dim",
+                     "maxabs", "ambiguous", "action", "codebase_version", "note"}))
+                if first:
+                    w.writeheader()
+            w.writerow({k: row.get(k) for k in w.fieldnames})
+            mf.flush()
+            print(f"[{i+1}/{len(subs)}] {sub}: {row.get('action')}", file=sys.stderr)
+
+
+if __name__ == "__main__":
+    main()

community_v3_migration/so_arm_frame.py

@@ -0,0 +1,68 @@
+"""SO-100/101 joint-frame conversion to physical degrees (post-#777 convention).
+
+Two calibration-free branches + one that needs an assumed canonical range:
+
+  * degrees_old  (bare robot_type `so100`/`so101`, |vals|>~180): PR #3879 old->new
+                 convention (sign flip shoulder_lift, +90 deg shoulder_lift/elbow_flex).
+                 EXACT.
+  * degrees_new  (`*_follower` recorded with use_degrees=True, not saturated): already
+                 degrees. EXACT.
+  * normalized   (`*_follower`, -100..100 joints / 0..100 gripper, saturates at bounds):
+                 already mid-range-zero; only the SCALE is missing (per-robot range_min/max
+                 is not stored) -> use assumed canonical per-joint spans below. APPROXIMATE.
+  * radians      -> untouched.
+
+Joint order per arm: shoulder_pan, shoulder_lift, elbow_flex, wrist_flex, wrist_roll, gripper.
+Bimanual (12-dim) tiles the 6-joint block twice.
+"""
+import numpy as np
+
+JOINT_ORDER = ["shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex", "wrist_roll", "gripper"]
+
+# --- PR #3879 (degrees). old(community frame) <-> new(v3.0 / post-#777) frame. ---
+SIGNS = np.array([1.0, -1.0, 1.0, 1.0, 1.0, 1.0], dtype=np.float64)
+OFFSETS_DEG = np.array([0.0, 90.0, 90.0, 0.0, 0.0, 0.0], dtype=np.float64)
+
+# --- Canonical per-joint spans (DEGREES) used ONLY to invert the -100..100 / 0..100
+#     normalization when per-robot calibration is unavailable. joints 0..4 (RANGE_M100_100):
+#     normalized +/-100 -> +/-HALF_RANGE. gripper (RANGE_0_100): 0..100 -> centered at 50,
+#     span FULL_RANGE. THESE ARE PLACEHOLDERS — run calibrate_canonical_ranges.py and paste
+#     the fitted values here before a production run. ---
+CANON_HALF_RANGE_DEG = np.array([100.0, 100.0, 100.0, 100.0, 100.0], dtype=np.float64)  # 5 arm joints
+CANON_GRIPPER_FULL_RANGE_DEG = 90.0
+CANON_IS_CALIBRATED = False  # flipped to True once you paste fitted values
+
+
+def _convert_arm(x: np.ndarray, encoding: str) -> np.ndarray:
+    """x: (..., 6) for a single SO arm -> degrees (..., 6)."""
+    x = np.asarray(x, dtype=np.float64)
+    if encoding == "radians":
+        return x
+    if encoding == "degrees_old":
+        return SIGNS * (x - OFFSETS_DEG)
+    if encoding == "degrees_new":
+        return x
+    if encoding == "normalized":
+        new_deg = np.empty_like(x)
+        new_deg[..., :5] = (x[..., :5] / 100.0) * CANON_HALF_RANGE_DEG
+        new_deg[..., 5] = (x[..., 5] / 100.0 - 0.5) * CANON_GRIPPER_FULL_RANGE_DEG
+        return new_deg
+    raise ValueError(f"unknown encoding: {encoding!r}")
+
+
+def to_degrees(arr, encoding: str, n_joints_per_arm: int = 6) -> np.ndarray:
+    """arr: (..., D) with D a multiple of 6. Returns float32 degrees, same shape."""
+    arr = np.asarray(arr, dtype=np.float64)
+    d = arr.shape[-1]
+    if d % n_joints_per_arm != 0:
+        raise ValueError(f"action/state dim {d} is not a multiple of {n_joints_per_arm}")
+    if encoding == "normalized" and not CANON_IS_CALIBRATED:
+        raise RuntimeError(
+            "CANON ranges are placeholders. Run calibrate_canonical_ranges.py and set "
+            "CANON_* + CANON_IS_CALIBRATED=True, or pass --allow-uncalibrated to accept them."
+        )
+    out = np.empty_like(arr)
+    for a in range(d // n_joints_per_arm):
+        sl = slice(a * n_joints_per_arm, (a + 1) * n_joints_per_arm)
+        out[..., sl] = _convert_arm(arr[..., sl], encoding)
+    return out.astype(np.float32)