| """Open-loop evaluation of a (text-conditioned) FACT checkpoint over many episodes. |
| |
| Same per-anchor protocol as scripts/eval.py (K=1 chunk, pure-noise start, |
| seed 500+i, normalized 62-D delta-base action MSE, arm/hand split, first-11-step |
| horizon-matched MSE, future-state MSE), applied to ``--anchors-per-episode`` |
| evenly spaced anchors of every requested episode. Latents come from the |
| per-episode cache directory, text from a umT5 text cache. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
|
|
| from fact import FactPolicy |
| from fact.data import ( |
| FactEpisode, |
| TextCache, |
| latent_cache_path, |
| load_episode_tasks, |
| ) |
| from fact.runtime_stats import TrexRuntimeStatistics |
|
|
| ARM = list(range(0, 9)) + list(range(31, 40)) |
| HAND = list(range(9, 31)) + list(range(40, 62)) |
|
|
|
|
| def parse_episodes(spec: str) -> list[int]: |
| out: list[int] = [] |
| for part in spec.split(","): |
| part = part.strip() |
| if "-" in part: |
| a, b = part.split("-") |
| out.extend(range(int(a), int(b) + 1)) |
| elif part: |
| out.append(int(part)) |
| return out |
|
|
|
|
| def main() -> None: |
| p = argparse.ArgumentParser() |
| p.add_argument("--checkpoint", required=True) |
| p.add_argument("--dataset-root", required=True) |
| p.add_argument("--cache-dir", required=True, help="dir with latents/episode_XXXXXX.npy") |
| p.add_argument("--text-cache", default=None, help="text_cache.npz (default <cache-dir>/text_cache.npz)") |
| p.add_argument("--stats-root", default=None, help="dataset root whose q01/q99 stats normalize inputs/outputs (default dataset-root)") |
| p.add_argument("--episodes", required=True, help="e.g. 0-9,20") |
| p.add_argument("--anchors-per-episode", type=int, default=8) |
| p.add_argument("--out", required=True) |
| p.add_argument("--bf16", action="store_true") |
| p.add_argument("--dump", default=None, help="npz path: raw (denormalized delta-base) predictions/GT per anchor") |
| p.add_argument("--tau-losses", action="store_true", help="also compute fixed-tau flow losses on K=4 anchors") |
| p.add_argument("--tau-grid", default="0.8,0.6,0.4,0.2,0.05") |
| args = p.parse_args() |
|
|
| device = torch.device("cuda") |
| model = FactPolicy.load(args.checkpoint).to(device).eval() |
| cfg = model.cfg |
| root = Path(args.dataset_root) |
| cache = Path(args.cache_dir) |
| text_cache = None |
| if cfg.use_text: |
| text_cache = TextCache(args.text_cache or cache / "text_cache.npz") |
| stats = TrexRuntimeStatistics.from_dataset(args.stats_root or root) |
| tasks = load_episode_tasks(root) |
|
|
| per_episode: dict[int, dict[str, float]] = {} |
| all_metrics: dict[str, list[float]] = {} |
| dump: dict[str, list] = {"gt": [], "pred_cascade": [], "pred_coarse_to_zero": [], "episode": [], "anchor": []} |
| tau_losses: dict[str, dict[str, list[float]]] = {} |
| denorm = lambda x: (np.asarray(x, np.float32) + 1.0) * 0.5 * (stats.action_q99 - stats.action_q01) + stats.action_q01 |
| for ep in parse_episodes(args.episodes): |
| episode = FactEpisode(root, ep, latent_cache_path(cache, ep), stats=stats, text_cache=text_cache, task=tasks.get(ep)) |
| anchors = episode.valid_anchor_times(1) |
| picks = [anchors[int(i)] for i in np.linspace(0, len(anchors) - 1, min(args.anchors_per_episode, len(anchors)))] |
| ep_metrics: dict[str, list[float]] = {} |
| if args.tau_losses: |
| k4 = episode.valid_anchor_times(4) |
| for i, anchor in enumerate([k4[int(j)] for j in np.linspace(0, len(k4) - 1, min(3, len(k4)))] if k4 else []): |
| sample4 = episode.build(anchor, blocks=4) |
| batch4 = {k: torch.as_tensor(v)[None].to(device) for k, v in sample4.items() |
| if k not in ("action_delta_raw", "reference_state", "obs_frame_index")} |
| for tau_value in (float(t) for t in args.tau_grid.split(",")): |
| for rep in range(2): |
| torch.manual_seed(100000 + rep * 1000 + i * 10) |
| with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16, enabled=args.bf16): |
| o = model.forward_core(batch4, tau=torch.full((1, 4), tau_value, device=device)) |
| for key in ("action_loss", "state_loss", "force_loss"): |
| tau_losses.setdefault(f"{tau_value:g}", {}).setdefault(key, []).append(float(o[key])) |
| for i, anchor in enumerate(picks): |
| sample = episode.build(anchor, blocks=1) |
| gt = torch.as_tensor(sample["action"][0, :, :62]) |
| gt_fs = torch.as_tensor(sample["future_state"][0, :, :62]) |
| batch = {k: torch.as_tensor(v)[None].to(device) for k, v in sample.items() |
| if k not in ("action_delta_raw", "reference_state", "obs_frame_index")} |
| for mode in ("cascade", "coarse_only", "coarse_to_zero"): |
| with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16, enabled=args.bf16): |
| r = model.sample(batch, seed=500 + i, mode=mode) |
| pred = r["action_pred"][0, :, :62].float().cpu() |
| ep_metrics.setdefault(mode, []).append(float(((pred - gt) ** 2).mean())) |
| ep_metrics.setdefault(f"{mode}_first11", []).append(float(((pred[:11] - gt[:11]) ** 2).mean())) |
| ep_metrics.setdefault(f"{mode}_arm", []).append(float(((pred[:, ARM] - gt[:, ARM]) ** 2).mean())) |
| ep_metrics.setdefault(f"{mode}_hand", []).append(float(((pred[:, HAND] - gt[:, HAND]) ** 2).mean())) |
| if mode == "cascade" and "future_state_pred" in r: |
| fp = r["future_state_pred"][0, :, :62].float().cpu() |
| ep_metrics.setdefault("future_state", []).append(float(((fp - gt_fs) ** 2).mean())) |
| if mode in ("cascade", "coarse_to_zero"): |
| dump[f"pred_{mode}"].append(denorm(pred.numpy())) |
| dump["gt"].append(sample["action_delta_raw"][0]) |
| dump["episode"].append(ep) |
| dump["anchor"].append(anchor) |
| per_episode[ep] = {k: float(np.mean(v)) for k, v in ep_metrics.items()} |
| for k, v in ep_metrics.items(): |
| all_metrics.setdefault(k, []).extend(v) |
| print(f"episode {ep:4d} ({episode.length} frames, {len(picks)} anchors): " |
| + " ".join(f"{k} {per_episode[ep][k]:.4f}" for k in ("cascade", "coarse_to_zero", "future_state") if k in per_episode[ep]), flush=True) |
|
|
| summary = {k: [float(np.mean(v)), float(np.std(v)), len(v)] for k, v in all_metrics.items()} |
| if args.dump: |
| np.savez_compressed(args.dump, **{k: np.asarray(v) for k, v in dump.items()}) |
| print("dumped", args.dump) |
| record = {"checkpoint": args.checkpoint, "dataset_root": str(root), "episodes": sorted(per_episode), |
| "anchors_per_episode": args.anchors_per_episode, "config": json.loads(cfg.to_json()), |
| "params_total": sum(p.numel() for p in model.parameters()), |
| "openloop_action_mse": summary, "per_episode": per_episode, |
| "forward_tau_losses": {t: {k: [float(np.mean(v)), float(np.std(v))] for k, v in d.items()} for t, d in tau_losses.items()}} |
| Path(args.out).write_text(json.dumps(record, indent=1)) |
| print("summary:", {k: round(v[0], 4) for k, v in summary.items() if not k.endswith(("_arm", "_hand"))}) |
| print("wrote", args.out) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
