Update Wall-X to 1.1.0 (#104)

wall_x/data/backends/lerobot/rotation_layout.py

@@ -0,0 +1,95 @@
+"""Layout helpers when config expects 6D rotation but LeRobot stores 3D Euler."""
+
+from __future__ import annotations
+
+import numpy as np
+
+from wall_x._vendor.x2robot_utils.geometry import euler_to_matrix_zyx_6d_nb
+
+LAYOUT_SKIP_KEYS = frozenset(
+    {"velocity_decomposed", "height", "head_actions", "action_padding"}
+)
+ROTATION_KEYWORD = "rotation"
+ROTATION_6D_KEYWORD = "6D"
+
+
+def layout_uses_6d_rotation(layout_config: dict) -> bool:
+    for key, dim in layout_config.items():
+        if key in LAYOUT_SKIP_KEYS:
+            continue
+        if ROTATION_KEYWORD in key and ROTATION_6D_KEYWORD in key and dim == 6:
+            return True
+    return False
+
+
+def euler_layout_dim(layout_config: dict) -> int:
+    """Vector width in LeRobot when rotation slices are still 3D Euler."""
+    total = 0
+    for key, dim in layout_config.items():
+        if key in LAYOUT_SKIP_KEYS:
+            continue
+        if ROTATION_KEYWORD in key and ROTATION_6D_KEYWORD in key and dim == 6:
+            total += 3
+        else:
+            total += int(dim)
+    return total
+
+
+def convert_euler_to_6d(vec: np.ndarray, layout_config: dict) -> np.ndarray:
+    """Rewrite [pos, euler(3), tail...] to [pos, rot6d(6), tail...] per layout."""
+    vec = np.asarray(vec, dtype=np.float64)
+    single = vec.ndim == 1
+    if single:
+        vec = vec[np.newaxis, :]
+
+    out_rows = []
+    for row in vec:
+        parts: list[np.ndarray] = []
+        raw_cur = 0
+        for key, dim in layout_config.items():
+            if key in LAYOUT_SKIP_KEYS:
+                continue
+            dim = int(dim)
+            if ROTATION_KEYWORD in key and ROTATION_6D_KEYWORD in key and dim == 6:
+                euler = row[raw_cur : raw_cur + 3]
+                rot6d = euler_to_matrix_zyx_6d_nb(euler.reshape(1, 3)).reshape(6)
+                parts.append(rot6d)
+                raw_cur += 3
+            else:
+                parts.append(row[raw_cur : raw_cur + dim])
+                raw_cur += dim
+        out_rows.append(np.concatenate(parts, axis=0))
+
+    out = np.stack(out_rows, axis=0)
+    return out[0] if single else out
+
+
+def maybe_convert_norm_stats_vector(
+    values,
+    layout_config: dict,
+    enabled: bool | None = None,
+):
+    """Convert a 1D norm-stat vector (q01/q99/mean/std) from Euler layout to 6D."""
+    if enabled is None:
+        enabled = layout_uses_6d_rotation(layout_config)
+    if not enabled or not layout_config:
+        return values
+    arr = np.asarray(values, dtype=np.float64)
+    if arr.ndim != 1:
+        return values
+    raw_dim = euler_layout_dim(layout_config)
+    if arr.shape[0] != raw_dim:
+        return values
+    return convert_euler_to_6d(arr, layout_config).astype(np.float32)
+
+
+def maybe_convert_euler_to_6d(
+    vec: np.ndarray, layout_config: dict, enabled: bool
+) -> np.ndarray:
+    if not enabled or not layout_config:
+        return vec
+    raw_dim = euler_layout_dim(layout_config)
+    arr = np.asarray(vec)
+    if arr.shape[-1] != raw_dim:
+        return vec
+    return convert_euler_to_6d(arr, layout_config).astype(np.float32)