File size: 3,595 Bytes
8101696 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 | """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"
#: Value of ``DataConfig.robot_command_layout`` that selects the conversions in
#: this module. The policy server reads that attribute instead of matching on a
#: class name.
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"}
|