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