mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-23 09:46:00 +00:00
CarolinePascal
85 lines
3.5 KiB
Python
85 lines
3.5 KiB
Python
"""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}
|