diff --git a/community_v3_migration/classify.py b/community_v3_migration/classify.py new file mode 100644 index 000000000..4769b2eee --- /dev/null +++ b/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} diff --git a/community_v3_migration/fix_dataset.py b/community_v3_migration/fix_dataset.py new file mode 100644 index 000000000..b93007d4c --- /dev/null +++ b/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)"} diff --git a/community_v3_migration/run_migration.py b/community_v3_migration/run_migration.py new file mode 100644 index 000000000..57a65b3ca --- /dev/null +++ b/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() diff --git a/community_v3_migration/so_arm_frame.py b/community_v3_migration/so_arm_frame.py new file mode 100644 index 000000000..dd0a6e344 --- /dev/null +++ b/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)