scripts/eval/trex_dexterity_eval.py

"""Dexterity-focused comparison: continuous co-denoised force (joint) vs
discrete VQ-conditioned force (force_only).

Per checkpoint, on identical anchors/seeds:
  * open-loop hand/arm MSE, finger-velocity MSE, per-step error growth;
  * fine-structure (2nd-difference) energy ratio vs GT — <1 means the hand
    trajectory is over-smoothed, losing dexterous micro-movement;
  * tactile sensitivity: |refined action delta| (hand dims) when tactile
    inputs are zeroed;
  * figures: finger-joint trajectory overlays, per-step error curves,
    force continuity (GT vs VQ-discretized view vs predicted future force).

Usage:
  python scripts/eval/trex_dexterity_eval.py --dataset-root data/trex_small_force \
     --run force_only=checkpoints/small_forceonly_6k/checkpoint-6000 \
     --run joint=checkpoints/small_joint_6k/checkpoint-6000 \
     --out-dir ablation_runs/dexterity_small
"""

from __future__ import annotations

import argparse
import gc
import json
import sys
from pathlib import Path

import numpy as np
import torch

sys.path.insert(0, str(Path(__file__).parent))
from trex_ablation_eval import ACTION_HORIZON, ChunkBuilder, load_pipeline  # noqa: E402

from groot.vla.model.n1_5.sim_policy import unsqueeze_dict_values  # noqa: E402
from groot.vla.model.trex_track_force.vla import TrexTrackForceVLA  # noqa: E402

ARM_DIMS = list(range(0, 9)) + list(range(31, 40))
HAND_DIMS = list(range(9, 31)) + list(range(40, 62))
GT_COLOR, OLD_COLOR, NEW_COLOR = "#878c98", "#3c7ec0", "#b54e7f"


def normalize_gt(builder, delta):
    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)
    return np.where(scale == 0, delta, gt)


def collect(policy, transform_eval, builder, anchors, prompt, device):
    """Open-loop predictions on each anchor; returns stacked arrays."""

    preds, gts, forces, gt_slot_forces = [], [], [], []
    deltas = []
    for index, anchor in enumerate(anchors):
        raw = builder.build(anchor, blocks=1, prompt=prompt,
                            with_deform=False, history_only_video=True)
        delta = np.asarray(raw.pop("action.eef62"), dtype=np.float32)
        gt = normalize_gt(builder, delta)[: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))
            batch[key] = value
        with torch.inference_mode():
            result = policy.sample(batch, seed=900 + index,
                                   run_force_refinement=True)
        preds.append(result["action_pred"][0, :, :62].float().cpu().numpy())
        gts.append(gt)
        gt_slot_forces.append(np.asarray(raw["current_force"])[0])   # [4,10,6]
        if "force_pred" in result:
            forces.append(result["force_pred"][0].float().cpu().numpy())
        deltas.append(delta[:ACTION_HORIZON])
    return (np.stack(preds), np.stack(gts), np.stack(gt_slot_forces),
            np.stack(forces) if forces else None)


def tactile_delta(policy, transform_eval, builder, anchor, prompt, device):
    """|refined action(real tactile) − refined action(zeroed tactile)|, hand dims."""

    raw = builder.build(anchor, blocks=1, prompt=prompt,
                        with_deform=False, history_only_video=True)
    raw.pop("action.eef62")
    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))
        batch[key] = value
    with torch.inference_mode():
        real = policy.sample(batch, seed=901, run_force_refinement=True)
        zeroed = dict(batch)
        zeroed["current_force"] = torch.zeros_like(batch["current_force"])
        zeroed["tactile_force_history"] = torch.zeros_like(
            batch["tactile_force_history"])
        blank = policy.sample(zeroed, seed=901, run_force_refinement=True)
    delta = (real["action_pred"] - blank["action_pred"])[0, :, :62].float()
    return float(delta[:, HAND_DIMS].abs().mean())


def vq_view_of_force(policy, history):
    """What the discrete pathway 'sees': VQ encode->decode reconstruction."""

    vq = policy.force_transformer.tactile_vq
    with torch.inference_mode():
        out = vq(torch.as_tensor(history, dtype=policy.dtype,
                                 device=policy.device)[None])
    recon = out.reconstruction
    return None if recon is None else recon[0].float().cpu().numpy()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset-root", required=True)
    parser.add_argument("--run", action="append", required=True)
    parser.add_argument("--out-dir", required=True)
    parser.add_argument("--anchors", type=int, default=24)
    args = parser.parse_args()

    out_dir = Path(args.out_dir); out_dir.mkdir(parents=True, exist_ok=True)
    device = torch.device("cuda")
    builder = ChunkBuilder(args.dataset_root)
    with open(Path(args.dataset_root) / "meta" / "tasks.jsonl") as handle:
        prompt = json.loads(handle.readline())["task"]
    anchors_all = builder.valid_anchor_times(1)
    anchors = [anchors_all[int(i)]
               for i in np.linspace(0, len(anchors_all) - 1, args.anchors)]

    report, curves, trajs = {}, {}, {}
    force_fig = {}
    for spec in args.run:
        name, _, ckpt = spec.partition("=")
        print(f"=== {name}: {ckpt} ===", flush=True)
        model = TrexTrackForceVLA.load_lora(str(ckpt)).to(
            device=device, dtype=torch.bfloat16)
        model.eval()
        policy = model.action_head
        transform_eval, _ = load_pipeline(Path(ckpt), training=False)

        preds, gts, gt_slots, force_pred = collect(
            policy, transform_eval, builder, anchors, prompt, device)
        err = (preds - gts) ** 2
        vel_p, vel_g = np.diff(preds, axis=1), np.diff(gts, axis=1)
        acc_p, acc_g = np.diff(preds, 2, axis=1), np.diff(gts, 2, axis=1)
        record = {
            "checkpoint": ckpt,
            "mse": float(err.mean()),
            "mse_hand": float(err[:, :, HAND_DIMS].mean()),
            "mse_arm": float(err[:, :, ARM_DIMS].mean()),
            "vel_mse_hand": float(
                ((vel_p - vel_g)[:, :, HAND_DIMS] ** 2).mean()),
            "fine_energy_ratio_hand": float(
                (acc_p[:, :, HAND_DIMS] ** 2).mean()
                / max((acc_g[:, :, HAND_DIMS] ** 2).mean(), 1e-9)),
            "mse_hand_per_anchor": err[:, :, HAND_DIMS].mean(
                axis=(1, 2)).tolist(),
        }
        curves[name] = err[:, :, HAND_DIMS].mean(axis=(0, 2)).tolist()
        # representative anchor: largest GT slot-force variance
        rep = int(np.argmax(gt_slots.var(axis=(1, 2, 3))))
        moving = np.argsort(gts[rep].ptp(axis=0)[HAND_DIMS])[-4:]
        trajs[name] = {
            "anchor_index": rep,
            "joints": [int(HAND_DIMS[j]) for j in moving],
            "pred": preds[rep][:, [HAND_DIMS[j] for j in moving]].tolist(),
            "gt": gts[rep][:, [HAND_DIMS[j] for j in moving]].tolist(),
        }
        raw_rep = builder.build(anchors[rep], blocks=1, prompt=prompt,
                                with_deform=False, history_only_video=True)
        history = np.asarray(raw_rep["force_history"])[0, 0]     # [16,10,6]
        vq_recon = vq_view_of_force(policy, history)
        finger = int(np.argmax(np.abs(gt_slots[rep]).sum(axis=(0, 2))))
        force_fig[name] = {
            "finger": finger,
            "gt_history": history[:, finger, :3].tolist(),
            "gt_slots": gt_slots[rep][:, finger, :3].tolist(),
            "vq_recon": (None if vq_recon is None
                         else vq_recon[:, finger, :3].tolist()),
            "force_pred": (None if force_pred is None
                           else force_pred[rep][:, finger, :3].tolist()),
        }
        record["tactile_delta_hand"] = tactile_delta(
            policy, transform_eval, builder, anchors[rep], prompt, device)
        report[name] = record
        del model, policy
        gc.collect(); torch.cuda.empty_cache()

    json.dump({"report": report, "curves": curves, "trajectories": trajs,
               "force_fig": force_fig},
              open(out_dir / "dexterity.json", "w"), indent=1)

    # ---------------- figures ----------------
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    names = list(report)
    palette = [OLD_COLOR, NEW_COLOR, "#6f66bd", "#c26a1e"]
    colors = {n: palette[i % len(palette)] for i, n in enumerate(names)}

    fig, axes = plt.subplots(1, 4, figsize=(16, 3.4))
    first = names[0]
    for k in range(4):
        ax = axes[k]
        joint_id = trajs[first]["joints"][k]
        ax.plot(np.array(trajs[first]["gt"])[:, k], color=GT_COLOR,
                ls="--", lw=2, label="GT")
        for name in names:
            ax.plot(np.array(trajs[name]["pred"])[:, k],
                    color=colors[name], lw=1.8, label=name)
        ax.set_title(f"hand dim {joint_id}", fontsize=10)
        ax.set_xlabel("step (20 Hz)")
    axes[0].set_ylabel("normalized q")
    axes[0].legend(fontsize=8)
    fig.suptitle("Finger-joint open-loop trajectories (most-moving hand dims)")
    fig.tight_layout(); fig.savefig(out_dir / "fig_finger_trajectories.png",
                                    dpi=150); plt.close(fig)

    fig, ax = plt.subplots(figsize=(6, 3.6))
    for name in names:
        ax.plot(curves[name], color=colors[name], lw=2, marker="o",
                ms=3, label=name)
    ax.set_xlabel("chunk step"); ax.set_ylabel("hand MSE")
    ax.set_title("Per-step hand error growth (mean over anchors)")
    ax.legend(); fig.tight_layout()
    fig.savefig(out_dir / "fig_per_step_error.png", dpi=150); plt.close(fig)

    fig, axes = plt.subplots(1, len(names), figsize=(6.5 * len(names), 3.6),
                             squeeze=False)
    slot_x = 16 + np.arange(4) * 1.0
    for ax, name in zip(axes[0], names):
        ff = force_fig[name]
        gt_hist = np.array(ff["gt_history"])
        ax.plot(np.arange(16), gt_hist[:, 2], color=GT_COLOR, lw=2,
                label="GT force (5 Hz history)")
        ax.plot(slot_x, np.array(ff["gt_slots"])[:, 2], color=GT_COLOR,
                ls="--", marker="o", label="GT future slots")
        if ff["vq_recon"] is not None:
            ax.step(np.arange(16), np.array(ff["vq_recon"])[:, 2],
                    color=OLD_COLOR, where="mid", lw=1.6,
                    label="VQ-discretized view")
        if ff["force_pred"] is not None:
            ax.plot(slot_x, np.array(ff["force_pred"])[:, 2],
                    color=NEW_COLOR, marker="s", lw=1.8,
                    label="predicted future force")
        ax.axvline(15.5, color="#cccccc", lw=1)
        ax.set_title(f"{name} · finger {ff['finger']} · Fz", fontsize=10)
        ax.set_xlabel("5 Hz frame (past | future)")
        ax.legend(fontsize=7)
    axes[0][0].set_ylabel("normalized force")
    fig.suptitle("Continuous force vs discrete VQ bottleneck")
    fig.tight_layout(); fig.savefig(out_dir / "fig_force_continuity.png",
                                    dpi=150); plt.close(fig)

    metrics = ["mse_hand", "vel_mse_hand", "fine_energy_ratio_hand",
               "tactile_delta_hand"]
    fig, axes = plt.subplots(1, len(metrics), figsize=(3.2 * len(metrics), 3.2))
    for ax, metric in zip(axes, metrics):
        vals = [report[n][metric] for n in names]
        ax.bar(range(len(names)), vals,
               color=[colors[n] for n in names], width=0.55)
        ax.set_xticks(range(len(names)), names, fontsize=8)
        ax.set_title(metric, fontsize=9)
        if metric == "fine_energy_ratio_hand":
            ax.axhline(1.0, color=GT_COLOR, ls="--", lw=1)
    fig.suptitle("Dexterity metrics (hand dims)")
    fig.tight_layout(); fig.savefig(out_dir / "fig_metrics.png", dpi=150)
    plt.close(fig)
    print("wrote", out_dir)


if __name__ == "__main__":
    main()