"""Open-loop evaluation of FACT with the baseline ablation protocol. Same anchors, seeds and metric as ``the T-Rex single-episode ablation protocol (results/ablation_baseline)``: * fixed-tau forward losses on 6 K=4 anchors x 4 noise repeats (seeds 100000 + repeat*1000 + anchor*10); * open-loop chunk reconstruction on 16 K=1 anchors (seed 500+i): normalized 62-D MSE of the sampled chunk vs. the ground-truth delta-base chunk, for ``cascade`` (6 coarse + 4 tactile steps), ``coarse_only`` (coarse action read out at tau=0.4, exactly the baseline's coarse_only), and the extra ``coarse_to_zero`` (coarse model carried to tau=0), with arm/hand splits. * open-loop future-state MSE (same anchors). """ from __future__ import annotations import argparse import json from pathlib import Path import numpy as np import torch from fact import FactEpisode, FactPolicy from fact.data import AR_BLOCKS from fact.force import ACTION_HORIZON ARM_DIMS = list(range(0, 9)) + list(range(31, 40)) HAND_DIMS = list(range(9, 31)) + list(range(40, 62)) def to_batch(sample: dict[str, np.ndarray], device: torch.device) -> dict[str, torch.Tensor]: batch = {} for key, value in sample.items(): if key in ("action_delta_raw", "reference_state", "obs_frame_index"): continue batch[key] = torch.as_tensor(value)[None].to(device) return batch def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--checkpoint", required=True) parser.add_argument("--dataset-root", required=True) parser.add_argument("--episode", type=int, default=0) parser.add_argument("--cache", 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("--baseline-json", default=None) parser.add_argument("--bf16", action="store_true") parser.add_argument("--heldout-fraction", type=float, default=None, help="evaluate only on anchors starting after (1-f) of the episode (unseen by a --train-fraction (1-f) run)") args = parser.parse_args() device = torch.device("cuda") model = FactPolicy.load(args.checkpoint).to(device).eval() episode = FactEpisode(args.dataset_root, args.episode, latent_cache=args.cache) t_min = None if args.heldout_fraction is None else episode.split_time(1.0 - args.heldout_fraction) anchors_k4 = episode.valid_anchor_times(AR_BLOCKS, t_min=t_min) forward_anchors = [anchors_k4[int(i)] for i in np.linspace(0, len(anchors_k4) - 1, args.forward_anchors)] anchors_k1 = episode.valid_anchor_times(1, t_min=t_min) openloop_anchors = [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(",")] record: dict[str, object] = { "checkpoint": args.checkpoint, "heldout_fraction": args.heldout_fraction, "num_k1_anchors_available": len(anchors_k1), "config": json.loads(model.cfg.to_json()), "params_total": sum(p.numel() for p in model.parameters()), "params_coarse": sum(p.numel() for p in model.coarse.parameters()), "params_force": 0 if model.force is None else sum(p.numel() for p in model.force.parameters()), } # ---------------- fixed-tau forward losses ---------------- forward_batches = [to_batch(episode.build(a, blocks=AR_BLOCKS), device) for a in forward_anchors] tau_table = {} for tau_value in tau_grid: metrics: dict[str, list[float]] = {} for repeat in range(args.tau_repeats): for index, batch in enumerate(forward_batches): torch.manual_seed(100000 + repeat * 1000 + index * 10) tau = torch.full((1, AR_BLOCKS), tau_value, device=device) with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16, enabled=args.bf16): out = model.forward_core(batch, tau=tau) for key in ("action_loss", "state_loss", "force_loss"): metrics.setdefault(key, []).append(float(out[key])) tau_table[f"{tau_value:g}"] = {k: [float(np.mean(v)), float(np.std(v))] for k, v in metrics.items()} record["forward_tau_losses"] = tau_table print("forward losses done", flush=True) # ---------------- open-loop reconstruction ---------------- modes = ["coarse_only", "coarse_to_zero"] if model.force is not None: modes.insert(0, "cascade") openloop: dict[str, list[float]] = {} split: dict[str, list[float]] = {} fstate_mse: dict[str, list[float]] = {} for index, anchor in enumerate(openloop_anchors): sample = episode.build(anchor, blocks=1) gt = torch.as_tensor(sample["action"][0, :, :62]) gt_fstate = torch.as_tensor(sample["future_state"][0, :, :62]) batch = to_batch(sample, device) for mode in modes: with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16, enabled=args.bf16): result = model.sample(batch, seed=500 + index, mode=mode) pred = result["action_pred"][0, :, :62].float().cpu() openloop.setdefault(mode, []).append(float(((pred - gt) ** 2).mean())) openloop.setdefault(f"{mode}_first11", []).append(float(((pred[:11] - gt[:11]) ** 2).mean())) split.setdefault(f"{mode}_arm", []).append(float(((pred[:, ARM_DIMS] - gt[:, ARM_DIMS]) ** 2).mean())) split.setdefault(f"{mode}_hand", []).append(float(((pred[:, HAND_DIMS] - gt[:, HAND_DIMS]) ** 2).mean())) if "future_state_pred" in result: fp = result["future_state_pred"][0, :, :62].float().cpu() fstate_mse.setdefault(mode, []).append(float(((fp - gt_fstate) ** 2).mean())) record["openloop_action_mse"] = {k: [float(np.mean(v)), float(np.std(v)), len(v)] for k, v in openloop.items()} record["openloop_action_mse_raw"] = openloop record["openloop_action_mse_split"] = {k: float(np.mean(v)) for k, v in split.items()} record["openloop_future_state_mse"] = {k: float(np.mean(v)) for k, v in fstate_mse.items()} print("open-loop done", flush=True) # ---------------- comparison table ---------------- lines = [] lines.append(f"FACT params: total {record['params_total']/1e6:.1f}M (coarse {record['params_coarse']/1e6:.1f}M, force {record['params_force']/1e6:.1f}M)") for mode in modes: m, s, n = record["openloop_action_mse"][mode] lines.append(f" open-loop {mode:15s} first 11 steps (0.55 s, T-Rex-horizon) mse {record['openloop_action_mse'][mode + '_first11'][0]:.4f}") lines.append( f" open-loop {mode:15s} mse {m:.4f} ± {s:.4f} (n={n}) arm {record['openloop_action_mse_split'][mode+'_arm']:.4f} hand {record['openloop_action_mse_split'][mode+'_hand']:.4f}" ) for tau_key, vals in tau_table.items(): lines.append(f" tau={tau_key:>4s} action {vals['action_loss'][0]:.4f} state {vals['state_loss'][0]:.4f} force {vals['force_loss'][0]:.4f}") if args.baseline_json: base = json.load(open(args.baseline_json)) lines.append("baseline (trex_track_force ablation, same protocol):") for name, rec in base.items(): ol = rec.get("openloop_action_mse", {}) parts = [f"{k} {v[0]:.4f}" for k, v in ol.items()] sp = rec.get("openloop_action_mse_split", {}) parts += [f"{k} {v:.4f}" for k, v in sp.items()] lines.append(f" {name:8s} " + " ".join(parts)) record["baseline"] = { name: {"openloop_action_mse": rec.get("openloop_action_mse"), "split": rec.get("openloop_action_mse_split")} for name, rec in base.items() } summary = "\n".join(lines) print(summary, flush=True) record["summary"] = summary Path(args.out).write_text(json.dumps(record, indent=1)) print("wrote", args.out) if __name__ == "__main__": main()