groot/vla/model/trex_track_force/runtime.py

"""Online 20 Hz / 5 Hz execution helpers for T-Rex Track-Force."""

from __future__ import annotations

from dataclasses import dataclass
import json
from pathlib import Path
from typing import Any

import numpy as np
import torch

from .force import FORCE_OFFSETS
from .policy import WANTrackForcePolicyHead
from .dataset import rotation_6d_to_matrix


def _matrix_to_rotation_6d(matrix: np.ndarray) -> np.ndarray:
    return np.concatenate((matrix[..., :, 0], matrix[..., :, 1]), axis=-1)


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)


@dataclass
class TactileRefinementState:
    coarse_action_at_split: torch.Tensor
    coarse_memory: torch.Tensor
    refined_action: torch.Tensor
    latest_offset: int = 0


class TrexTrackForceOnlineController:
    """Cache coarse output and refresh only suffixes at steps 0/4/8/12."""

    def __init__(self, policy: WANTrackForcePolicyHead) -> None:
        self.policy = policy
        self.state: TactileRefinementState | None = None

    @torch.no_grad()
    def start_chunk(
        self,
        model_inputs: dict[str, Any],
        *,
        seed: int | None = None,
    ) -> torch.Tensor:
        result = self.policy.sample(
            model_inputs, seed=seed, return_refinement_state=True
        )
        self.state = TactileRefinementState(
            coarse_action_at_split=result["coarse_action_at_split"].detach(),
            coarse_memory=result["coarse_memory"].detach(),
            refined_action=result["action_pred"].detach(),
            latest_offset=0,
        )
        return self.state.refined_action

    @torch.no_grad()
    def refresh_tactile(
        self,
        current_force: torch.Tensor,
        tactile_history_or_codes: torch.Tensor,
        *,
        executed_steps: int,
    ) -> torch.Tensor:
        if self.state is None:
            raise RuntimeError("start_chunk must be called before tactile refresh")
        if executed_steps not in FORCE_OFFSETS:
            raise ValueError(f"executed_steps must be one of {FORCE_OFFSETS}")
        if executed_steps < self.state.latest_offset:
            raise ValueError("tactile refresh offsets must be monotonic")
        refreshed = self.policy.refine_action_suffix(
            self.state.coarse_action_at_split,
            self.state.coarse_memory,
            current_force,
            tactile_history_or_codes,
            update_offset=executed_steps,
        )
        # Preserve commands already emitted from earlier refinements.
        refreshed[:, :executed_steps] = self.state.refined_action[:, :executed_steps]
        self.state.refined_action = refreshed.detach()
        self.state.latest_offset = executed_steps
        return self.state.refined_action

    def action_at(self, step: int) -> torch.Tensor:
        if self.state is None:
            raise RuntimeError("no active action chunk")
        if not 0 <= step < self.state.refined_action.shape[1]:
            raise IndexError(step)
        return self.state.refined_action[:, step]


__all__ = [
    "TactileRefinementState",
    "TrexRuntimeStatistics",
    "TrexTrackForceOnlineController",
    "delta_base_to_absolute",
]