File size: 7,499 Bytes
af8c481 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 | """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()
|