| |
| """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 |
| from src.dualrate.data import render_frame_text |
| from src.dualrate.model import DualRatePhysicsModel |
| from src.dualrate.normalization import StateNormalizer |
| from src.dualrate.predictor import DualRatePredictor |
| from src.dualrate.rollout_adapter import DualRateRolloutEvaluator |
|
|
| 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: |
| 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() |
|
|
