lerobot/community_v3_migration/classify.py

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

# so100/so101 (+ _follower/_bimanual), so_follower, and bimanual bi_so* (bi_so_follower,
# bi_so100_follower, ...; 12-dim).
SO_PREFIXES = ("so100", "so101", "so_", "bi_so")
SO_EXACT: set[str] = set()
# 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 is_so_robot_type(rt: str) -> bool:
    """True if the recorded ``robot_type`` denotes an in-scope SO-100/101 arm."""
    return bool(rt) and (rt.startswith(SO_PREFIXES) or rt in SO_EXACT) and rt not in NEVER_FIX


SO_JOINTS = ("shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex", "wrist_roll", "gripper")


def is_end_effector(info: dict) -> bool:
    """True if action/observation.state are task-space end-effector features (e.g. ``ee_x``,
    ``ee_roll``) rather than joint angles. Such datasets are out of scope for the joint fix."""
    feats = info.get("features", {})
    for key in ("action", "observation.state"):
        names = [str(n).lower() for n in (feats.get(key, {}).get("names") or [])]
        if any(n.startswith("ee_") or "end_effector" in n or "eef" in n for n in names):
            return True
        if {"x", "y", "z"} <= set(names):
            return True
    return False


def is_mislabeled_so(info: dict) -> bool:
    """True when ``robot_type`` claims SO but the features prove it isn't a standard 6-DOF SO arm:
    the joint dim isn't a multiple of 6, or (when names are present) the first 6 joints don't match
    the canonical SO set. Such datasets keep their joints untouched and are relabeled 'unknown'
    rather than being degrees-converted on a wrong assumption."""
    feats = info.get("features", {})
    dims = [feats[c]["shape"][0] for c in ("action", "observation.state") if feats.get(c, {}).get("shape")]
    if any(d % 6 != 0 for d in dims):
        return True
    for key in ("action", "observation.state"):
        names = [str(n).lower() for n in (feats.get(key, {}).get("names") or [])]
        if len(names) >= 6 and not all(SO_JOINTS[i] in names[i] for i in range(6)):
            return True
    return False


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 encoding_from_bounds(lo, hi, rt: str) -> dict:
    """Detect the SO-arm joint encoding from per-joint global min/max and the robot_type name.

    Layout-agnostic (v2.1 episodes_stats or v3.0 stats.json both reduce to lo/hi here), so it is
    the single source of truth for the degrees_old / degrees_new / normalized / radians decision.
    """
    lo = np.asarray(lo, dtype=float)
    hi = np.asarray(hi, dtype=float)
    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 {"encoding": enc, "maxabs": round(maxabs, 2), "saturates": sat, "ambiguous": ambiguous}


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}

    if is_end_effector(info):
        return {**out, "is_so": False, "encoding": "end_effector",
                "note": "task-space end-effector features"}

    if is_so_robot_type(rt) and is_mislabeled_so(info):
        return {**out, "is_so": False, "encoding": "non_so", "mislabeled_so": True,
                "note": "robot_type claims SO but joint dim/names don't match a 6-DOF SO arm"}

    is_so = is_so_robot_type(rt)
    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"}

    return {**out, "is_so": True, "stats_key": key_used, **encoding_from_bounds(lo, hi, rt)}