#!/usr/bin/env python3 """Evaluate a dual-rate checkpoint: single-step + autoregressive rollout. Uses the stock RolloutEvaluator metrics pipeline (via DualRateRolloutEvaluator) so numbers are apples-to-apples with the LFM2 text baseline (22.64px). Scenes come from the numeric cache (val split); OOD scenes are streamed directly from the packed tarballs (they are never cached). """ import argparse import json import sys import tarfile import time from pathlib import Path import numpy as np import torch sys.path.insert(0, str(Path(__file__).parent.parent)) sys.path.insert(0, str(Path(__file__).parent)) from src.dualrate.config import DualRateConfig # noqa: E402 from src.dualrate.data import render_frame_text # noqa: E402 from src.dualrate.model import DualRatePhysicsModel # noqa: E402 from src.dualrate.normalization import StateNormalizer # noqa: E402 from src.dualrate.predictor import DualRatePredictor # noqa: E402 from src.dualrate.rollout_adapter import DualRateRolloutEvaluator # noqa: E402 OOD_TYPES = ["pong", "bowling", "ramp_roll", "angry_birds", "hourglass", "newtons_cradle"] def frames_to_dicts(states, ids, start_idx, count): """Rebuild evaluator-compatible GT frame dicts from cached arrays.""" out = [] for i in range(start_idx, min(start_idx + count, states.shape[0])): objs = [] for j in range(states.shape[1]): px, py, vx, vy, a, av = (float(x) for x in states[i, j]) objs.append({"id": int(ids[j]), "position": {"x": px, "y": py}, "velocity": {"x": vx, "y": vy}, "angle": a, "angular_velocity": av}) out.append({"frame": i + 1, "objects": objs}) return out def build_initial_context(header_text, states, ids, n_ctx_frames): parts = [header_text] for i in range(n_ctx_frames): parts.append(render_frame_text(i + 1, states[i], ids)) parts.append("Predict next frame:") return "".join(parts) def load_scene_npz(path): z = np.load(path) return (z["states"], z["static"], z["ids"], z["gravity"], float(z["dt"]), str(z["header_text"])) def stream_ood_scenes(packed_dir, scenario, limit): sys.path.insert(0, str(Path(__file__).parent)) from build_dualrate_cache import scene_to_arrays out = [] tar_path = Path(packed_dir) / f"{scenario}.tar.gz" with tarfile.open(tar_path, "r:gz") as tf: for member in tf: if len(out) >= limit: break if not member.name.endswith(".jsonl"): continue lines = tf.extractfile(member).read().decode().splitlines() try: _, states, static, ids, gravity, dt, header_text = scene_to_arrays(lines) out.append((states, static, ids, gravity, float(dt), header_text)) except Exception: # noqa: BLE001 continue return out def eval_scenes(scenes, predictor, evaluator, n_ctx, max_steps): results, times = [], [] for states, static, ids, gravity, dt, header_text in scenes: if states.shape[0] < n_ctx + 2: continue ctx = build_initial_context(header_text, states, ids, n_ctx) gt = frames_to_dicts(states, ids, n_ctx, max_steps) predictor.reset_from_arrays(static, ids, gravity, dt, states[n_ctx - 1], start_frame=n_ctx + 1) t0 = time.time() r = evaluator.evaluate_rollout(model=None, initial_context=ctx, ground_truth_frames=gt, max_steps=max_steps) times.append((time.time() - t0) / max(r["num_steps"], 1)) results.append(r) return results, times def summarize(results, times, max_steps): def finite_mean(vals): v = [x for x in vals if np.isfinite(x)] return float(np.mean(v)) if v else float("inf") step1_pos = finite_mean([r["step_errors"][0] for r in results if r["step_errors"]]) step1_vel = finite_mean([r["velocity_errors"][0] for r in results if r["velocity_errors"]]) all_pos = finite_mean([e for r in results for e in r["step_errors"]]) n_steps = max((len(r["step_errors"]) for r in results), default=0) per_step = [finite_mean([r["step_errors"][i] for r in results if len(r["step_errors"]) > i]) for i in range(n_steps)] div_steps = [r["divergence_step"] for r in results] div_rate = float(np.mean([d is not None for d in div_steps])) return { "n_scenes": len(results), "single_step_pos_mse": step1_pos, "single_step_pos_err_px": float(np.sqrt(step1_pos)) if np.isfinite(step1_pos) else None, "single_step_vel_mse": step1_vel, "rollout_mean_pos_mse": all_pos, "per_step_pos_mse": per_step, "divergence_rate": div_rate, "mean_divergence_step": finite_mean([d for d in div_steps if d is not None] or [np.nan]), "sec_per_frame": float(np.mean(times)) if times else None, "max_steps": max_steps, } def main(): ap = argparse.ArgumentParser() ap.add_argument("--ckpt", required=True) ap.add_argument("--scenes", type=int, default=100) ap.add_argument("--max-steps", type=int, default=100) ap.add_argument("--n-ctx", type=int, default=5) ap.add_argument("--tick-k", type=int, default=None) ap.add_argument("--denoise-steps", type=int, default=None) ap.add_argument("--ood", action="store_true") ap.add_argument("--ood-scenes", type=int, default=20) ap.add_argument("--packed-dir", default="data_scenarios_packed/val") ap.add_argument("--out", default=None) args = ap.parse_args() ck = torch.load(args.ckpt, map_location="cpu", weights_only=False) cfg = DualRateConfig.from_dict(ck["config"]) if args.tick_k: cfg.tick_k = args.tick_k if args.denoise_steps: cfg.denoise_steps = args.denoise_steps norm = StateNormalizer(**{k: ck["normalizer"][k] for k in ("state_offset", "state_scale", "delta_scale")}) from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained(cfg.model_name) model = DualRatePhysicsModel(cfg, readout_token_id=tokenizer.pad_token_id).cuda() missing, unexpected = model.load_state_dict(ck["model"], strict=False) missing = [m for m in missing if not m.startswith("semantic.lm.")] assert not missing and not unexpected, (missing, unexpected) model.eval() predictor = DualRatePredictor(model, tokenizer, cfg, norm, device="cuda") evaluator = DualRateRolloutEvaluator(predictor) report = {"ckpt": args.ckpt, "tick_k": cfg.tick_k, "denoise_steps": cfg.denoise_steps, "n_ctx": args.n_ctx} val_files = sorted(Path("data_dualrate/cache/val").glob("*.npz")) rng = np.random.default_rng(0) picks = [val_files[i] for i in rng.permutation(len(val_files))[:args.scenes]] scenes = [load_scene_npz(p) for p in picks] results, times = eval_scenes(scenes, predictor, evaluator, args.n_ctx, args.max_steps) report["seen"] = summarize(results, times, args.max_steps) print("SEEN:", json.dumps(report["seen"], indent=2), flush=True) if args.ood: report["ood"] = {} for sc in OOD_TYPES: scenes = stream_ood_scenes(args.packed_dir, sc, args.ood_scenes) results, times = eval_scenes(scenes, predictor, evaluator, args.n_ctx, args.max_steps) report["ood"][sc] = summarize(results, times, args.max_steps) print(f"OOD {sc}:", json.dumps(report["ood"][sc]), flush=True) out = args.out or (str(Path(args.ckpt).parent / "eval_report.json")) Path(out).write_text(json.dumps(report, indent=2)) print("saved", out) if __name__ == "__main__": main()