| """xtrainer2 command layout used between StarVLA and the real robot. |
| |
| The converted LeRobot dataset stores each arm as ``xyz + quaternion + gripper``. |
| StarVLA / PyTorch3D read that quaternion as ``w, x, y, z``. The robot HDF5 |
| streams and the robot controller use ``x, y, z, w``. |
| |
| The training transform already maps quaternion <-> rotation-6D. After |
| un-normalization the server therefore has 16-D **wxyz** commands. Sending |
| those 16 numbers to the robot without this reorder will execute the wrong |
| orientation. |
| |
| Do not confuse these widths: |
| |
| * 14 – joint angles in the source HDF5 (not used for this VLA) |
| * 16 – stored / robot EEF command (after this conversion, xyzw) |
| * 20 – model tensor after rotation-6D (horizon is a separate axis) |
| """ |
|
|
| from __future__ import annotations |
|
|
| import numpy as np |
|
|
| from starVLA.dataloader.relative_pose import compose_quaternion |
|
|
| QUATERNION_STARTS = (3, 11) |
| POSITION_SLICES = ((0, 3), (8, 11)) |
| ACTION_DIM = 16 |
| ROBOT_TYPE = "xtrainer2_bimanual_eef" |
| SOURCE_QUATERNION_FORMAT = "xyzw" |
| MODEL_QUATERNION_FORMAT = "wxyz" |
|
|
| |
| |
| |
| COMMAND_LAYOUT = "xtrainer2_bimanual_eef_xyzw" |
|
|
|
|
| def _as_float_array(values: np.ndarray) -> np.ndarray: |
| array = np.asarray(values, dtype=np.float32) |
| if array.shape[-1] != ACTION_DIM: |
| raise ValueError(f"expected last dimension {ACTION_DIM}, got {array.shape}") |
| return array |
|
|
|
|
| def wxyz_to_xyzw_block(block: np.ndarray) -> np.ndarray: |
| """``w, x, y, z`` -> ``x, y, z, w``.""" |
| return np.concatenate([block[..., 1:4], block[..., :1]], axis=-1) |
|
|
|
|
| def xyzw_to_wxyz_block(block: np.ndarray) -> np.ndarray: |
| """``x, y, z, w`` -> ``w, x, y, z``.""" |
| return np.concatenate([block[..., 3:4], block[..., :3]], axis=-1) |
|
|
|
|
| def model_actions_to_robot(values: np.ndarray) -> np.ndarray: |
| """Convert un-normalized StarVLA actions (wxyz) to robot commands (xyzw).""" |
| converted = np.array(_as_float_array(values), copy=True) |
| for start in QUATERNION_STARTS: |
| converted[..., start : start + 4] = wxyz_to_xyzw_block(converted[..., start : start + 4]) |
| return converted |
|
|
|
|
| def robot_state_to_model(values: np.ndarray) -> np.ndarray: |
| """Convert robot observations (xyzw) to the layout StarVLA was trained on.""" |
| converted = np.array(_as_float_array(values), copy=True) |
| for start in QUATERNION_STARTS: |
| converted[..., start : start + 4] = xyzw_to_wxyz_block(converted[..., start : start + 4]) |
| return converted |
|
|
|
|
| def compose_relative_eef_wxyz(relative_actions: np.ndarray, state_wxyz: np.ndarray) -> np.ndarray: |
| """Turn rel-mode 16-D wxyz actions into absolute 16-D wxyz poses. |
| |
| Positions add the current EE; rotations left-multiply the current |
| quaternion; grippers are already absolute and are left unchanged. |
| """ |
| relative = np.array(_as_float_array(relative_actions), copy=True) |
| state = np.asarray(state_wxyz, dtype=np.float32).reshape(-1, ACTION_DIM)[0] |
| absolute = relative |
| for start, end in POSITION_SLICES: |
| absolute[..., start:end] = relative[..., start:end] + state[start:end] |
| for start in QUATERNION_STARTS: |
| absolute[..., start : start + 4] = compose_quaternion( |
| state[start : start + 4], |
| relative[..., start : start + 4], |
| ) |
| return absolute |
|
|
|
|
| def is_relative_action_mode(action_mode: object | None) -> bool: |
| if action_mode is None: |
| return False |
| return str(action_mode).lower() in {"rel", "relative"} |
|
|
