groot/vla/model/trex_track_force/transform.py

"""Model-side transform helpers for fixed-size track/force columns."""

from __future__ import annotations

from typing import Any

import numpy as np
from pydantic import Field

from groot.vla.model.dreamzero.transform.dreamzero_cotrain import (
    DefaultDataCollator,
    DreamTransform,
)

from .force import (
    ACTION_HORIZON,
    FORCE_HISTORY_FRAMES,
    FORCE_OFFSETS,
    FORCE_SENSOR_COUNT,
    FORCE_SENSOR_DIM,
)
from .track import NUM_TRACK_POINTS, TRACK_HORIZON


TRACK_XY_ALIASES = (
    "track_xy",
    "tracks",
    "track.xy",
    "observation.track_xy",
    "observation.tracks_xy",
)
TRACK_VISIBILITY_ALIASES = (
    "track_visibility",
    "visibility",
    "vis",
    "track.visibility",
    "observation.track_visibility",
    "observation.tracks_visibility",
)
FORCE_CURRENT_ALIASES = (
    "current_force",
    "force_current",
    "observation.force_current",
)
FORCE_HISTORY_ALIASES = (
    "force_history",
    "tactile_force_history",
    "force_history_raw",
    "observation.tactile_force",
    "force_history_vq",
    "force_vq_codes",
    "observation.force_history_vq",
)
FORCE_HISTORY_PADDING_MASK_ALIASES = (
    "force_history_padding_mask",
    "tactile_force_history_padding_mask",
)


def _first_present(data: dict[str, Any], names: tuple[str, ...]) -> Any | None:
    for name in names:
        if name in data:
            return data[name]
    return None


def _reshape_track_xy(values: Any) -> np.ndarray:
    xy = np.asarray(values, dtype=np.float32)
    if xy.ndim == 2 and xy.shape[-1] == NUM_TRACK_POINTS * 2:
        xy = xy.reshape(xy.shape[0], NUM_TRACK_POINTS, 2)
    if xy.ndim != 3 or xy.shape[1:] != (NUM_TRACK_POINTS, 2):
        raise ValueError(
            "track xy must be [T,250,2] or flattened [T,500], got "
            f"{xy.shape}"
        )
    if not np.isfinite(xy).all():
        raise ValueError("track xy contains NaN or infinity")
    if xy.min(initial=0.0) < -1e-4 or xy.max(initial=1.0) > 1.0001:
        raise ValueError("track xy must be normalized to [0,1]")
    return xy


def _reshape_track_visibility(values: Any, time: int) -> np.ndarray:
    visibility = np.asarray(values)
    if visibility.ndim == 1 and visibility.size == time * NUM_TRACK_POINTS:
        visibility = visibility.reshape(time, NUM_TRACK_POINTS)
    if visibility.shape != (time, NUM_TRACK_POINTS):
        raise ValueError(
            f"track visibility must be [T,250], got {visibility.shape}"
        )
    return visibility.astype(np.float32)


def prepare_track_windows(
    track_xy: Any,
    track_visibility: Any,
    *,
    num_blocks: int,
    history_frames: int = FORCE_HISTORY_FRAMES,
    future_horizon: int = TRACK_HORIZON,
) -> dict[str, np.ndarray]:
    """Split a contiguous track timeline into clean-past and future windows.

    The first future step immediately follows the final clean-past step.
    Boundary indices are clamped, which is deterministic for online prefixes.
    """

    if num_blocks <= 0 or history_frames <= 0:
        raise ValueError("num_blocks and history_frames must be positive")
    xy = _reshape_track_xy(track_xy)
    visibility = _reshape_track_visibility(track_visibility, xy.shape[0])
    total_time = xy.shape[0]
    minimum = history_frames + num_blocks * future_horizon
    if total_time < future_horizon:
        raise ValueError(
            f"track timeline has {total_time} frames; need at least {future_horizon}"
        )

    # A full training window is [H past][K*16 future]. For a shorter online
    # prefix, anchor at its first frame and use deterministic edge padding.
    first_anchor = history_frames - 1 if total_time >= minimum else 0
    past_xy, past_vis, future_xy, future_vis = [], [], [], []
    for block_index in range(num_blocks):
        anchor = first_anchor + block_index * future_horizon
        past_indices = np.arange(
            anchor - history_frames + 1, anchor + 1, dtype=np.int64
        )
        future_indices = np.arange(
            anchor + 1, anchor + 1 + future_horizon, dtype=np.int64
        )
        past_indices = np.clip(past_indices, 0, total_time - 1)
        future_indices = np.clip(future_indices, 0, total_time - 1)
        past_xy.append(xy[past_indices])
        past_vis.append(visibility[past_indices])
        future_xy.append(xy[future_indices])
        future_vis.append(visibility[future_indices])
    return {
        "track_past_xy": np.stack(past_xy),
        "track_past_visibility": np.stack(past_vis),
        "track_future_xy": np.stack(future_xy),
        "track_future_visibility": np.stack(future_vis),
    }


class TrexTrackForceTransform(DreamTransform):
    """DreamTransform extension that preserves track and force model inputs."""

    track_history_frames: int = FORCE_HISTORY_FRAMES
    track_horizon: int = TRACK_HORIZON
    track_points: int = NUM_TRACK_POINTS
    require_track_inputs: bool = True
    require_force_inputs: bool = True
    use_deform_tactile: bool = False
    track_xy_aliases: tuple[str, ...] = Field(default=TRACK_XY_ALIASES)
    track_visibility_aliases: tuple[str, ...] = Field(
        default=TRACK_VISIBILITY_ALIASES
    )
    force_current_aliases: tuple[str, ...] = Field(default=FORCE_CURRENT_ALIASES)
    force_history_aliases: tuple[str, ...] = Field(default=FORCE_HISTORY_ALIASES)
    force_history_padding_mask_aliases: tuple[str, ...] = Field(
        default=FORCE_HISTORY_PADDING_MASK_ALIASES
    )

    def _num_blocks(self, data: dict[str, Any]) -> int:
        if "action" not in data:
            return 1
        action = np.asarray(data["action"])
        if action.shape[0] % ACTION_HORIZON:
            raise ValueError("action timeline must contain complete 16-step blocks")
        return max(action.shape[0] // ACTION_HORIZON, 1)

    def _prepare_tracks(self, data: dict[str, Any], blocks: int) -> dict[str, np.ndarray]:
        explicit_keys = (
            "track_past_xy",
            "track_past_visibility",
            "track_future_xy",
            "track_future_visibility",
        )
        if all(key in data for key in explicit_keys):
            result = {
                key: np.asarray(data[key], dtype=np.float32) for key in explicit_keys
            }
            expected = {
                "track_past_xy": (
                    blocks,
                    self.track_history_frames,
                    self.track_points,
                    2,
                ),
                "track_past_visibility": (
                    blocks,
                    self.track_history_frames,
                    self.track_points,
                ),
                "track_future_xy": (
                    blocks,
                    self.track_horizon,
                    self.track_points,
                    2,
                ),
                "track_future_visibility": (
                    blocks,
                    self.track_horizon,
                    self.track_points,
                ),
            }
            for key, shape in expected.items():
                if result[key].shape != shape:
                    raise ValueError(f"{key} is {result[key].shape}, expected {shape}")
            return result

        xy = _first_present(data, self.track_xy_aliases)
        visibility = _first_present(data, self.track_visibility_aliases)
        if xy is None or visibility is None:
            if self.require_track_inputs:
                raise KeyError(
                    "250-point track xy/visibility columns were not provided"
                )
            return {}
        return prepare_track_windows(
            xy,
            visibility,
            num_blocks=blocks,
            history_frames=self.track_history_frames,
            future_horizon=self.track_horizon,
        )

    def _prepare_force(self, data: dict[str, Any], blocks: int) -> dict[str, np.ndarray]:
        current = _first_present(data, self.force_current_aliases)
        history = _first_present(data, self.force_history_aliases)
        history_padding_mask = _first_present(
            data, self.force_history_padding_mask_aliases
        )
        if history is None:
            if self.require_force_inputs:
                raise KeyError(
                    "raw 16x10x6 tactile history or 10 precomputed codes are required"
                )
            return {}
        history_array = np.asarray(history)
        is_raw = np.issubdtype(history_array.dtype, np.floating)
        force_updates = len(FORCE_OFFSETS)
        if is_raw:
            if history_array.shape == (
                FORCE_HISTORY_FRAMES,
                FORCE_SENSOR_COUNT * FORCE_SENSOR_DIM,
            ):
                history_array = history_array.reshape(
                    FORCE_HISTORY_FRAMES, FORCE_SENSOR_COUNT, FORCE_SENSOR_DIM
                )
            if history_array.shape == (
                blocks,
                FORCE_HISTORY_FRAMES,
                FORCE_SENSOR_COUNT * FORCE_SENSOR_DIM,
            ):
                history_array = history_array.reshape(
                    blocks,
                    FORCE_HISTORY_FRAMES,
                    FORCE_SENSOR_COUNT,
                    FORCE_SENSOR_DIM,
                )
            if history_array.shape == (
                FORCE_HISTORY_FRAMES,
                FORCE_SENSOR_COUNT,
                FORCE_SENSOR_DIM,
            ):
                history_array = np.broadcast_to(
                    history_array,
                    (
                        blocks,
                        FORCE_HISTORY_FRAMES,
                        FORCE_SENSOR_COUNT,
                        FORCE_SENSOR_DIM,
                    ),
                ).copy()
            expected_history = (
                blocks,
                FORCE_HISTORY_FRAMES,
                FORCE_SENSOR_COUNT,
                FORCE_SENSOR_DIM,
            )
            expected_update_history = (
                blocks,
                force_updates,
                FORCE_HISTORY_FRAMES,
                FORCE_SENSOR_COUNT,
                FORCE_SENSOR_DIM,
            )
            if history_array.shape not in (expected_history, expected_update_history):
                raise ValueError(
                    "raw tactile history must be [blocks,16,10,6] or "
                    f"[blocks,{force_updates},16,10,6], "
                    f"got {history_array.shape}"
                )
            history_array = history_array.astype(np.float32)
            expected_padding_shape = history_array.shape[:-2]
            if history_padding_mask is None:
                padding_array = np.zeros(expected_padding_shape, dtype=bool)
            else:
                padding_array = np.asarray(history_padding_mask, dtype=bool)
                if padding_array.shape == (FORCE_HISTORY_FRAMES,):
                    padding_array = np.broadcast_to(
                        padding_array, expected_padding_shape
                    ).copy()
                if padding_array.shape != expected_padding_shape:
                    raise ValueError(
                        "force history padding mask must match the raw history "
                        f"time axes {expected_padding_shape}, got "
                        f"{padding_array.shape}"
                    )
        else:
            if history_array.shape == (FORCE_SENSOR_COUNT,):
                history_array = np.broadcast_to(
                    history_array, (blocks, FORCE_SENSOR_COUNT)
                ).copy()
            expected_codes = (blocks, FORCE_SENSOR_COUNT)
            expected_update_codes = (blocks, force_updates, FORCE_SENSOR_COUNT)
            if history_array.shape not in (expected_codes, expected_update_codes):
                raise ValueError(
                    "precomputed tactile codes must be [blocks,10] or "
                    f"[blocks,{force_updates},10], "
                    f"got {history_array.shape}"
                )
            history_array = history_array.astype(np.int64)

        if current is None and is_raw:
            current = (
                history_array[:, :, -1]
                if history_array.ndim == 6
                else history_array[:, -1]
            )
        if current is None:
            if self.require_force_inputs:
                raise KeyError("current 10x6 force is required with precomputed codes")
            return {}
        current_array = np.asarray(current, dtype=np.float32)
        if current_array.shape == (FORCE_SENSOR_COUNT * FORCE_SENSOR_DIM,):
            current_array = current_array.reshape(FORCE_SENSOR_COUNT, FORCE_SENSOR_DIM)
        if current_array.shape == (FORCE_SENSOR_COUNT, FORCE_SENSOR_DIM):
            current_array = np.broadcast_to(
                current_array, (blocks, FORCE_SENSOR_COUNT, FORCE_SENSOR_DIM)
            ).copy()
        expected_current = (
            blocks,
            FORCE_SENSOR_COUNT,
            FORCE_SENSOR_DIM,
        )
        expected_update_current = (
            blocks,
            force_updates,
            FORCE_SENSOR_COUNT,
            FORCE_SENSOR_DIM,
        )
        if current_array.shape not in (expected_current, expected_update_current):
            raise ValueError(
                "current force must be [blocks,10,6] or "
                f"[blocks,{force_updates},10,6], got {current_array.shape}"
            )

        result = {"current_force": current_array}
        result[
            "tactile_force_history" if is_raw else "force_history_vq"
        ] = history_array
        if is_raw:
            result["tactile_force_history_padding_mask"] = padding_array
        return result

    def apply_single(self, data: dict) -> dict:
        blocks = self._num_blocks(data)
        track_data = self._prepare_tracks(data, blocks)
        force_data = self._prepare_force(data, blocks)
        transformed = super().apply_single(data)
        images = np.asarray(transformed["images"])
        expected_training_frames = 1 + blocks * 8
        if self.training and images.shape[0] == expected_training_frames:
            # Preserve DreamZero's 4-block/33-frame autoregressive geometry:
            # one clean conditioning frame plus eight 10 Hz targets per block.
            transformed["history_images"] = images[:1]
            transformed["future_images"] = images[1:]
        elif images.shape[0] >= 1 and not self.training:
            transformed["history_images"] = images
        else:
            raise ValueError(
                "T-Rex Track-Force expects one conditioning frame plus "
                f"eight frames per autoregressive block ({expected_training_frames} "
                "training frames), or one or more history frames for inference; "
                f"got {images.shape[0]}"
            )
        transformed.update(track_data)
        transformed.update(force_data)
        if self.use_deform_tactile:
            if "deform_current" not in data:
                raise KeyError(
                    "use_deform_tactile=True but the dataset provided no "
                    "deform_current frames (enable load_deform_tactile)"
                )
            transformed["deform_current"] = np.asarray(
                data["deform_current"], dtype=np.uint8
            )
        return transformed


TrackForceTransform = TrexTrackForceTransform
TrackForceDataCollator = DefaultDataCollator


__all__ = [
    "TrackForceDataCollator",
    "TrackForceTransform",
    "TrexTrackForceTransform",
    "prepare_track_windows",
]