| |
| """Check that a gr00t_views dataset's `action` column really is in the LeRobot order its |
| meta/modality.json claims. |
| |
| RoboCasa's simulator emits actions as |
| [eef_pos(3), eef_rot(3), gripper(1), base(3), torso(1), base_mode(1)] |
| while a gr00t_views dataset declares |
| [base_motion(4), control_mode(1), eef_pos(3), eef_rot(3), gripper_close(1)] |
| |
| Any builder that copies simulator actions through without reordering produces a dataset whose |
| numbers are right but whose columns are mislabelled. Training then puts arm motion on the |
| base-motion channel and a constant on the gripper channel, so the policy drives the robot away |
| and never closes the gripper. The loss still converges -- nothing errors -- which is why this |
| has to be checked explicitly. |
| |
| The two layouts are told apart by which dimensions are constant, no reference dataset needed: |
| |
| LeRobot dims 0-3 all zero, dim 4 constant, dim 11 two-valued (the gripper) |
| robosuite dims 0-5 continuous, dim 6 two-valued (the gripper), dims 7-10 zero, dim 11 constant |
| |
| python check_action_layout.py --dataset <gr00t_views dataset> [--episodes 20] |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import glob |
| import json |
| import os |
| import sys |
|
|
| import numpy as np |
| import pyarrow.parquet as pq |
|
|
| LEROBOT_NAMES = ["base_x", "base_y", "base_th", "torso", "ctrl_mode", |
| "eef_x", "eef_y", "eef_z", "rot_x", "rot_y", "rot_z", "gripper"] |
| ROBOSUITE_NAMES = ["eef_x", "eef_y", "eef_z", "rot_x", "rot_y", "rot_z", |
| "gripper", "base_x", "base_y", "base_th", "torso", "base_mode"] |
| |
| FIX = [7, 8, 9, 10, 11, 0, 1, 2, 3, 4, 5, 6] |
|
|
|
|
| def load_actions(dataset: str, n_episodes: int) -> np.ndarray: |
| paths = sorted(glob.glob(os.path.join(dataset, "data", "chunk-*", "*.parquet"))) |
| if not paths: |
| sys.exit(f"no parquet under {dataset}/data/chunk-*/") |
| rows = [] |
| for p in paths[:n_episodes]: |
| rows.append(np.stack(pq.read_table(p, columns=["action"]).column("action").to_pylist())) |
| return np.vstack(rows), len(paths), min(n_episodes, len(paths)) |
|
|
|
|
| def classify(a: np.ndarray) -> tuple[str, list[str]]: |
| """Which layout do the constant/continuous dimensions match?""" |
| const = [i for i in range(a.shape[1]) if len(np.unique(np.round(a[:, i], 4))) == 1] |
| zero = [i for i in const if abs(a[:, i]).max() < 1e-6] |
| twoval = [i for i in range(a.shape[1]) if len(np.unique(np.round(a[:, i], 4))) == 2] |
| notes = [f" all-zero dims : {zero}", f" constant dims : {const}", |
| f" two-valued : {twoval} (the gripper is two-valued)"] |
| lerobot = set([0, 1, 2, 3]).issubset(zero) and 4 in const and 11 in twoval |
| robosuite = set([7, 8, 9, 10]).issubset(zero) and 11 in const and 6 in twoval |
| if lerobot and not robosuite: |
| return "lerobot", notes |
| if robosuite and not lerobot: |
| return "robosuite", notes |
| return "unknown", notes |
|
|
|
|
| def main() -> int: |
| ap = argparse.ArgumentParser(description=__doc__, |
| formatter_class=argparse.RawDescriptionHelpFormatter) |
| ap.add_argument("--dataset", required=True, help="gr00t_views dataset directory") |
| ap.add_argument("--episodes", type=int, default=20, help="episodes to sample (default 20)") |
| a = ap.parse_args() |
|
|
| actions, n_total, n_read = load_actions(a.dataset, a.episodes) |
| print(f"dataset : {a.dataset}") |
| print(f"sampled : {n_read} of {n_total} episodes, {actions.shape[0]} frames x {actions.shape[1]} dims\n") |
|
|
| if actions.shape[1] != 12: |
| print(f"FAIL expected 12 action dims, found {actions.shape[1]}") |
| return 2 |
|
|
| verdict, notes = classify(actions) |
| names = LEROBOT_NAMES if verdict == "lerobot" else ROBOSUITE_NAMES |
| for i in range(12): |
| col = actions[:, i] |
| uniq = len(np.unique(np.round(col, 4))) |
| tag = " <- constant" if uniq == 1 else (" <- two-valued" if uniq == 2 else "") |
| label = names[i] if verdict != "unknown" else "?" |
| print(f" dim{i:2d} {label:10s} {col.min():8.3f} .. {col.max():7.3f} " |
| f"|mean| {abs(col).mean():6.3f} uniq {uniq:5d}{tag}") |
| print() |
| print("\n".join(notes)) |
| print() |
|
|
| modality = os.path.join(a.dataset, "meta", "modality.json") |
| if os.path.exists(modality): |
| act = json.load(open(modality)).get("action", {}) |
| print(" modality.json action slices: " |
| + ", ".join(f"{k} [{v['start']}:{v['end']}]" for k, v in act.items())) |
| print() |
|
|
| if verdict == "lerobot": |
| print("PASS action column is in LeRobot order, matching modality.json.") |
| return 0 |
| if verdict == "robosuite": |
| print("FAIL action column is in RoboCasa/robosuite order but modality.json declares") |
| print(" LeRobot order. Training would read arm motion as base motion and a") |
| print(" constant as the gripper.") |
| print(f" Repair: actions = actions[:, {FIX}]") |
| print(" Fix the builder and rebuild; do not patch the parquet in place unless you") |
| print(" also confirm nothing else consumed it.") |
| return 1 |
| print("UNKNOWN neither layout matched. Inspect the dimensions above by hand:") |
| print(" LeRobot expects dims 0-3 zero, dim 4 constant, dim 11 two-valued.") |
| print(" An all-constant or degenerate episode sample can also land here --") |
| print(" re-run with a larger --episodes before concluding.") |
| return 3 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
