6d: continuous rotation representation, bimanual layouts, shared norm helpers

egomimic/rldb/zarr/action_chunk_transforms.py

@@ -28,9 +28,11 @@
     _matrix_to_xyz,
     _matrix_to_xyzwxyz,
     _matrix_to_xyzypr,
+    _rot6d_to_ypr,
     _xyz_to_matrix,
     _xyzwxyz_to_matrix,
     _xyzypr_to_matrix,
+    _ypr_to_rot6d,
     wxyz_to_xyzw,
     xyzw_to_wxyz,
 )
@@ -387,6 +389,109 @@ def transform(self, batch: dict) -> dict:
         return batch
 
 
+def _as_array(value):
+    """(numpy array, was_tensor) for a batch value; tensors are copied to numpy."""
+    is_tensor = isinstance(value, torch.Tensor)
+    arr = value.detach().cpu().numpy() if is_tensor else np.asarray(value)
+    return arr, is_tensor
+
+
+def _like_input(out: np.ndarray, is_tensor: bool):
+    return torch.from_numpy(out) if is_tensor else out
+
+
+class CartesianYPRToRot6D(Transform):
+    """Convert a bimanual cartesian vector from per-arm xyz+ypr(+gripper) to
+    per-arm xyz+rot6d(+gripper).
+
+    ``rot6d`` is the continuous 6D rotation representation = the first two
+    columns of the rotation matrix, packed as [col0(3), col1(3)] (see
+    :func:`egomimic.utils.pose_utils._ypr_to_rot6d`). This matches the column
+    convention of the pi0.5 32D action blocks, so the resulting per-arm layout
+    packs into them with no rotation math in the model.
+
+    Input layouts (last dim), for action chunks ``(T, D)`` and single proprio
+    poses ``(D,)`` alike:
+      12 -> [L xyz ypr, R xyz ypr]       -> 18 [L xyz 6d, R xyz 6d]
+      14 -> [L xyz ypr g, R xyz ypr g]   -> 20 [L xyz 6d g, R xyz 6d g]
+
+    Preserves the numpy/tensor type of the input.
+    """
+
+    def __init__(
+        self, action_key: str = "actions_cartesian", output_key: str | None = None
+    ):
+        self.action_key = action_key
+        self.output_key = output_key or action_key
+
+    def transform(self, batch: dict) -> dict:
+        arr, is_tensor = _as_array(batch[self.action_key])
+        D = arr.shape[-1]
+        if D == 14:
+            l_xyz, l_ypr, l_g = arr[..., 0:3], arr[..., 3:6], arr[..., 6:7]
+            r_xyz, r_ypr, r_g = arr[..., 7:10], arr[..., 10:13], arr[..., 13:14]
+            out = np.concatenate(
+                [l_xyz, _ypr_to_rot6d(l_ypr), l_g, r_xyz, _ypr_to_rot6d(r_ypr), r_g],
+                axis=-1,
+            )
+        elif D == 12:
+            l_xyz, l_ypr = arr[..., 0:3], arr[..., 3:6]
+            r_xyz, r_ypr = arr[..., 6:9], arr[..., 9:12]
+            out = np.concatenate(
+                [l_xyz, _ypr_to_rot6d(l_ypr), r_xyz, _ypr_to_rot6d(r_ypr)],
+                axis=-1,
+            )
+        else:
+            raise ValueError(
+                f"CartesianYPRToRot6D expects last-dim 12 or 14, got {arr.shape} "
+                f"for '{self.action_key}'"
+            )
+        batch[self.output_key] = _like_input(out, is_tensor)
+        return batch
+
+
+class CartesianRot6DToYPR(Transform):
+    """Inverse of :class:`CartesianYPRToRot6D`: per-arm xyz+rot6d(+gripper) ->
+    xyz+ypr(+gripper). Gram-Schmidt re-orthonormalizes the two columns, so a
+    non-orthonormal model prediction still yields a proper rotation.
+
+    Input layouts (last dim):
+      18 -> [L xyz 6d, R xyz 6d]         -> 12 [L xyz ypr, R xyz ypr]
+      20 -> [L xyz 6d g, R xyz 6d g]     -> 14 [L xyz ypr g, R xyz ypr g]
+    """
+
+    def __init__(
+        self, action_key: str = "actions_cartesian", output_key: str | None = None
+    ):
+        self.action_key = action_key
+        self.output_key = output_key or action_key
+
+    def transform(self, batch: dict) -> dict:
+        arr, is_tensor = _as_array(batch[self.action_key])
+        D = arr.shape[-1]
+        if D == 20:
+            l_xyz, l_6d, l_g = arr[..., 0:3], arr[..., 3:9], arr[..., 9:10]
+            r_xyz, r_6d, r_g = arr[..., 10:13], arr[..., 13:19], arr[..., 19:20]
+            out = np.concatenate(
+                [l_xyz, _rot6d_to_ypr(l_6d), l_g, r_xyz, _rot6d_to_ypr(r_6d), r_g],
+                axis=-1,
+            )
+        elif D == 18:
+            l_xyz, l_6d = arr[..., 0:3], arr[..., 3:9]
+            r_xyz, r_6d = arr[..., 9:12], arr[..., 12:18]
+            out = np.concatenate(
+                [l_xyz, _rot6d_to_ypr(l_6d), r_xyz, _rot6d_to_ypr(r_6d)],
+                axis=-1,
+            )
+        else:
+            raise ValueError(
+                f"CartesianRot6DToYPR expects last-dim 18 or 20, got {arr.shape} "
+                f"for '{self.action_key}'"
+            )
+        batch[self.output_key] = _like_input(out, is_tensor)
+        return batch
+
+
 class CartesianWithGripperCoordinateTransform(Transform):
     def __init__(
         self,

egomimic/rldb/zarr/zarr_dataset_multi.py

@@ -44,6 +44,10 @@
 # from action_chunk_transforms import Transform
 from egomimic.rldb.filters import DatasetFilter
 from egomimic.rldb.resolve_memo import memoized
+from egomimic.utils.action_utils import (
+    _apply_norm_one,
+    _apply_unnorm_one,
+)
 from egomimic.utils.env import load_env
 
 
@@ -1494,58 +1498,10 @@ def cache_stats(self, save_cache_dir: str):
     # ---- normalize / unnormalize ----
 
     def _apply_norm_one(self, tensor, stats):
-        if self.norm_mode == "zscore":
-            mean = torch.as_tensor(
-                stats["mean"], device=tensor.device, dtype=torch.float32
-            )
-            std = torch.as_tensor(
-                stats["std"], device=tensor.device, dtype=torch.float32
-            )
-            return (tensor - mean) / (std + 1e-6)
-        if self.norm_mode == "minmax":
-            mn = torch.as_tensor(
-                stats["min"], device=tensor.device, dtype=torch.float32
-            )
-            mx = torch.as_tensor(
-                stats["max"], device=tensor.device, dtype=torch.float32
-            )
-            return 2.0 * ((tensor - mn) / (mx - mn + 1e-6)) - 1.0
-        if self.norm_mode == "quantile":
-            q1 = torch.as_tensor(
-                stats["quantile_1"], device=tensor.device, dtype=torch.float32
-            )
-            q99 = torch.as_tensor(
-                stats["quantile_99"], device=tensor.device, dtype=torch.float32
-            )
-            return 2.0 * ((tensor - q1) / (q99 - q1 + 1e-6)) - 1.0
-        raise ValueError(f"Invalid normalization mode: {self.norm_mode}")
+        return _apply_norm_one(tensor, stats, self.norm_mode)
 
     def _apply_unnorm_one(self, tensor, stats):
-        if self.norm_mode == "zscore":
-            mean = torch.as_tensor(
-                stats["mean"], device=tensor.device, dtype=torch.float32
-            )
-            std = torch.as_tensor(
-                stats["std"], device=tensor.device, dtype=torch.float32
-            )
-            return tensor * (std + 1e-6) + mean
-        if self.norm_mode == "minmax":
-            mn = torch.as_tensor(
-                stats["min"], device=tensor.device, dtype=torch.float32
-            )
-            mx = torch.as_tensor(
-                stats["max"], device=tensor.device, dtype=torch.float32
-            )
-            return (tensor + 1) * 0.5 * (mx - mn + 1e-6) + mn
-        if self.norm_mode == "quantile":
-            q1 = torch.as_tensor(
-                stats["quantile_1"], device=tensor.device, dtype=torch.float32
-            )
-            q99 = torch.as_tensor(
-                stats["quantile_99"], device=tensor.device, dtype=torch.float32
-            )
-            return (tensor + 1) * 0.5 * (q99 - q1 + 1e-6) + q1
-        raise ValueError(f"Invalid normalization mode: {self.norm_mode}")
+        return _apply_unnorm_one(tensor, stats, self.norm_mode)
 
     def normalize(self, data: dict, embodiment_id: int) -> dict:
         if not self.norm_stats.get(embodiment_id):

egomimic/utils/action_utils.py

@@ -1,4 +1,4 @@
-from typing import Dict, Tuple
+from typing import Any, Dict, Tuple
 
 import torch
 
@@ -43,6 +43,54 @@ def _pad32(x: torch.Tensor) -> torch.Tensor:
     return x[..., :32]
 
 
+def _stat_tensor(stats: dict[str, Any], key: str, ref: torch.Tensor) -> torch.Tensor:
+    value = torch.as_tensor(stats[key], device=ref.device, dtype=torch.float32)
+    return value.to(dtype=ref.dtype if ref.is_floating_point() else torch.float32)
+
+
+def _norm_center_scale(
+    stats: dict[str, Any], norm_mode: str, ref: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """``(center, half_range)`` so that ``normalized = (x - center) / half_range``
+    (zscore: ``(mean, std)``)."""
+    if norm_mode == "zscore":
+        return _stat_tensor(stats, "mean", ref), _stat_tensor(stats, "std", ref)
+    if norm_mode == "minmax":
+        lo, hi = _stat_tensor(stats, "min", ref), _stat_tensor(stats, "max", ref)
+    elif norm_mode == "quantile":
+        lo = _stat_tensor(stats, "quantile_1", ref)
+        hi = _stat_tensor(stats, "quantile_99", ref)
+    else:
+        raise ValueError(f"Invalid normalization mode: {norm_mode}")
+    return 0.5 * (lo + hi), 0.5 * (hi - lo)
+
+
+def _eps(norm_mode: str) -> float:
+    # Keeps the historical denominators: std + 1e-6 and (hi - lo) + 1e-6.
+    return 1e-6 if norm_mode == "zscore" else 0.5e-6
+
+
+def _apply_norm_one(
+    tensor: torch.Tensor,
+    stats: dict[str, Any],
+    norm_mode: str,
+) -> torch.Tensor:
+    """Normalize with per-key stats (zscore, or [-1, 1] for minmax/quantile)."""
+    center, scale = _norm_center_scale(stats, norm_mode, tensor)
+    return (tensor - center) / (scale + _eps(norm_mode))
+
+
+def _apply_unnorm_one(
+    tensor: torch.Tensor,
+    stats: dict[str, Any],
+    norm_mode: str,
+) -> torch.Tensor:
+    """Inverse of :func:`_apply_norm_one`."""
+    center, scale = _norm_center_scale(stats, norm_mode, tensor)
+    return tensor * (scale + _eps(norm_mode)) + center
+
+
+# ---------- rotation helpers (shared with the eval metrics) ----------
 def _ypr_to_matrix(ypr: torch.Tensor, degrees: bool = False) -> torch.Tensor:
     if degrees:
         ypr = ypr * (torch.pi / 180.0)

egomimic/utils/pose_utils.py

@@ -130,6 +130,103 @@ def _xyzypr_to_matrix(xyzypr: np.ndarray) -> np.ndarray:
     return mats
 
 
+def _ypr_to_rot6d(ypr: np.ndarray) -> np.ndarray:
+    """Convert euler ypr to the continuous 6D rotation representation.
+
+    args:
+        ypr: (..., 3) array of [yaw, pitch, roll] (radians, ZYX convention)
+    returns:
+        (..., 6) array = first two columns of the rotation matrix,
+        concatenated as [col0(3), col1(3)].
+
+    Matches the column convention of the torch helpers in
+    ``egomimic.utils.action_utils`` (``_ypr_to_matrix`` = Rz@Ry@Rx and
+    ``_reconstruct_R_from_cols`` taking ``R[..., 0]`` / ``R[..., 1]``).
+    """
+    ypr = np.asarray(ypr)
+    if ypr.shape[-1] != 3:
+        raise ValueError(f"Expected (..., 3) ypr, got shape {ypr.shape}")
+    dtype = ypr.dtype if np.issubdtype(ypr.dtype, np.floating) else np.float64
+    shape = ypr.shape[:-1]
+    flat = ypr.reshape(-1, 3).astype(np.float64)
+    mats = R.from_euler("ZYX", flat, degrees=False).as_matrix()  # (N, 3, 3)
+    six = np.concatenate([mats[:, :, 0], mats[:, :, 1]], axis=-1)  # cols 0, 1
+    return six.reshape(*shape, 6).astype(dtype, copy=False)
+
+
+def _rot6d_to_ypr(six: np.ndarray) -> np.ndarray:
+    """Inverse of :func:`_ypr_to_rot6d`.
+
+    args:
+        six: (..., 6) array = [col0(3), col1(3)] of a rotation matrix.
+    returns:
+        (..., 3) array of [yaw, pitch, roll] (radians, ZYX convention).
+
+    Reconstructs a proper rotation via Gram-Schmidt (mirroring
+    ``_reconstruct_R_from_cols`` in ``action_utils``) before extracting euler
+    angles, so a non-orthonormal model prediction still yields a valid pose
+    and ``_rot6d_to_ypr(_ypr_to_rot6d(ypr)) == ypr``.
+    """
+    six = np.asarray(six)
+    if six.shape[-1] != 6:
+        raise ValueError(f"Expected (..., 6) rot6d, got shape {six.shape}")
+    dtype = six.dtype if np.issubdtype(six.dtype, np.floating) else np.float64
+    shape = six.shape[:-1]
+    flat = six.reshape(-1, 6).astype(np.float64)
+    c1 = flat[:, 0:3]
+    c2 = flat[:, 3:6]
+    eps = 1e-8
+    c1n = c1 / np.clip(np.linalg.norm(c1, axis=-1, keepdims=True), eps, None)
+    proj = np.sum(c2 * c1n, axis=-1, keepdims=True) * c1n
+    c2o = c2 - proj
+    c2n = c2o / np.clip(np.linalg.norm(c2o, axis=-1, keepdims=True), eps, None)
+    c3n = np.cross(c1n, c2n)
+    mats = np.stack([c1n, c2n, c3n], axis=-1)  # columns
+    ypr = R.from_matrix(mats).as_euler("ZYX", degrees=False)
+    return ypr.reshape(*shape, 3).astype(dtype, copy=False)
+
+
+# [left arm | right arm]; per-arm blocks are one of:
+#   ypr: xyz(3) + ypr(3)            [+ gripper(1)]
+#   6d:  xyz(3) + col0(3) + col1(3) [+ gripper(1)]
+# Mapping width -> {xyz, rot, grip} channel indices. xyz/grip are bounded,
+# linearly-interpolated channels; the rot channels are either Euler (wrap at
+# +-pi) or continuous 6D columns (bounded in ~[-1, 1]), so quantile bounds on
+# them are meaningless. The norm-stat bounds check and the eval metrics both
+# consume this to know which channel is which.
+BIMANUAL_CARTESIAN_LAYOUTS = {
+    12: {  # human ypr:  [L xyz ypr | R xyz ypr]
+        "xyz": (0, 1, 2, 6, 7, 8),
+        "rot": (3, 4, 5, 9, 10, 11),
+        "grip": (),
+    },
+    14: {  # robot ypr:  [L xyz ypr g | R xyz ypr g]
+        "xyz": (0, 1, 2, 7, 8, 9),
+        "rot": (3, 4, 5, 10, 11, 12),
+        "grip": (6, 13),
+    },
+    18: {  # human 6d:   [L xyz c0 c1 | R xyz c0 c1]
+        "xyz": (0, 1, 2, 9, 10, 11),
+        "rot": (3, 4, 5, 6, 7, 8, 12, 13, 14, 15, 16, 17),
+        "grip": (),
+    },
+    20: {  # robot 6d:   [L xyz c0 c1 g | R xyz c0 c1 g]
+        "xyz": (0, 1, 2, 10, 11, 12),
+        "rot": (3, 4, 5, 6, 7, 8, 13, 14, 15, 16, 17, 18),
+        "grip": (9, 19),
+    },
+}
+
+
+def bimanual_cartesian_layout(width: int) -> dict | None:
+    """Index layout for a bimanual cartesian action/proprio vector.
+
+    Returns a dict with ``xyz`` / ``rot`` / ``grip`` index tuples, or ``None``
+    if ``width`` is not a recognized native width (12/14 ypr, 18/20 6D).
+    """
+    return BIMANUAL_CARTESIAN_LAYOUTS.get(int(width))
+
+
 def _matrix_to_xyzwxyz(mats: np.ndarray) -> np.ndarray:
     """
     args: