fix(robotwin): preserve native wxyz quaternion convention in EEF paths

README.md

@@ -73,6 +73,10 @@ All architectures are selected via `configs/model/<framework>.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/` |

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

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)]

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,22 +278,25 @@ 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)
     right_grip = np.asarray(right_gripper, dtype=np.float32).reshape(-1)
     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)
 
 

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,15 +713,15 @@ 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)
 
         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,

openwam/dataloader/transforms/rotation.py

@@ -153,19 +153,18 @@ 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.
 # ---------------------------------------------------------------------------
 
 
 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)

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,15 +136,15 @@ 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,
     )
     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 ``<checkpoint_path>.partial/shards_v1/`` after each
+            shard into ``<checkpoint_path>.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.

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
 

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),
         },

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