migration: honor matching SO joint names, convert leading SO block

When the leading joint names match the canonical SO set, treat the dataset as a
genuine SO arm and convert those joints to degrees even if extra columns follow
(bbox, etc.), passing the trailing non-joint columns through untouched. Encoding
detection now looks only at the SO slice so appended columns can't skew it. Only
when no leading 6-DOF SO block can be substantiated is the dataset relabeled
'unknown' and migrated structurally.

community_v3_migration/classify.py

@@ -41,20 +41,32 @@ def is_end_effector(info: dict) -> bool:
     return False
 
 
+def _leading_so_joints(names: list[str], dim: int) -> int:
+    """Number of LEADING joints (a multiple of 6) that match the SO joint order in blocks of 6.
+    When names are absent, fall back to the full dim if it's already a multiple of 6, else 0."""
+    if not names:
+        return dim if dim and dim % 6 == 0 else 0
+    k = 0
+    while (k + 1) * 6 <= len(names) and all(SO_JOINTS[i] in names[k * 6 + i] for i in range(6)):
+        k += 1
+    return k * 6
+
+
+def so_joint_count(info: dict, key: str) -> int:
+    """Leading SO-arm joint count for one feature (``action`` / ``observation.state``). Trailing
+    non-SO columns (bbox, appended EE pose, ...) are excluded so only the genuine SO joints are
+    ever degrees-converted."""
+    feat = info.get("features", {}).get(key, {})
+    dim = (feat.get("shape") or [0])[0]
+    names = [str(n).lower() for n in (feat.get("names") or [])]
+    return _leading_so_joints(names, dim)
+
+
 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
+    """True when ``robot_type`` claims SO but no leading 6-DOF SO joint block can be substantiated
+    from action/observation.state (wrong dim, or names that don't match the SO set). When the first
+    6 joint names DO match, the SO block is honored (and processed) even if extra columns follow."""
+    return max(so_joint_count(info, "action"), so_joint_count(info, "observation.state")) == 0
 
 
 def load_info(root: Path) -> dict:
@@ -138,4 +150,6 @@ def classify(root) -> dict:
         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)}
+    n = so_joint_count(info, key_used) or len(hi)  # ignore trailing non-joint columns
+    return {**out, "is_so": True, "stats_key": key_used, "so_dim": n,
+            **encoding_from_bounds(lo[:n], hi[:n], rt)}

community_v3_migration/fix_dataset.py

@@ -8,7 +8,7 @@ import numpy as np
 import pandas as pd
 
 import so_arm_frame
-from classify import classify
+from classify import classify, load_info, so_joint_count
 
 VALUE_COLS = ("observation.state", "action")
 
@@ -24,14 +24,16 @@ def _set_robot_type(root: Path, robot_type: str) -> None:
     info_path.write_text(json.dumps(info, indent=4))
 
 
-def _rewrite_parquet(root: Path, encoding: str) -> None:
+def _rewrite_parquet(root: Path, encoding: str, so_dims: dict) -> 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 = so_arm_frame.to_degrees(_stack(df[col].values), encoding, n_joints_per_arm=6)
-                df[col] = list(conv.astype(np.float32))
+            n = so_dims.get(col, 0)
+            if col in df.columns and n:
+                full = _stack(df[col].values)  # (N, D)
+                full[:, :n] = so_arm_frame.to_degrees(full[:, :n], encoding, n_joints_per_arm=6)
+                df[col] = list(full.astype(np.float32))
                 changed = True
         if changed:
             df.to_parquet(pq, index=False)
@@ -96,7 +98,12 @@ def fix_dataset_in_place(root) -> dict:
     # drop stray files that would otherwise be uploaded
     for junk in (root / "meta").glob("info.json.bak"):
         junk.unlink()
-    _rewrite_parquet(root, enc)
+    info = load_info(root)
+    so_dims = {c: so_joint_count(info, c) for c in VALUE_COLS}
+    _rewrite_parquet(root, enc, so_dims)
     _regen_episode_stats(root)
+    full_dims = {c: (info.get("features", {}).get(c, {}).get("shape") or [0])[0] for c in VALUE_COLS}
+    partial = any(0 < so_dims[c] < full_dims[c] for c in VALUE_COLS)
+    tail = " (leading SO joints only; trailing non-joint columns left unchanged)" if partial else ""
     return {**cls, "converted": True,
-            "action": f"structural v2.1->v3.0 + joint values converted ({enc} -> degrees)"}
+            "action": f"structural v2.1->v3.0 + joint values converted ({enc} -> degrees){tail}"}