diff --git a/community_v3_migration/classify.py b/community_v3_migration/classify.py index dfd76e730..0552f59d8 100644 --- a/community_v3_migration/classify.py +++ b/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)} diff --git a/community_v3_migration/fix_dataset.py b/community_v3_migration/fix_dataset.py index 4bbb0569a..9c18d003b 100644 --- a/community_v3_migration/fix_dataset.py +++ b/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}"}