scripts/eval/trex_ablation_eval.py

"""Controlled ablation evaluation for the T-Rex Track-Force cascade.

For each named checkpoint, on identical data:

* fixed-tau forward losses (action / dynamics / track / force flow MSE);
* open-loop chunk reconstruction: normalized 62-D action MSE of ``sample()``
  against the ground-truth delta-base chunk (cascade and coarse-only);
* Stage-1 self-attention mass of action/obs queries over key token groups;
* Stage-2 force-transformer attention mass of action queries over
  F6 / VQ-history / deform / coarse-memory tokens;
* tactile sensitivity: |refine(real tactile) - refine(tactile masked)|.

Usage:
    python scripts/eval/trex_ablation_eval.py \
        --dataset-root data/trex_mini_force \
        --run full=checkpoints/ablate_full_3k/checkpoint-3000 \
        --out ablation_eval.json
"""

from __future__ import annotations

import argparse
import gc
import json
from pathlib import Path

import numpy as np
import torch
from hydra.utils import instantiate
from omegaconf import OmegaConf

from groot.vla.data.schema import DatasetMetadata, EmbodimentTag
from groot.vla.experiment.trex_eval_utils import TrexEpisode
from groot.vla.model.trex_track_force.attention import TokenType
from groot.vla.model.trex_track_force.dataset import (
    DEFORM_VIDEO_KEYS,
    eef62_delta_base,
    nearest_timestamp_indices,
    uniform_target_times,
)
from groot.vla.model.trex_track_force.force import (
    ACTION_HORIZON,
    FORCE_HISTORY_FRAMES,
    FORCE_OFFSETS,
    euler_flow_step,
    pad_action_62_to_64,
)
from groot.vla.model.trex_track_force.runtime import TrexRuntimeStatistics
from groot.vla.model.trex_track_force.track import TRACK_HORIZON
from groot.vla.model.n1_5.sim_policy import unsqueeze_dict_values
from groot.vla.model.trex_track_force.vla import TrexTrackForceVLA

ACTION_RATE_HZ = 20.0
TACTILE_RATE_HZ = 5.0
VIDEO_RATE_HZ = 10.0
AR_BLOCKS = 4
VIDEO_FRAMES_PER_BLOCK = 8
VIDEO_KEYS = ("video.head_left", "video.left_wrist", "video.right_wrist")


def _column(episode: TrexEpisode, name: str, dtype=np.float32) -> np.ndarray:
    values = episode.table.column(name).to_numpy(zero_copy_only=False)
    values = np.asarray(values)
    while values.dtype == object:
        values = np.stack([np.stack(row) for row in values])
    return values.astype(dtype)


def _sample(timestamps, anchor, offsets, rate):
    return nearest_timestamp_indices(
        timestamps, uniform_target_times(anchor, offsets, rate)
    )


def _read_frames(episode: TrexEpisode, key: str, indices: np.ndarray) -> np.ndarray:
    import decord

    reader = decord.VideoReader(episode.video_dirs[key], num_threads=1)
    return reader.get_batch([int(i) for i in indices]).asnumpy().astype(np.uint8)


class ChunkBuilder:
    """Builds training-format (K-block) raw samples from one episode."""

    def __init__(self, dataset_root: str, episode_index: int = 0) -> None:
        self.root = dataset_root
        self.episode = TrexEpisode(dataset_root, episode_index)
        self.stats = TrexRuntimeStatistics.from_dataset(dataset_root)
        self.timestamps = _column(self.episode, "timestamp", np.float64).reshape(-1)
        self.state = _column(self.episode, "observation.state_eef62")
        self.action_abs = _column(self.episode, "action.eef62_absolute")
        self.track_xy = _column(self.episode, "observation.track_xy")
        self.track_vis = _column(self.episode, "observation.track_visibility")
        self.force = _column(self.episode, "observation.tactile_force").reshape(
            -1, 10, 6
        )
        self._deform: np.ndarray | None = None

    def valid_anchor_times(self, blocks: int) -> list[float]:
        future = max(
            blocks * ACTION_HORIZON / ACTION_RATE_HZ,
            blocks * VIDEO_FRAMES_PER_BLOCK / VIDEO_RATE_HZ,
        )
        grid = np.arange(
            self.timestamps[0], self.timestamps[-1] + 1e-9, 1.0 / ACTION_RATE_HZ
        )
        return [float(t) for t in grid if t + future <= self.timestamps[-1]]

    def deform_frames(self, size: int = 96) -> np.ndarray:
        if self._deform is None:
            import cv2
            import decord

            streams = []
            for key in DEFORM_VIDEO_KEYS:
                path = (
                    Path(self.root) / "videos" / "chunk-000"
                    / f"observation.images.{key}"
                    / f"episode_{self.episode.episode_index:06d}.mp4"
                )
                reader = decord.VideoReader(str(path), num_threads=1)
                frames = reader.get_batch(range(len(reader))).asnumpy()
                frames = np.stack(
                    [cv2.resize(f, (size, size), interpolation=cv2.INTER_AREA)
                     for f in frames]
                )
                streams.append(frames.astype(np.uint8))
            self._deform = np.stack(streams, axis=1)
        return self._deform

    def _norm_force(self, selection) -> np.ndarray:
        values = self.stats.normalize_force(self.force[selection.indices])
        values[selection.padding_mask] = 0.0
        return values

    def build(self, anchor: float, *, blocks: int, prompt: str,
              with_deform: bool, history_only_video: bool = False) -> dict:
        ts = self.timestamps
        block_anchors = [
            anchor + b * ACTION_HORIZON / ACTION_RATE_HZ for b in range(blocks)
        ]
        action_sel = _sample(ts, anchor, range(blocks * ACTION_HORIZON), ACTION_RATE_HZ)
        state_sel = _sample(
            ts, anchor, range(0, blocks * ACTION_HORIZON, ACTION_HORIZON),
            ACTION_RATE_HZ,
        )
        reference = self.state[state_sel.indices]
        absolute = self.action_abs[action_sel.indices].reshape(
            blocks, ACTION_HORIZON, 62
        )
        delta = np.stack(
            [eef62_delta_base(reference[b], absolute[b]) for b in range(blocks)]
        ).reshape(blocks * ACTION_HORIZON, 62)

        past_sels = [
            _sample(ts, b, range(-(FORCE_HISTORY_FRAMES - 1), 1), ACTION_RATE_HZ)
            for b in block_anchors
        ]
        future_sels = [
            _sample(ts, b, range(TRACK_HORIZON), ACTION_RATE_HZ)
            for b in block_anchors
        ]
        force_sels = [
            [
                _sample(ts, b + off / ACTION_RATE_HZ,
                        range(-(FORCE_HISTORY_FRAMES - 1), 1), TACTILE_RATE_HZ)
                for off in FORCE_OFFSETS
            ]
            for b in block_anchors
        ]
        force_history = np.stack(
            [[self._norm_force(sel) for sel in block_sel] for block_sel in force_sels]
        )

        video_hist = _sample(ts, anchor, range(1), VIDEO_RATE_HZ)
        if history_only_video:
            video_indices = video_hist.indices
        else:
            video_future = _sample(
                ts, anchor, range(1, blocks * VIDEO_FRAMES_PER_BLOCK + 1),
                VIDEO_RATE_HZ,
            )
            video_indices = np.concatenate(
                (video_hist.indices, video_future.indices)
            )

        raw: dict[str, object] = {
            key: _read_frames(self.episode, key, video_indices)
            for key in VIDEO_KEYS
        }
        raw.update(
            {
                "state.eef62": reference.astype(np.float32),
                "action.eef62": delta.astype(np.float32),
                "track_past_xy": np.stack(
                    [self.track_xy[s.indices] for s in past_sels]
                ),
                "track_past_visibility": np.stack(
                    [self.track_vis[s.indices] * (~s.padding_mask[:, None])
                     for s in past_sels]
                ),
                "track_future_xy": np.stack(
                    [self.track_xy[s.indices] for s in future_sels]
                ),
                "track_future_visibility": np.stack(
                    [self.track_vis[s.indices] for s in future_sels]
                ),
                "current_force": force_history[:, :, -1],
                "force_history": force_history,
                "force_history_padding_mask": np.stack(
                    [[sel.padding_mask for sel in block_sel]
                     for block_sel in force_sels]
                ),
                "annotation.task": prompt,
            }
        )
        if with_deform:
            deform = self.deform_frames()
            refresh = np.stack(
                [[int(sel.indices[-1]) for sel in block_sel]
                 for block_sel in force_sels]
            )
            raw["deform_current"] = deform[refresh]
        return raw


def load_pipeline(checkpoint: Path, *, training: bool):
    cfg_dir = checkpoint / "experiment_cfg"
    if not cfg_dir.exists():
        cfg_dir = checkpoint.parent / "experiment_cfg"
    cfg = OmegaConf.load(cfg_dir / "conf.yaml")
    with open(cfg_dir / "metadata.json", "r", encoding="utf-8") as handle:
        metadata = DatasetMetadata.model_validate(
            json.load(handle)[EmbodimentTag.TREX.value]
        )
    transform = instantiate(cfg.transforms["trex"])
    transform.set_metadata(metadata)
    transform.train() if training else transform.eval()
    collator = instantiate(cfg.data_collator)
    return transform, collator


def to_batch(sample: dict, collator, device, dtype) -> dict:
    batch = collator([sample])
    out = {}
    for key, value in batch.items():
        if torch.is_tensor(value):
            value = (
                value.to(device=device, dtype=dtype)
                if value.is_floating_point()
                else value.to(device=device)
            )
        out[key] = value
    return out


def stage1_attention_masses(policy, batch, *, layers, tau_value):
    from groot.vla.model.trex_track_force import blocks as block_module

    records: dict[int, dict] = {}

    def make_patched(layer_index, module):
        def patched(x, *, layout, rope_frequencies, allow_matrix=None, **_):
            batch_size, length = x.shape[:2]
            heads, head_dim = module.num_heads, module.head_dim
            query = module.norm_q(module.q(x)).view(batch_size, length, heads, head_dim)
            key = module.norm_k(module.k(x)).view(batch_size, length, heads, head_dim)
            value = module.v(x).view(batch_size, length, heads, head_dim)
            query = block_module.apply_multimodal_rope(
                query, rope_frequencies
            ).type_as(value)
            key = block_module.apply_multimodal_rope(
                key, rope_frequencies
            ).type_as(value)
            scores = torch.einsum("blhd,bmhd->bhlm", query, key) * head_dim**-0.5
            scores = scores.float().masked_fill(
                ~allow_matrix.view(1, 1, length, length), float("-inf")
            )
            attention = scores.softmax(dim=-1)
            token_type, _ = layout.token_metadata(device=x.device)
            layer_record = {}
            for query_group, query_type in (
                ("action", TokenType.ACTION), ("obs", TokenType.OBS),
            ):
                rows = (token_type == int(query_type)).nonzero(as_tuple=True)[0]
                row_attention = attention[:, :, rows]
                masses = {}
                for name, key_type in (
                    ("cond_obs", TokenType.CONDITIONING_OBS),
                    ("obs", TokenType.OBS),
                    ("action", TokenType.ACTION),
                    ("state", TokenType.STATE),
                    ("track_past", TokenType.TRACK_PAST),
                    ("track_future", TokenType.TRACK_FUTURE),
                ):
                    cols = (token_type == int(key_type)).nonzero(as_tuple=True)[0]
                    masses[name] = (
                        float(row_attention[..., cols].sum(dim=-1).mean().item())
                        if cols.numel()
                        else 0.0
                    )
                layer_record[query_group] = masses
            records[layer_index] = layer_record
            output = torch.einsum(
                "bhlm,bmhd->blhd", attention.to(value.dtype), value
            ).reshape(batch_size, length, module.dim)
            return module.o(output), None

        return patched

    originals = {}
    for layer_index in layers:
        module = policy.model.blocks[layer_index].self_attn
        originals[layer_index] = module.forward
        module.forward = make_patched(layer_index, module)
    try:
        tau = torch.full(
            (1, AR_BLOCKS), tau_value, device=policy.device, dtype=policy.dtype
        )
        with torch.inference_mode():
            policy.forward_core(batch, tau=tau)
    finally:
        for layer_index, forward in originals.items():
            policy.model.blocks[layer_index].self_attn.forward = forward
    return records


class Stage2Recorder:
    """Wrap each force-transformer layer to record action-query attention."""

    def __init__(self, force):
        self.force = force
        self.groups = {"action": force.action_horizon,
                       "f6": force.force_sensor_count,
                       "vq": force.force_sensor_count}
        if force.use_deform_tactile:
            self.groups["deform"] = force.force_sensor_count
        self.records: list[dict[str, float]] = []
        self._originals = []

    def __enter__(self):
        for layer in self.force.transformer.layers:
            self._originals.append((layer, layer.forward))
            layer.forward = self._make(layer)
        return self

    def __exit__(self, *exc):
        for layer, forward in self._originals:
            layer.forward = forward

    def _make(self, layer):
        groups = self.groups
        records = self.records

        def forward(src, src_mask=None, src_key_padding_mask=None, is_causal=False):
            x = src
            normed = layer.norm1(x)
            attn_out, weights = layer.self_attn(
                normed, normed, normed,
                attn_mask=src_mask,
                key_padding_mask=src_key_padding_mask,
                need_weights=True,
                average_attn_weights=True,
            )
            action_rows = weights[:, : groups["action"]]
            cursor, masses = 0, {}
            for name, size in groups.items():
                masses[name] = float(
                    action_rows[..., cursor:cursor + size].sum(-1).mean().item()
                )
                cursor += size
            masses["memory"] = float(action_rows[..., cursor:].sum(-1).mean().item())
            records.append(masses)
            x = x + layer.dropout1(attn_out)
            x = x + layer._ff_block(layer.norm2(x))
            return x

        return forward


@torch.inference_mode()
def refine_with_keep(policy, state_batch, refinement, *, keep: bool,
                     deform_images=None):
    """Mirror refine_action_suffix but force the tactile keep mask."""

    force = policy.force_transformer
    action = pad_action_62_to_64(refinement["coarse_action"]).clone()
    action[..., policy.config.physical_action_dim:] = 0
    keep_mask = torch.full(
        (action.shape[0],), 1.0 if keep else 0.0,
        device=policy.device, dtype=policy.dtype,
    )
    for step in policy.schedule.iter_steps("force"):
        flow = force(
            action,
            step.tau,
            refinement["current_force"],
            refinement["history"],
            coarse_memory=refinement["memory"],
            update_offset=0,
            tactile_keep_mask=keep_mask,
            tactile_history_valid_mask=refinement["history_valid"],
            deform_images=deform_images,
        )
        action = euler_flow_step(action, flow, step.tau, step.tau_next)
        action[..., policy.config.physical_action_dim:] = 0
    return action


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset-root", required=True)
    parser.add_argument("--run", action="append", required=True)
    parser.add_argument("--out", required=True)
    parser.add_argument("--forward-anchors", type=int, default=6)
    parser.add_argument("--openloop-anchors", type=int, default=16)
    parser.add_argument("--tau-grid", default="0.8,0.6,0.4,0.2,0.05")
    parser.add_argument("--tau-repeats", type=int, default=4)
    parser.add_argument("--attn-layers", default="0,14,29")
    parser.add_argument("--prompt", default=None)
    args = parser.parse_args()

    device = torch.device("cuda")
    builder = ChunkBuilder(args.dataset_root)
    prompt = args.prompt
    if prompt is None:
        with open(Path(args.dataset_root) / "meta" / "tasks.jsonl") as handle:
            prompt = json.loads(handle.readline())["task"]

    anchors_k4 = builder.valid_anchor_times(AR_BLOCKS)
    forward_anchor_times = [
        anchors_k4[int(i)]
        for i in np.linspace(0, len(anchors_k4) - 1, args.forward_anchors)
    ]
    anchors_k1 = builder.valid_anchor_times(1)
    openloop_anchor_times = [
        anchors_k1[int(i)]
        for i in np.linspace(0, len(anchors_k1) - 1, args.openloop_anchors)
    ]
    tau_grid = [float(v) for v in args.tau_grid.split(",")]
    attn_layers = [int(v) for v in args.attn_layers.split(",")]

    results: dict[str, dict] = {}
    for spec in args.run:
        name, _, checkpoint = spec.partition("=")
        checkpoint_path = Path(checkpoint)
        print(f"=== {name}: {checkpoint_path} ===", flush=True)
        model = TrexTrackForceVLA.load_lora(str(checkpoint_path))
        model = model.to(device=device, dtype=torch.bfloat16)
        model.eval()
        policy = model.action_head
        use_force = policy.config.use_force
        use_deform = policy.config.use_deform_tactile
        transform, collator = load_pipeline(checkpoint_path, training=True)
        transform_eval, _ = load_pipeline(checkpoint_path, training=False)

        record: dict[str, object] = {
            "checkpoint": str(checkpoint_path),
            "use_track": policy.config.use_track,
            "use_force": use_force,
            "use_deform_tactile": use_deform,
        }

        # ---------------- controlled forward losses ----------------------
        forward_batches = []
        for anchor in forward_anchor_times:
            raw = builder.build(anchor, blocks=AR_BLOCKS, prompt=prompt,
                                with_deform=use_deform)
            forward_batches.append(
                to_batch(transform(dict(raw)), collator, device, torch.bfloat16)
            )
        tau_table = {}
        for tau_value in tau_grid:
            metrics: dict[str, list[float]] = {}
            for repeat in range(args.tau_repeats):
                for batch_index, batch in enumerate(forward_batches):
                    torch.manual_seed(100000 + repeat * 1000 + batch_index * 10)
                    tau = torch.full((1, AR_BLOCKS), tau_value,
                                     device=device, dtype=torch.bfloat16)
                    with torch.inference_mode():
                        out = policy.forward_core(batch, tau=tau)
                    for key in ("action_loss", "dynamics_loss",
                                "track_flow_loss", "force_loss"):
                        metrics.setdefault(key, []).append(float(out[key]))
            tau_table[f"{tau_value:g}"] = {
                key: [float(np.mean(v)), float(np.std(v))]
                for key, v in metrics.items()
            }
        record["forward_tau_losses"] = tau_table
        print(f"  forward losses done", flush=True)

        # ---------------- open-loop chunk reconstruction -----------------
        arm_dims = list(range(0, 9)) + list(range(31, 40))
        hand_dims = list(range(9, 31)) + list(range(40, 62))
        openloop: dict[str, list[float]] = {}
        split_metrics: dict[str, list[float]] = {}
        refinement_state = None
        for anchor_index, anchor in enumerate(openloop_anchor_times):
            raw = builder.build(anchor, blocks=1, prompt=prompt,
                                with_deform=use_deform, history_only_video=True)
            delta = np.asarray(raw.pop("action.eef62"), dtype=np.float32)
            scale = builder.stats.action_q99 - builder.stats.action_q01
            gt = np.clip(
                2.0 * (delta - builder.stats.action_q01)
                / np.where(scale == 0, 1.0, scale)
                - 1.0,
                -1.0,
                1.0,
            )
            gt = np.where(scale == 0, delta, gt)
            gt = torch.as_tensor(gt[:ACTION_HORIZON])
            collated = transform_eval(unsqueeze_dict_values(dict(raw)))
            batch = {}
            for key, value in collated.items():
                if torch.is_tensor(value):
                    value = (
                        value.to(device=device, dtype=torch.bfloat16)
                        if value.is_floating_point()
                        else value.to(device=device)
                    )
                batch[key] = value
            if use_deform and "deform_current" in batch:
                batch["deform_current"] = batch["deform_current"][:, :1, :1]
            modes = [("cascade", True)] if use_force else []
            modes.append(("coarse_only", False))
            for mode_name, refine in modes:
                with torch.inference_mode():
                    result = policy.sample(
                        batch, seed=500 + anchor_index,
                        run_force_refinement=refine,
                        return_refinement_state=refine and anchor_index == 0,
                    )
                if refine and anchor_index == 0:
                    current, history, valid = policy._extract_force_inputs(batch, 1)
                    refinement_state = {
                        "coarse_action": result["coarse_action_at_split"],
                        "memory": result["coarse_memory"],
                        "current_force": current[:, 0, 0].to(policy.device,
                                                             policy.dtype),
                        "history": history[:, 0, 0].to(policy.device),
                        "history_valid": (
                            None if valid is None
                            else valid[:, 0, 0].to(policy.device)
                        ),
                        "deform": (
                            policy._extract_deform_images(batch, 1, 1)
                            if use_deform else None
                        ),
                    }
                pred = result["action_pred"][0, :, :62].float().cpu()
                openloop.setdefault(mode_name, []).append(
                    float(((pred - gt) ** 2).mean())
                )
                if mode_name == "cascade" and "force_pred" in result:
                    gt_slots = batch["current_force"][0, 0].float().cpu()
                    pred_slots = result["force_pred"][0].float().cpu()
                    openloop.setdefault("force_future", []).append(
                        float(((pred_slots[1:] - gt_slots[1:]) ** 2).mean())
                    )
                prefix = "cascade" if mode_name == "cascade" else "coarse"
                split_metrics.setdefault(f"{prefix}_arm", []).append(
                    float(((pred[:, arm_dims] - gt[:, arm_dims]) ** 2).mean())
                )
                split_metrics.setdefault(f"{prefix}_hand", []).append(
                    float(((pred[:, hand_dims] - gt[:, hand_dims]) ** 2).mean())
                )
        record["openloop_action_mse"] = {
            key: [float(np.mean(v)), float(np.std(v)), len(v)]
            for key, v in openloop.items()
        }
        record["openloop_action_mse_raw"] = openloop
        record["openloop_action_mse_split"] = {
            key: float(np.mean(v)) for key, v in split_metrics.items()
        }
        print(f"  open-loop done", flush=True)

        # ---------------- attention masses --------------------------------
        record["stage1_attention"] = {
            str(layer): masses
            for layer, masses in stage1_attention_masses(
                policy, forward_batches[0], layers=attn_layers, tau_value=0.4
            ).items()
        }
        stage2_probe_ok = use_force and hasattr(
            policy.force_transformer, "transformer"
        )
        if stage2_probe_ok and refinement_state is not None:
            with Stage2Recorder(policy.force_transformer) as recorder:
                refine_with_keep(policy, None, refinement_state, keep=True,
                                 deform_images=refinement_state["deform"])
            per_layer = recorder.records[: len(
                policy.force_transformer.transformer.layers
            )]
            record["stage2_attention"] = per_layer

            refined_real = refine_with_keep(
                policy, None, refinement_state, keep=True,
                deform_images=refinement_state["deform"],
            )
            refined_masked = refine_with_keep(
                policy, None, refinement_state, keep=False,
                deform_images=refinement_state["deform"],
            )
            coarse = pad_action_62_to_64(refinement_state["coarse_action"])
            delta_tactile = (refined_real - refined_masked)[..., :62].float()
            delta_refine = (refined_real - coarse)[..., :62].float()
            record["tactile_sensitivity"] = {
                "mean_abs_delta_vs_masked": float(delta_tactile.abs().mean()),
                "max_abs_delta_vs_masked": float(delta_tactile.abs().max()),
                "mean_abs_refinement": float(delta_refine.abs().mean()),
            }
        print(f"  attention/sensitivity done", flush=True)

        results[name] = record
        del model, policy, forward_batches
        gc.collect()
        torch.cuda.empty_cache()

    with open(args.out, "w", encoding="utf-8") as handle:
        json.dump(results, handle, indent=1)
    print(f"wrote {args.out}")


if __name__ == "__main__":
    main()