diff --git a/README.md b/README.md index 04890a250..dfbe8a1cc 100644 --- a/README.md +++ b/README.md @@ -73,6 +73,10 @@ All architectures are selected via `configs/model/.yaml` with `archit ### Benchmarks and Evaluation +RoboTwin EEF uses native **wxyz** quaternions. Checkpoints trained with the earlier +xyzw interpretation require their matching legacy revision; see the +[quaternion convention and migration note](benchmarks/robotwin/README.md#eef-quaternion-convention-and-checkpoint-compatibility). + | Benchmark | Status | Notes | |---|---|---| | RoboTwin eval | Supported | All 50 tasks; see `benchmarks/robotwin/` | diff --git a/benchmarks/robotwin/README.md b/benchmarks/robotwin/README.md index e9d21a84f..8fc937d65 100644 --- a/benchmarks/robotwin/README.md +++ b/benchmarks/robotwin/README.md @@ -35,6 +35,22 @@ git clone https://github.com/RoboTwin-Platform/RoboTwin.git /path/to/RoboTwin | mplib | 0.2.1 | | torch (RoboTwin env) | 2.4.1 (the policy server uses OpenWAM's own env) | +### EEF quaternion convention and checkpoint compatibility + +Native RoboTwin `endpose` observations, HDF5 fields, and `ee` actions use +scalar-first **wxyz** quaternions. The EEF loader, statistics computation, and +benchmark adapter convert these to/from the first two columns of the physical +rotation matrix (rot6d). + +Earlier revisions interpreted these native tuples as xyzw. This correction +changes the EEF representation and is **not a drop-in change for checkpoints +trained with that legacy convention**. Keep those checkpoints with their +matching legacy training/evaluation revision until they are explicitly migrated +or retrained. Regenerate normalization statistics from native HDF5 for training +with the corrected representation; existing final stats files are not migrated +automatically. The stats script uses a new partial-shard namespace so resumed +computation does not reuse legacy shards. Joint/qpos mode is unaffected. + ## 2. Start the Policy Server ```bash diff --git a/benchmarks/robotwin/policy_config.yml b/benchmarks/robotwin/policy_config.yml index 47557a6a4..12ad94fd7 100644 --- a/benchmarks/robotwin/policy_config.yml +++ b/benchmarks/robotwin/policy_config.yml @@ -25,7 +25,7 @@ request_timeout: 300 # WebSocket timeout in seconds # Must match the action_mode used at training time. # # ee — End-effector: the server returns 20D (xyz+rot6d+grip)×2, which is -# automatically converted to 16D (xyz+quat_xyzw+grip)×2 before calling +# automatically converted to 16D (xyz+quat_wxyz+grip)×2 before calling # take_action(action_type='ee'). Use with EEF-mode training (action_mode: eef). Default. # # qpos — 14D joint angles: [l_joints(6), l_grip(1), r_joints(6), r_grip(1)] diff --git a/benchmarks/utils/action_conversion.py b/benchmarks/utils/action_conversion.py index 8c3feda30..6d6794062 100644 --- a/benchmarks/utils/action_conversion.py +++ b/benchmarks/utils/action_conversion.py @@ -5,10 +5,9 @@ [l_xyz(3), l_rot6d(6), l_grip(1), r_xyz(3), r_rot6d(6), r_grip(1)] -Many robot benchmarks (RoboTwin, SimplerEnv, etc.) consume 16-dim end-effector -actions with quaternion rotations:: +RoboTwin consumes 16-dim end-effector actions with wxyz quaternion rotations:: - [l_xyz(3), l_quat_xyzw(4), l_grip(1), r_xyz(3), r_quat_xyzw(4), r_grip(1)] + [l_xyz(3), l_quat_wxyz(4), l_grip(1), r_xyz(3), r_quat_wxyz(4), r_grip(1)] These helpers are pure numpy and have no dependency on the ``openwam`` package, so they can run in any benchmark client's Python environment. @@ -26,7 +25,7 @@ def quat_xyzw_to_rot6d(quat: np.ndarray) -> np.ndarray: - """Convert xyzw quaternion(s) to 6D rotation, matching RoboTwinDataset.""" + """Convert xyzw quaternion(s) to 6D rotation (first two matrix columns).""" q = np.asarray(quat, dtype=np.float64) if q.shape[-1] != 4: raise ValueError(f"quat must end with dimension 4, got shape {q.shape}") @@ -104,12 +103,13 @@ def eef20d_to_ee16d(action: np.ndarray) -> np.ndarray: """Convert 20D EEF action (xyz+rot6d+grip)×2 to 16D ee action (xyz+quat+grip)×2. OpenWAM EEF (20D): ``[l_xyz(3), l_rot6d(6), l_grip(1), r_xyz(3), r_rot6d(6), r_grip(1)]`` - RoboTwin ee (16D): ``[l_xyz(3), l_quat_xyzw(4), l_grip(1), r_xyz(3), r_quat_xyzw(4), r_grip(1)]`` + RoboTwin ee (16D): ``[l_xyz(3), l_quat_wxyz(4), l_grip(1), r_xyz(3), r_quat_wxyz(4), r_grip(1)]`` """ l_xyz, l_r6d, l_grip = action[0:3], action[3:9], action[9:10] r_xyz, r_r6d, r_grip = action[10:13], action[13:19], action[19:20] - l_quat = rot6d_to_quat_xyzw(l_r6d) - r_quat = rot6d_to_quat_xyzw(r_r6d) + # RoboTwin consumes scalar-first wxyz; the shared decoder returns xyzw. + l_quat = rot6d_to_quat_xyzw(l_r6d)[[3, 0, 1, 2]] + r_quat = rot6d_to_quat_xyzw(r_r6d)[[3, 0, 1, 2]] return np.concatenate([l_xyz, l_quat, l_grip, r_xyz, r_quat, r_grip]).astype(np.float32) @@ -278,13 +278,13 @@ def robotwin_endpose_to_eef20d( This mirrors ``RoboTwinDataset._read_eef_actions``: ``[left_xyz, left_rot6d, left_grip, right_xyz, right_rot6d, right_grip]``. - RoboTwin endpose quaternions are xyzw. + RoboTwin endpose quaternions are wxyz. """ left_ep = np.asarray(left_endpose, dtype=np.float32).reshape(-1) right_ep = np.asarray(right_endpose, dtype=np.float32).reshape(-1) if left_ep.shape[0] != 7 or right_ep.shape[0] != 7: raise ValueError( - f"RoboTwin endpose fields must be 7D xyz+quat_xyzw; got left={left_ep.shape}, right={right_ep.shape}" + f"RoboTwin endpose fields must be 7D xyz+quat_wxyz; got left={left_ep.shape}, right={right_ep.shape}" ) left_grip = np.asarray(left_gripper, dtype=np.float32).reshape(-1) @@ -292,8 +292,11 @@ def robotwin_endpose_to_eef20d( if left_grip.size < 1 or right_grip.size < 1: raise ValueError("RoboTwin gripper fields must contain at least one scalar value.") - left = np.concatenate([left_ep[:3], quat_xyzw_to_rot6d(left_ep[3:]), left_grip[:1]], axis=-1) - right = np.concatenate([right_ep[:3], quat_xyzw_to_rot6d(right_ep[3:]), right_grip[:1]], axis=-1) + # Reorder native xyz+wxyz poses for the shared scalar-last encoder. + left_quat_xyzw = left_ep[[4, 5, 6, 3]] + right_quat_xyzw = right_ep[[4, 5, 6, 3]] + left = np.concatenate([left_ep[:3], quat_xyzw_to_rot6d(left_quat_xyzw), left_grip[:1]], axis=-1) + right = np.concatenate([right_ep[:3], quat_xyzw_to_rot6d(right_quat_xyzw), right_grip[:1]], axis=-1) return np.concatenate([left, right], axis=-1).astype(np.float32) diff --git a/openwam/dataloader/robotwin.py b/openwam/dataloader/robotwin.py index 48362ee3e..715521ab4 100644 --- a/openwam/dataloader/robotwin.py +++ b/openwam/dataloader/robotwin.py @@ -34,7 +34,7 @@ Normalizer, load_mode_stats, ) -from openwam.dataloader.transforms.rotation import quat_xyzw_to_rotation_6d +from openwam.dataloader.transforms.rotation import RotationType, convert_rotation from openwam.dataloader.transforms.video import VideoColorJitter, color_jitter_enabled from openwam.dataloader.utils.unify_action import ( UNIFY_DIM, @@ -713,7 +713,7 @@ def _read_eef_actions(self, f, start: int, end: int) -> np.ndarray: right_xyz(3), right_rot6d(6), right_grip(1)] Gripper values are raw continuous values from HDF5 (1=open, 0=closed). """ - left_ep = f["endpose/left_endpose"][start:end] # (T, 7): xyz + quat_xyzw + left_ep = f["endpose/left_endpose"][start:end] # (T, 7): xyz + quat_wxyz right_ep = f["endpose/right_endpose"][start:end] left_grip = f["endpose/left_gripper"][start:end].astype(np.float64) right_grip = f["endpose/right_gripper"][start:end].astype(np.float64) @@ -721,7 +721,7 @@ def _read_eef_actions(self, f, start: int, end: int) -> np.ndarray: left = np.concatenate( [ left_ep[:, :3], - quat_xyzw_to_rotation_6d(left_ep[:, 3:]), + convert_rotation(left_ep[:, 3:], RotationType.QUATERNION, RotationType.ROTATION_6D), left_grip[:, None] if left_grip.ndim == 1 else left_grip, ], axis=-1, @@ -730,7 +730,7 @@ def _read_eef_actions(self, f, start: int, end: int) -> np.ndarray: right = np.concatenate( [ right_ep[:, :3], - quat_xyzw_to_rotation_6d(right_ep[:, 3:]), + convert_rotation(right_ep[:, 3:], RotationType.QUATERNION, RotationType.ROTATION_6D), right_grip[:, None] if right_grip.ndim == 1 else right_grip, ], axis=-1, diff --git a/openwam/dataloader/transforms/rotation.py b/openwam/dataloader/transforms/rotation.py index f3452cf19..7b62257dc 100644 --- a/openwam/dataloader/transforms/rotation.py +++ b/openwam/dataloader/transforms/rotation.py @@ -153,11 +153,10 @@ def convert_rotation( # --------------------------------------------------------------------------- -# Convenience helpers for RoboTwin HDF5 xyzw quaternion ↔ 6D rotation +# Convenience helpers for xyzw quaternion ↔ 6D rotation # -# RoboTwin HDF5 stores quaternions as xyzw (scipy default) while -# OpenWAM's convert_rotation() uses wxyz. These thin wrappers handle -# the convention swap so callers don't have to think about it. +# scipy uses xyzw while convert_rotation() uses wxyz. Native RoboTwin +# endpose fields already use wxyz and should call convert_rotation() directly. # --------------------------------------------------------------------------- @@ -165,7 +164,7 @@ def quat_xyzw_to_rotation_6d(q: np.ndarray) -> np.ndarray: """Convert xyzw quaternion to 6D rotation representation. Args: - q: (..., 4) quaternion in xyzw format (scipy / RoboTwin convention). + q: (..., 4) quaternion in xyzw format (scipy convention). Returns: (..., 6) 6D rotation (first two columns of the rotation matrix). @@ -181,7 +180,7 @@ def rotation_6d_to_quat_xyzw(r6d: np.ndarray) -> np.ndarray: r6d: (..., 6) 6D rotation representation. Returns: - (..., 4) quaternion in xyzw format (scipy / RoboTwin convention). + (..., 4) quaternion in xyzw format (scipy convention). """ wxyz = convert_rotation(r6d, RotationType.ROTATION_6D, RotationType.QUATERNION) return np.concatenate([wxyz[..., 1:], wxyz[..., :1]], axis=-1) diff --git a/openwam/dataloader/utils/stats_computation/robotwin_stats_computation.py b/openwam/dataloader/utils/stats_computation/robotwin_stats_computation.py index 3f5325069..8c848d51a 100644 --- a/openwam/dataloader/utils/stats_computation/robotwin_stats_computation.py +++ b/openwam/dataloader/utils/stats_computation/robotwin_stats_computation.py @@ -43,7 +43,8 @@ _EEF_ACTION_DIM = 20 # [xyz(3) + rot6d(6) + grip(1)] x 2 arms _MODES = ("joint", "eef") -_SHARD_DIR_NAME = "shards_v1" +# Do not reuse EEF shards computed by interpreting native wxyz as xyzw. +_SHARD_DIR_NAME = "shards_wxyz_v2" # --------------------------------------------------------------------------- @@ -107,8 +108,8 @@ def _shard_path_for_root(checkpoint_path: str, data_root: str) -> str: The absolute ``data_root`` is the whole cache key: expanding/shrinking the task list naturally reuses or ignores shards by recomputing this path for - the *current* task_roots. ``shards_v1`` is the shard-format namespace; bump - it if the NPZ payload schema ever changes. + the *current* task_roots. The shard namespace versions both the payload + schema and rotation convention, so legacy EEF statistics cannot be reused. """ canonical = os.path.abspath(data_root) digest = hashlib.sha1(canonical.encode("utf-8")).hexdigest()[:16] @@ -125,7 +126,7 @@ def _read_eef_actions(f) -> Optional[np.ndarray]: if "endpose/left_endpose" not in f: return None # Local import avoids a hard dependency when only joint is needed in tests - from openwam.dataloader.transforms.rotation import quat_xyzw_to_rotation_6d + from openwam.dataloader.transforms.rotation import RotationType, convert_rotation left_ep = f["endpose/left_endpose"][()].astype(np.float64) right_ep = f["endpose/right_endpose"][()].astype(np.float64) @@ -135,7 +136,7 @@ def _read_eef_actions(f) -> Optional[np.ndarray]: left = np.concatenate( [ left_ep[:, :3], - quat_xyzw_to_rotation_6d(left_ep[:, 3:]).astype(np.float64), + convert_rotation(left_ep[:, 3:], RotationType.QUATERNION, RotationType.ROTATION_6D).astype(np.float64), left_grip[:, None] if left_grip.ndim == 1 else left_grip, ], axis=-1, @@ -143,7 +144,7 @@ def _read_eef_actions(f) -> Optional[np.ndarray]: right = np.concatenate( [ right_ep[:, :3], - quat_xyzw_to_rotation_6d(right_ep[:, 3:]).astype(np.float64), + convert_rotation(right_ep[:, 3:], RotationType.QUATERNION, RotationType.ROTATION_6D).astype(np.float64), right_grip[:, None] if right_grip.ndim == 1 else right_grip, ], axis=-1, @@ -373,7 +374,7 @@ def compute_multitask_robotwin_stats( tasks: Optional internal task restriction. Defaults to every task discovered on disk. checkpoint_path: When set, the function persists a per-task NPZ - shard into ``.partial/shards_v1/`` after each + shard into ``.partial/shards_wxyz_v2/`` after each task-root is processed. A subsequent call recomputes the shard path from the current ``data_root`` and skips it when the shard already exists. diff --git a/tests/test_robotwin_dataloader.py b/tests/test_robotwin_dataloader.py index 94582bd2c..ce071a929 100644 --- a/tests/test_robotwin_dataloader.py +++ b/tests/test_robotwin_dataloader.py @@ -28,7 +28,7 @@ def _encode_jpeg(height=16, width=16, seed=0): def _create_mock_episode(path, T=20, action_dim=14, seed=0): """Create a mock RoboTwin HDF5 episode with both joint and endpose keys.""" rng = np.random.default_rng(seed) - quats = Rotation.random(T, random_state=seed).as_quat() # (T, 4) xyzw + quats = Rotation.random(T, random_state=seed).as_quat()[:, [3, 0, 1, 2]] # native wxyz # Pre-encode JPEG frames jpeg_frames = [_encode_jpeg(16, 16, seed=seed + i) for i in range(T)] @@ -861,7 +861,7 @@ def test_multitask_action_stats_can_resume_from_partial_checkpoint(): checkpoint_path=checkpoint_path, ) assert partial_stats["num_timesteps"] > 0 - assert os.path.isdir(f"{checkpoint_path}.partial/shards_v1") + assert os.path.isdir(f"{checkpoint_path}.partial/shards_wxyz_v2") resumed_stats = compute_multitask_robotwin_stats( dataset_dir=tmpdir, @@ -904,7 +904,7 @@ def test_resume_does_not_recompute_already_checkpointed_shards(): checkpoint_path=checkpoint_path, ) - shards_dir = f"{checkpoint_path}.partial/shards_v1" + shards_dir = f"{checkpoint_path}.partial/shards_wxyz_v2" shards_before = {name: os.stat(os.path.join(shards_dir, name)).st_mtime_ns for name in os.listdir(shards_dir)} assert len(shards_before) == 1 diff --git a/tests/test_robotwin_interface_proprio.py b/tests/test_robotwin_interface_proprio.py index a20aa3365..995b11711 100644 --- a/tests/test_robotwin_interface_proprio.py +++ b/tests/test_robotwin_interface_proprio.py @@ -5,7 +5,6 @@ from benchmarks.robotwin import openwam2robotwin_interface as iface from benchmarks.utils import action_conversion -from openwam.dataloader.transforms.rotation import quat_xyzw_to_rotation_6d class _Model: @@ -15,17 +14,17 @@ def __init__(self, action_type: str): def test_robotwin_endpose_conversion_matches_dataset_eef_layout(): - left = np.array([0.1, 0.2, 0.3, 0.0, 0.0, 0.0, 1.0], dtype=np.float32) - right = np.array([0.4, 0.5, 0.6, 0.0, 0.0, 0.0, 1.0], dtype=np.float32) + left = np.array([0.1, 0.2, 0.3, 1.0, 0.0, 0.0, 0.0], dtype=np.float32) + right = np.array([0.4, 0.5, 0.6, 1.0, 0.0, 0.0, 0.0], dtype=np.float32) eef = action_conversion.robotwin_endpose_to_eef20d(left, right, 0.25, 0.75) expected = np.concatenate( [ left[:3], - quat_xyzw_to_rotation_6d(left[3:]), + np.array([1, 0, 0, 0, 1, 0], dtype=np.float32), np.array([0.25], dtype=np.float32), right[:3], - quat_xyzw_to_rotation_6d(right[3:]), + np.array([1, 0, 0, 0, 1, 0], dtype=np.float32), np.array([0.75], dtype=np.float32), ] ).astype(np.float32) @@ -37,8 +36,8 @@ def test_robotwin_endpose_conversion_matches_dataset_eef_layout(): def test_ee_action_type_sends_20d_eef_proprio(): observation = { "endpose": { - "left_endpose": np.array([0.1, 0.2, 0.3, 0.0, 0.0, 0.0, 1.0], dtype=np.float32), - "right_endpose": np.array([0.4, 0.5, 0.6, 0.0, 0.0, 0.0, 1.0], dtype=np.float32), + "left_endpose": np.array([0.1, 0.2, 0.3, 1.0, 0.0, 0.0, 0.0], dtype=np.float32), + "right_endpose": np.array([0.4, 0.5, 0.6, 1.0, 0.0, 0.0, 0.0], dtype=np.float32), "left_gripper": np.array(0.25, dtype=np.float32), "right_gripper": np.array(0.75, dtype=np.float32), }, diff --git a/tests/test_robotwin_quaternion_convention.py b/tests/test_robotwin_quaternion_convention.py new file mode 100644 index 000000000..69d634ba3 --- /dev/null +++ b/tests/test_robotwin_quaternion_convention.py @@ -0,0 +1,92 @@ +"""Native RoboTwin wxyz must describe the same physical rotation in every EEF path.""" + +import hashlib +from pathlib import Path + +import h5py +import numpy as np +import pytest +from scipy.spatial.transform import Rotation + +from benchmarks.utils import action_conversion +from openwam.dataloader.robotwin import RoboTwinDataset +from openwam.dataloader.utils.stats_computation import robotwin_stats_computation as stats + + +@pytest.mark.parametrize( + "quat", + [ + [1, 0, 0, 0], # Identity, not Rx(180). + [2**-0.5, 2**-0.5, 0, 0], + [2**-0.5, 0, 2**-0.5, 0], + [2**-0.5, 0, 0, 2**-0.5], + [0, 1, 0, 0], + [0, 0, 1, 0], + [0, 0, 0, 1], + [0.5, -0.5, 0.5, 0.5], + ], + ids=["identity", "x90", "y90", "z90", "x180", "y180", "z180", "mixed"], +) +@pytest.mark.parametrize("sign", [1, -1]) +def test_native_wxyz_matches_physical_rotation_across_eef_paths(tmp_path, quat, sign): + left_q = sign * np.asarray(quat, dtype=np.float64) + right_q = left_q[[0, 3, 1, 2]] + left = np.r_[[0.1, 0.2, 0.3], left_q] + right = np.r_[[0.4, 0.5, 0.6], right_q] + rotations = Rotation.from_quat(np.stack([left_q, right_q])[:, [1, 2, 3, 0]]).as_matrix() + rot6d = np.concatenate([rotations[:, :, 0], rotations[:, :, 1]], axis=-1) + expected = np.r_[left[:3], rot6d[0], 0.25, right[:3], rot6d[1], 0.75] + + # Live observation encoding: compare physical matrices, not an inverse pair. + proprio = action_conversion.robotwin_endpose_to_eef20d(left, right, 0.25, 0.75) + np.testing.assert_allclose(proprio, expected, atol=1e-6) + + # Both the training loader and normalization-stat reader consume native HDF5. + with h5py.File(tmp_path / "episode.hdf5", "w") as f: + f.create_dataset("endpose/left_endpose", data=[left, left]) + f.create_dataset("endpose/right_endpose", data=[right, right]) + f.create_dataset("endpose/left_gripper", data=[0.25, 0.25]) + f.create_dataset("endpose/right_gripper", data=[0.75, 0.75]) + dataset = RoboTwinDataset.__new__(RoboTwinDataset) + np.testing.assert_allclose(dataset._read_eef_actions(f, 1, 2), expected[None], atol=1e-6) + np.testing.assert_allclose(stats._read_eef_actions(f), np.stack([expected, expected]), atol=1e-6) + + # Decode independently constructed physical rot6d, not the encoder's output. + action = action_conversion.eef20d_to_ee16d(expected.astype(np.float32)) + np.testing.assert_allclose( + action[[0, 1, 2, 7, 8, 9, 10, 15]], + expected[[0, 1, 2, 9, 10, 11, 12, 19]], + ) + native_quats = np.stack([action[3:7], action[11:15]]) + actual_rotations = Rotation.from_quat(native_quats[:, [1, 2, 3, 0]]).as_matrix() + np.testing.assert_allclose(actual_rotations, rotations, atol=1e-6) + assert proprio.dtype == action.dtype == np.float32 + + +def test_stats_recompute_instead_of_reusing_legacy_quaternion_shards(tmp_path): + root = tmp_path / "task_a" / "test-robot_clean_50" / "data" + root.mkdir(parents=True) + native_identity = np.array([[0, 0, 0, 1, 0, 0, 0]], dtype=np.float64) + with h5py.File(root / "episode0.hdf5", "w") as f: + for arm in ("left", "right"): + f.create_dataset(f"endpose/{arm}_endpose", data=native_identity) + f.create_dataset(f"endpose/{arm}_gripper", data=[0.5]) + + checkpoint = str(tmp_path / "stats.npy") + digest = hashlib.sha1(str(root.resolve()).encode("utf-8")).hexdigest()[:16] + legacy_shard = Path(f"{checkpoint}.partial/shards_v1/{digest}.npz") + legacy_shard.parent.mkdir(parents=True) + # An old shard must not be read, even if its task-root cache key matches. + legacy_eef = np.array([[0, 0, 0, 1, 0, 0, 0, -1, 0, 0.5] * 2], dtype=np.float64) + np.savez(legacy_shard, joint=np.empty((0, 14)), eef=legacy_eef) + legacy_bytes = legacy_shard.read_bytes() + stats.compute_multitask_robotwin_stats( + dataset_dir=str(tmp_path), + embodiment="test-robot", + variant="clean_50", + tasks=["task_a"], + checkpoint_path=checkpoint, + ) + with np.load(stats._shard_path_for_root(checkpoint, str(root)), allow_pickle=False) as payload: + np.testing.assert_allclose(payload["eef"][0, 3:9], [1, 0, 0, 0, 1, 0]) + assert legacy_shard.read_bytes() == legacy_bytes