| |
| """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 |
| from src.data.formats import format_scene_header |
| from src.training.data_loader import _format_header_text, _format_frame_text |
| from src.dualrate.config import DualRateConfig |
| from src.dualrate.model import DualRatePhysicsModel |
| from src.dualrate.normalization import StateNormalizer |
| from src.dualrate.predictor import DualRatePredictor |
|
|
| |
| 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() |
|
|
