#!/usr/bin/env python3 """Exact replication of PhysicsLLMEngine's evaluate_llm_mse.py protocol. Same 6 scenes (identical generator + simulation code, same seeds/difficulties), same prompt format (header + 10 context frames + "Predict next frame:"), same autoregressive loop (5 steps, prompt re-rendered from the growing frame list each step), same metric (np.mean((pred_pos - gt_pos)**2) per step). Differences vs theirs (documented, favoring neither side): - their LFM2 predicts positions only; fed-back frames carry forward the previous velocity. Our model predicts the full state, so fed-back frames contain our predicted velocities. - our semantic latent refreshes every K frames (dual rate); with K=8 and 5 prediction steps the latent is computed once from the initial context. --tick-k 1 refreshes it from the re-rendered prompt every step. """ import argparse import json import sys from pathlib import Path import numpy as np import torch sys.path.insert(0, str(Path(__file__).parent.parent)) from src.physics.scenario_generator import generate_scenario # noqa: E402 from src.data.formats import format_scene_header # noqa: E402 from src.training.data_loader import _format_header_text, _format_frame_text # noqa: E402 from src.dualrate.config import DualRateConfig # 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 # (scenario_type, difficulty, seed) exactly as in PhysicsLLMEngine/llm_mse_results.json PLLE_SCENES = [ ("pong", 3, 40723578), ("bowling", 3, 59859168), ("ramp_roll", 3, 95887616), ("angry_birds", 2, 12800155), ("hourglass", 2, 15862948), ("newtons_cradle", 3, 63938622), ] PLLE_LLM_AVG = {"pong": 5027.39, "bowling": 25.95, "ramp_roll": 28.28, "angry_birds": 30.21, "hourglass": 45.64, "newtons_cradle": 15.78} def sim_get_frame_dict(sim): """Identical to PhysicsLLMEngine evaluate_llm_mse._sim_get_frame_dict.""" state = sim.get_state() threshold = 1.0 moving = sum(1 for obj in state["objects"] if abs(obj["velocity"]["x"]) > threshold or abs(obj["velocity"]["y"]) > threshold) total = len(state["objects"]) if moving == 0: desc = "All objects are at rest" elif moving == total: desc = "All objects are in motion" else: desc = f"{moving} of {total} objects are moving" state["description"] = f"Frame {state['frame']}: {desc}." return state def build_context_prompt(header_dict, frame_dicts): header_text = _format_header_text(header_dict) frame_texts = [_format_frame_text(f) for f in frame_dicts] return header_text + "".join(frame_texts) + "Predict next frame:" def frame_dict_state(frame_dict, id_order): s = np.zeros((len(id_order), 6), dtype=np.float32) by_id = {o["id"]: o for o in frame_dict["objects"]} for i, oid in enumerate(id_order): o = by_id[oid] s[i] = [o["position"]["x"], o["position"]["y"], o["velocity"]["x"], o["velocity"]["y"], o.get("angle", 0.0), o.get("angular_velocity", 0.0)] return s def pred_frame_dict(state, id_order, prev_frame_dict, frame_num): """Our predicted full state as a frame dict for autoregressive re-rendering.""" prev_by_id = {o["id"]: o for o in prev_frame_dict["objects"]} objects = [] for i, oid in enumerate(id_order): px, py, vx, vy, a, av = (float(x) for x in state[i]) prev = prev_by_id[oid] objects.append({"id": oid, "type": prev.get("type", "unknown"), "position": {"x": round(px, 4), "y": round(py, 4)}, "velocity": {"x": round(vx, 4), "y": round(vy, 4)}, "angle": round(a, 4), "angular_velocity": round(av, 4), "material": prev.get("material", {})}) threshold = 1.0 moving = sum(1 for o in objects if abs(o["velocity"]["x"]) > threshold or abs(o["velocity"]["y"]) > threshold) if moving == 0: desc = "All objects are at rest" elif moving == len(objects): desc = "All objects are in motion" else: desc = f"{moving} of {len(objects)} objects are moving" return {"frame": frame_num, "description": f"Frame {frame_num}: {desc}.", "objects": objects} def evaluate_scene(predictor, scenario_type, difficulty, seed, context_frames=10, predict_frames=5): sim, metadata = generate_scenario(seed=seed, scenario_type=scenario_type, difficulty=difficulty) header_dict = format_scene_header(sim, seed, metadata=metadata) all_frames = [] for _ in range(context_frames + predict_frames): sim.step() all_frames.append(sim_get_frame_dict(sim)) context_dicts = list(all_frames[:context_frames]) gt_dicts = all_frames[context_frames:] id_order = [o["id"] for o in header_dict["objects"]] init_state = frame_dict_state(context_dicts[-1], id_order) predictor.reset(header_dict, init_state, start_frame=context_dicts[-1]["frame"] + 1) per_step_mse = [] for step_idx in range(predict_frames): prompt = build_context_prompt(header_dict, context_dicts) predictor.step(prompt) pred_pos = predictor.state[:, :2].copy() gt_pos = frame_dict_state(gt_dicts[step_idx], id_order)[:, :2] per_step_mse.append(float(np.mean((pred_pos - gt_pos) ** 2))) context_dicts.append(pred_frame_dict(predictor.state, id_order, context_dicts[-1], gt_dicts[step_idx]["frame"])) return {"scenario_type": scenario_type, "difficulty": difficulty, "seed": seed, "num_objects": len(id_order), "per_step_mse": per_step_mse, "avg_mse": float(np.mean(per_step_mse)), "success": True} def main(): ap = argparse.ArgumentParser() ap.add_argument("--ckpt", required=True) ap.add_argument("--tick-k", type=int, default=None) ap.add_argument("--max-ctx-tokens", type=int, default=4096) 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 cfg.max_ctx_tokens = args.max_ctx_tokens 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() model.load_state_dict(ck["model"], strict=False) model.eval() predictor = DualRatePredictor(model, tokenizer, cfg, norm, device="cuda") results = [] for scenario_type, difficulty, seed in PLLE_SCENES: r = evaluate_scene(predictor, scenario_type, difficulty, seed) theirs = PLLE_LLM_AVG[scenario_type] print(f"{scenario_type:15s} ours {r['avg_mse']:10.3f} theirs {theirs:8.2f} " f"per-step {[round(x, 2) for x in r['per_step_mse']]}", flush=True) results.append(r) ours_all = float(np.mean([r["avg_mse"] for r in results])) per_step = np.array([r["per_step_mse"] for r in results]).mean(0) print(f"\nOVERALL ours {ours_all:.3f} vs PhysicsLLMEngine 862.21") print("per-step ours ", [round(x, 2) for x in per_step]) print("per-step theirs [80.32, 308.6, 732.58, 1274.23, 1915.31]") out = args.out or str(Path(args.ckpt).parent / f"plle_exact_k{cfg.tick_k}_{Path(args.ckpt).stem}.json") Path(out).write_text(json.dumps({ "protocol": "PhysicsLLMEngine evaluate_llm_mse exact (ctx=10, predict=5)", "ckpt": args.ckpt, "tick_k": cfg.tick_k, "results": results, "overall_avg_mse": ours_all, "per_step_mean": per_step.tolist(), "plle_reference": {"overall": 862.21, "per_step": [80.32, 308.6, 732.58, 1274.23, 1915.31], "per_type": PLLE_LLM_AVG}, }, indent=2)) print("saved", out) if __name__ == "__main__": main()