groot/vla/experiment/trex_eval_utils.py

"""Shared T-Rex causal video inference helpers (training callback + offline scripts)."""

from __future__ import annotations

import glob
import os
from typing import TYPE_CHECKING

import cv2
import numpy as np
import pyarrow.parquet as pq
import torch
from einops import rearrange
from groot.vla.model.n1_5.sim_policy import unsqueeze_dict_values

if TYPE_CHECKING:
    from groot.vla.data.transform import ComposedModalityTransform

VIDEO_KEYS = [
    "video.head_left",
    "video.left_wrist",
    "video.right_wrist",
]
STATE_KEYS = {
    "state.left_arm": (0, 7),
    "state.left_hand": (7, 29),
    "state.right_arm": (29, 36),
    "state.right_hand": (36, 58),
}
VIDEO_FOLDERS = {
    "video.head_left": "observation.images.head_left",
    "video.left_wrist": "observation.images.left_wrist",
    "video.right_wrist": "observation.images.right_wrist",
}

RELATIVE_OFFSETS = [-23, -16, -8, 0]
ACTION_HORIZON = 24


def unwrap_model(model: torch.nn.Module) -> torch.nn.Module:
    return model.module if hasattr(model, "module") else model


def get_expected_video_resolution(model: torch.nn.Module) -> tuple[int, int]:
    cfg = model.action_head.config
    target_h = getattr(cfg, "target_video_height", None)
    target_w = getattr(cfg, "target_video_width", None)
    if target_h is not None and target_w is not None:
        return int(target_h), int(target_w)
    return 160, 320


def resize_frames(frames: np.ndarray, target_h: int, target_w: int) -> np.ndarray:
    if frames.ndim == 3:
        if (frames.shape[0], frames.shape[1]) == (target_h, target_w):
            return frames
        return cv2.resize(frames, (target_w, target_h), interpolation=cv2.INTER_LINEAR)
    return np.stack(
        [cv2.resize(f, (target_w, target_h), interpolation=cv2.INTER_LINEAR) for f in frames],
        axis=0,
    )


class TrexEpisode:
    """One T-Rex episode from LeRobot v2 layout."""

    def __init__(self, dataset_root: str, episode_index: int):
        pq_path = os.path.join(
            dataset_root,
            "data",
            f"chunk-{episode_index // 1000:03d}",
            f"episode_{episode_index:06d}.parquet",
        )
        if not os.path.isfile(pq_path):
            raise FileNotFoundError(pq_path)
        self.table = pq.read_table(pq_path)
        self.length = self.table.num_rows
        self.episode_index = episode_index
        self.root = dataset_root
        self._frame_cache: dict[str, dict[int, np.ndarray]] = {}

        self.video_dirs: dict[str, str] = {}
        for key, folder in VIDEO_FOLDERS.items():
            pattern = os.path.join(
                dataset_root,
                "videos",
                "**",
                folder,
                f"episode_{episode_index:06d}.mp4",
            )
            hits = sorted(glob.glob(pattern, recursive=True))
            if not hits:
                raise FileNotFoundError(f"No video for {key}: {pattern}")
            self.video_dirs[key] = hits[0]
            self._frame_cache[key] = {}

    def get_task(self, row: int) -> str:
        try:
            return str(self.table.column("annotation.task")[row].as_py())
        except Exception:
            return ""

    def get_state(self, row: int) -> np.ndarray:
        return np.array(self.table.column("observation.state")[row].as_py(), dtype=np.float64)

    def get_frames(
        self, rows: list[int] | np.ndarray, video_key: str
    ) -> np.ndarray:
        indices = np.asarray(rows, dtype=np.int64).reshape(-1)
        if video_key not in self.video_dirs:
            raise KeyError(video_key)
        if indices.size == 0:
            raise ValueError("rows must not be empty")
        if indices.min() < 0 or indices.max() >= self.length:
            raise IndexError("video frame index outside episode")
        cache = self._frame_cache[video_key]
        missing = sorted(set(indices.tolist()) - set(cache))
        if missing:
            cap = cv2.VideoCapture(self.video_dirs[video_key])
            if not cap.isOpened():
                raise RuntimeError(f"Failed to open {self.video_dirs[video_key]}")
            previous = -2
            try:
                for row in missing:
                    if row != previous + 1:
                        cap.set(cv2.CAP_PROP_POS_FRAMES, row)
                    ok, frame = cap.read()
                    if not ok:
                        raise RuntimeError(
                            f"Failed frame {row} from {self.video_dirs[video_key]}"
                        )
                    cache[row] = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                    previous = row
            finally:
                cap.release()
        return np.stack([cache[int(row)] for row in indices])

    def get_frame(self, row: int, video_key: str) -> np.ndarray:
        return self.get_frames([row], video_key)[0]


def build_chunk_schedule(total_frames: int, num_chunks: int) -> list[list[int]]:
    chunks: list[list[int]] = []
    anchor = 0
    for i in range(num_chunks):
        if i == 0:
            indices = [0]
        else:
            indices = [max(anchor + off, 0) for off in RELATIVE_OFFSETS]
        if indices[-1] >= total_frames:
            break
        chunks.append(indices)
        anchor += ACTION_HORIZON
    return chunks


def build_obs(
    episode: TrexEpisode,
    frame_indices: list[int],
    prompt: str,
    video_height: int,
    video_width: int,
    *,
    pre_resize: bool = True,
) -> dict:
    obs: dict = {}
    anchor = frame_indices[-1]
    state = episode.get_state(anchor)

    for key in VIDEO_KEYS:
        frames = episode.get_frames(frame_indices, key)
        frames = frames.astype(np.uint8)
        if pre_resize:
            frames = resize_frames(frames, video_height, video_width)
        if len(frame_indices) == 1:
            obs[key] = frames[0:1]  # (1, H, W, 3); transforms require a time dimension
        else:
            obs[key] = frames  # (T, H, W, 3)

    for key, (start, end) in STATE_KEYS.items():
        obs[key] = state[start:end].reshape(1, -1).astype(np.float64)

    obs["annotation.task"] = prompt
    return obs


def reset_action_head_causal_state(action_head) -> None:
    action_head.current_start_frame = 0
    action_head.kv_cache1 = None
    action_head.kv_cache_neg = None
    if hasattr(action_head, "last_language"):
        action_head.last_language = None


def decode_video_latents(model: torch.nn.Module, video_chunks: list[torch.Tensor]) -> np.ndarray:
    if not video_chunks:
        raise ValueError("No video chunks to decode")
    cat = torch.cat(video_chunks, dim=2)
    ah = unwrap_model(model).action_head
    frames = ah.vae.decode(
        cat,
        tiled=ah.tiled,
        tile_size=(ah.tile_size_height, ah.tile_size_width),
        tile_stride=(ah.tile_stride_height, ah.tile_stride_width),
    )
    frames = rearrange(frames, "B C T H W -> B T H W C")[0]
    return ((frames.float() + 1) * 127.5).clip(0, 255).cpu().numpy().astype(np.uint8)


def run_trex_causal_video_prediction(
    model: torch.nn.Module,
    eval_transform: ComposedModalityTransform,
    episode: TrexEpisode,
    *,
    prompt: str,
    num_chunks: int,
    eval_bf16: bool = True,
    pre_resize_video: bool = False,
) -> np.ndarray:
    """Run causal video reconstruction and return decoded uint8 frames (T, H, W, 3)."""
    unwrapped = unwrap_model(model)
    video_height, video_width = get_expected_video_resolution(unwrapped)
    schedule = build_chunk_schedule(episode.length, num_chunks)
    if not schedule:
        raise RuntimeError(f"Episode {episode.episode_index} is too short for causal inference")

    reset_action_head_causal_state(unwrapped.action_head)
    video_chunks: list[torch.Tensor] = []

    for frame_indices in schedule:
        obs = build_obs(
            episode,
            frame_indices,
            prompt,
            video_height,
            video_width,
            pre_resize=pre_resize_video,
        )
        # Match GrootSimPolicy inference: add batch dim before eval_transform so
        # DreamTransform.apply_batch splits state/video consistently.
        obs = unsqueeze_dict_values(obs)
        normalized_input = eval_transform(obs)
        for key, value in list(normalized_input.items()):
            if torch.is_tensor(value) and value.dtype == torch.float32 and eval_bf16:
                normalized_input[key] = value.to(dtype=torch.bfloat16)

        with torch.inference_mode():
            model_pred = unwrapped.lazy_joint_video_action_causal(normalized_input)
        video_chunks.append(model_pred["video_pred"])

    return decode_video_latents(model, video_chunks)