"""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()