File size: 7,844 Bytes
7e86b86 | 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 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 | #!/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()
|