| """Dataset normalisation statistics and delta-base inversion for online execution.""" |
|
|
| from __future__ import annotations |
|
|
| import json |
| from dataclasses import dataclass |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
| from .data_utils import matrix_to_rotation_6d as _matrix_to_rotation_6d |
| from .data_utils import rotation_6d_to_matrix |
|
|
|
|
| def delta_base_to_absolute( |
| reference_state: np.ndarray, delta_base_action: np.ndarray |
| ) -> np.ndarray: |
| """Invert T-Rex delta-base arm poses; hand targets are already absolute.""" |
|
|
| reference = np.asarray(reference_state, dtype=np.float64) |
| delta = np.asarray(delta_base_action, dtype=np.float64) |
| if reference.shape != (62,) or delta.shape[-1] != 62: |
| raise ValueError("expected reference [62] and delta action [...,62]") |
| output = np.empty_like(delta) |
| for pose_slice, hand_slice in ( |
| (slice(0, 9), slice(9, 31)), |
| (slice(31, 40), slice(40, 62)), |
| ): |
| reference_pose = reference[pose_slice] |
| delta_pose = delta[..., pose_slice] |
| reference_rotation = rotation_6d_to_matrix(reference_pose[3:9]) |
| delta_rotation = rotation_6d_to_matrix(delta_pose[..., 3:9]) |
| target_translation = ( |
| reference_pose[:3] |
| + np.einsum("ij,...j->...i", reference_rotation, delta_pose[..., :3]) |
| ) |
| target_rotation = np.einsum( |
| "ij,...jk->...ik", reference_rotation, delta_rotation |
| ) |
| output[..., pose_slice] = np.concatenate( |
| (target_translation, _matrix_to_rotation_6d(target_rotation)), axis=-1 |
| ) |
| output[..., hand_slice] = delta[..., hand_slice] |
| return output.astype(np.float32) |
|
|
|
|
| @dataclass(frozen=True) |
| class TrexRuntimeStatistics: |
| action_q01: np.ndarray |
| action_q99: np.ndarray |
| state_q01: np.ndarray |
| state_q99: np.ndarray |
| force_q01: np.ndarray |
| force_q99: np.ndarray |
|
|
| @classmethod |
| def from_dataset(cls, dataset_root: str | Path) -> "TrexRuntimeStatistics": |
| root = Path(dataset_root) |
| with (root / "meta" / "relative_stats_dreamzero.json").open() as handle: |
| action = json.load(handle)["eef62"] |
| with (root / "meta" / "stats.json").open() as handle: |
| stats = json.load(handle) |
| force = stats["observation.tactile_force"] |
| state = stats["observation.state_eef62"] |
| return cls( |
| action_q01=np.asarray(action["q01"], dtype=np.float32), |
| action_q99=np.asarray(action["q99"], dtype=np.float32), |
| state_q01=np.asarray(state["q01"], dtype=np.float32), |
| state_q99=np.asarray(state["q99"], dtype=np.float32), |
| force_q01=np.asarray(force["q01"], dtype=np.float32).reshape(10, 6), |
| force_q99=np.asarray(force["q99"], dtype=np.float32).reshape(10, 6), |
| ) |
|
|
| def normalize_force(self, force: np.ndarray) -> np.ndarray: |
| values = np.asarray(force, dtype=np.float32) |
| scale = self.force_q99 - self.force_q01 |
| normalized = np.where( |
| scale != 0, |
| 2.0 * (values - self.force_q01) / np.where(scale == 0, 1.0, scale) |
| - 1.0, |
| values, |
| ) |
| return np.clip(normalized, -1.0, 1.0).astype(np.float32) |
|
|
| def normalize_state(self, state: np.ndarray) -> np.ndarray: |
| values = np.asarray(state, dtype=np.float32) |
| scale = self.state_q99 - self.state_q01 |
| normalized = np.where( |
| scale != 0, |
| 2.0 * (values - self.state_q01) / np.where(scale == 0, 1.0, scale) |
| - 1.0, |
| values, |
| ) |
| return np.clip(normalized, -1.0, 1.0).astype(np.float32) |
|
|
| def denormalize_action(self, normalized_action: np.ndarray) -> np.ndarray: |
| values = np.asarray(normalized_action, dtype=np.float32)[..., :62] |
| return ( |
| (values + 1.0) * 0.5 * (self.action_q99 - self.action_q01) |
| + self.action_q01 |
| ).astype(np.float32) |
|
|
|
|
|
|
| __all__ = ["TrexRuntimeStatistics", "delta_base_to_absolute"] |
|
|
