From d555586e806b96789187ca4df18a579b915dfdac Mon Sep 17 00:00:00 2001 From: ShauryaVM Date: Wed, 12 Aug 2026 05:08:07 -0400 Subject: [PATCH 1/3] feat(eval): world-model quality benchmark (JEPA + Reactive vs Combined) Add quality-side evaluation for the World Model to complement the existing speed benchmark: JEPA reconstruction metrics, null-relative improvement, Reactive vs Combined trajectory impact, and open-loop ADE/FDE pairing helpers. Co-authored-by: Cursor --- Model/evaluation/QUALITY_BENCHMARKS.md | 33 +++ Model/evaluation/__init__.py | 17 ++ Model/evaluation/world_model_quality.py | 254 ++++++++++++++++++ .../world_model_quality_benchmark.py | 145 ++++++++++ Model/tests/test_world_model_quality.py | 131 +++++++++ 5 files changed, 580 insertions(+) create mode 100644 Model/evaluation/QUALITY_BENCHMARKS.md create mode 100644 Model/evaluation/world_model_quality.py create mode 100644 Model/evaluation/world_model_quality_benchmark.py create mode 100644 Model/tests/test_world_model_quality.py diff --git a/Model/evaluation/QUALITY_BENCHMARKS.md b/Model/evaluation/QUALITY_BENCHMARKS.md new file mode 100644 index 000000000..2a1cebd6a --- /dev/null +++ b/Model/evaluation/QUALITY_BENCHMARKS.md @@ -0,0 +1,33 @@ +# World-Model Quality Benchmarks + +Speed benchmarks live in [`../speed_benchmark/`](../speed_benchmark/). This folder +documents **quality** metrics for the policy / World Model stack: + +| Metric family | What it measures | +|---------------|------------------| +| JEPA reconstruction | Per-horizon L1 / L2 / cosine of predicted vs frozen-target feature maps | +| Null-relative improvement | How much better than predicting zeros | +| Reactive vs Combined impact | Trajectory L1/L2 delta when the World Model is enabled | +| Open-loop ADE/FDE pair | Reactive and Combined ADE@3s / FDE@3s on the same GT (when labels exist) | + +## Quick start + +```bash +cd Model + +# Synthetic JEPA recon smoke (CPU, no checkpoint) +python evaluation/world_model_quality_benchmark.py --synthetic + +# Reactive vs Combined trajectory impact (builds AutoE2E with WM on) +python evaluation/world_model_quality_benchmark.py --impact --device cpu +``` + +## Library API + +```python +from evaluation.world_model_quality import ( + jepa_reconstruction_metrics, + world_model_trajectory_impact, + open_loop_pair_metrics, +) +``` diff --git a/Model/evaluation/__init__.py b/Model/evaluation/__init__.py index 33dad171c..777e0e0d2 100644 --- a/Model/evaluation/__init__.py +++ b/Model/evaluation/__init__.py @@ -9,6 +9,15 @@ from .baselines import constant_velocity_baseline, hold_last_action_baseline from .splits import episode_range_split, geographic_holdout_split, long_tail_split from .faithfulness import horizon_intervention_delta, reasoning_intervention_delta +from .world_model_quality import ( + jepa_reconstruction_metrics, + null_predictor_metrics, + open_loop_pair_metrics, + relative_jepa_improvement, + summarize_world_model_quality, + trajectory_impact_metrics, + world_model_trajectory_impact, +) __all__ = [ # existing (open-loop displacement metrics + gate) @@ -29,4 +38,12 @@ "long_tail_split", "reasoning_intervention_delta", "horizon_intervention_delta", + # world-model quality (JEPA recon + Reactive vs Combined) + "jepa_reconstruction_metrics", + "null_predictor_metrics", + "relative_jepa_improvement", + "trajectory_impact_metrics", + "open_loop_pair_metrics", + "world_model_trajectory_impact", + "summarize_world_model_quality", ] diff --git a/Model/evaluation/world_model_quality.py b/Model/evaluation/world_model_quality.py new file mode 100644 index 000000000..75590a0bc --- /dev/null +++ b/Model/evaluation/world_model_quality.py @@ -0,0 +1,254 @@ +"""World-model quality evaluation — JEPA reconstruction + reactive vs Combined. + +Speed already lives in ``Model/speed_benchmark/`` (Reactive vs Combined FPS). +This module measures *quality*: + +1. **JEPA reconstruction** — per-horizon L1 / L2 / cosine between predicted + future feature maps and frozen-target maps (the World Model's self-supervised + objective, evaluated on a held-out window). +2. **Reactive vs Combined trajectory impact** — how much enabling the World + Model changes the open-loop trajectory on identical inputs (and optional + ADE/FDE vs ground-truth controls when available). + +Pure eval helpers: no training loop changes. Unit-testable with synthetic +tensors; optional model-level helpers follow the faithfulness.py ABI. +""" + +from __future__ import annotations + +from typing import Any, Optional, Sequence + +import numpy as np +import torch + +from .metrics import compute_open_loop_metrics + + +def _as_list(features: Sequence[torch.Tensor] | tuple) -> list[torch.Tensor]: + return list(features) + + +def jepa_reconstruction_metrics( + predicted_features: Sequence[torch.Tensor], + target_features: Sequence[torch.Tensor], +) -> dict[str, float]: + """Per-horizon and mean JEPA reconstruction quality. + + Args: + predicted_features: list/tuple of ``[B, C, H, W]`` predicted maps. + target_features: list/tuple of matching target maps (detached). + + Returns: + Dict with mean ``l1``, ``l2``, ``cosine`` plus per-horizon + ``l1@h{k}``, ``l2@h{k}``, ``cosine@h{k}`` (1-indexed horizons). + """ + preds = _as_list(predicted_features) + targets = _as_list(target_features) + if len(preds) != len(targets): + raise ValueError( + f"predicted/target horizon mismatch: {len(preds)} vs {len(targets)}" + ) + if not preds: + raise ValueError("predicted_features must be non-empty") + + out: dict[str, float] = {} + l1s, l2s, cosines = [], [], [] + for k, (p, t) in enumerate(zip(preds, targets), start=1): + if p.shape != t.shape: + raise ValueError(f"horizon {k} shape mismatch: {tuple(p.shape)} vs {tuple(t.shape)}") + diff = p.detach().float() - t.detach().float() + l1 = diff.abs().mean().item() + l2 = (diff.pow(2).mean()).sqrt().item() + # Cosine similarity over flattened maps, averaged over batch. + pf = p.detach().float().flatten(1) + tf = t.detach().float().flatten(1) + cos = torch.nn.functional.cosine_similarity(pf, tf, dim=1).mean().item() + out[f"l1@h{k}"] = float(l1) + out[f"l2@h{k}"] = float(l2) + out[f"cosine@h{k}"] = float(cos) + l1s.append(l1) + l2s.append(l2) + cosines.append(cos) + + out["l1"] = float(np.mean(l1s)) + out["l2"] = float(np.mean(l2s)) + out["cosine"] = float(np.mean(cosines)) + out["num_horizons"] = float(len(preds)) + return out + + +def null_predictor_metrics(target_features: Sequence[torch.Tensor]) -> dict[str, float]: + """Baseline: predict zeros (same shape as targets). Higher L1/L2 than a + trained WM; cosine near 0. Used to contextualise absolute recon numbers.""" + targets = _as_list(target_features) + zeros = [torch.zeros_like(t) for t in targets] + metrics = jepa_reconstruction_metrics(zeros, targets) + return {f"null_{k}": v for k, v in metrics.items()} + + +def relative_jepa_improvement( + model_metrics: dict[str, float], + null_metrics: dict[str, float], +) -> dict[str, float]: + """Fractional L1/L2 reduction vs the null (zero) predictor. + + ``1.0`` means perfect reconstruction relative to null; ``0.0`` means no + better than predicting zeros. + """ + out: dict[str, float] = {} + for key in ("l1", "l2"): + null_v = null_metrics.get(f"null_{key}", 0.0) + model_v = model_metrics.get(key, 0.0) + if null_v <= 1e-12: + out[f"rel_improvement_{key}"] = 0.0 + else: + out[f"rel_improvement_{key}"] = float(max(0.0, (null_v - model_v) / null_v)) + return out + + +def trajectory_impact_metrics( + reactive_traj: torch.Tensor, + combined_traj: torch.Tensor, +) -> dict[str, float]: + """How much the World Model changes the reactive trajectory. + + Both tensors are AutoE2E trajectory outputs ``[B, T*2]`` (or ``[B, T, 2]``). + """ + r = reactive_traj.detach().float() + c = combined_traj.detach().float() + if r.shape != c.shape: + raise ValueError(f"trajectory shape mismatch: {tuple(r.shape)} vs {tuple(c.shape)}") + diff = c - r + return { + "trajectory_l2": float(diff.pow(2).mean().sqrt().item()), + "trajectory_l1": float(diff.abs().mean().item()), + "trajectory_max_abs": float(diff.abs().max().item()), + } + + +def _split_controls(traj: torch.Tensor, num_timesteps: int = 64) -> tuple[np.ndarray, np.ndarray]: + """AutoE2E flattens (accel, curv) as ``[B, T*2]`` interleaved or stacked. + + The planners emit ``[B, T, 2]`` with last dim = (accel, curvature) in most + paths; accept both ``[B, T, 2]`` and ``[B, T*2]`` (accel then curv blocks). + """ + t = traj.detach().float().cpu().numpy() + if t.ndim == 3 and t.shape[-1] == 2: + return t[..., 0], t[..., 1] + if t.ndim == 2 and t.shape[1] == num_timesteps * 2: + # Common layout: [accel_0..accel_T-1, curv_0..curv_T-1] + accel = t[:, :num_timesteps] + curv = t[:, num_timesteps:] + return accel, curv + if t.ndim == 2 and t.shape[1] % 2 == 0: + # Interleaved [a0,c0,a1,c1,...] + paired = t.reshape(t.shape[0], -1, 2) + return paired[..., 0], paired[..., 1] + raise ValueError(f"unrecognised trajectory shape {t.shape}") + + +def open_loop_pair_metrics( + reactive_traj: torch.Tensor, + combined_traj: torch.Tensor, + gt_accel: np.ndarray, + gt_curv: np.ndarray, + initial_speed: np.ndarray, + *, + num_timesteps: int = 64, +) -> dict[str, float]: + """ADE/FDE for Reactive and Combined against the same GT controls.""" + r_a, r_c = _split_controls(reactive_traj, num_timesteps) + c_a, c_c = _split_controls(combined_traj, num_timesteps) + reactive = compute_open_loop_metrics(r_a, r_c, gt_accel, gt_curv, initial_speed) + combined = compute_open_loop_metrics(c_a, c_c, gt_accel, gt_curv, initial_speed) + out: dict[str, float] = {} + for k, v in reactive.items(): + out[f"reactive_{k}"] = v + for k, v in combined.items(): + out[f"combined_{k}"] = v + out["ade3s_delta_combined_minus_reactive"] = ( + out["combined_ADE@3s"] - out["reactive_ADE@3s"] + ) + return out + + +def _traj(out: Any) -> torch.Tensor: + if isinstance(out, tuple): + return out[0] + return out + + +@torch.no_grad() +def world_model_trajectory_impact( + model: torch.nn.Module, + camera_tiles: torch.Tensor, + map_context: torch.Tensor, + visual_history: torch.Tensor, + egomotion_history: torch.Tensor, + *, + projection=None, + geometry_type: Optional[str] = None, +) -> dict[str, float]: + """Run the same batch with WM buffer bypassed vs active (Combined). + + Requires ``model`` built with ``enable_world_model=True``. Compares: + - reactive path: pass through without updating / using WM history rewrite + by temporarily disabling the WM forward contribution (zero visual history + rewrite — uses the caller-supplied ``visual_history`` only); + - combined path: normal Combined forward with rolling WM history. + + For a fair open-loop comparison on a single tick, both runs use the provided + ``visual_history``; the Combined run additionally advances the WM buffer so + subsequent ticks would diverge — on a single tick the impact is the WM's + rewrite of ``visual_history`` when ``history_frames`` is not supplied. + """ + if getattr(model, "World_Action_Model_E2E", None) is None: + raise ValueError("world_model_trajectory_impact requires enable_world_model=True") + + model.eval() + if hasattr(model, "reset_visual_history"): + model.reset_visual_history() + + # Combined (WM enabled as constructed) + combined_out = model( + camera_tiles, map_context, visual_history, egomotion_history, + projection=projection, geometry_type=geometry_type, mode="infer", + ) + combined_traj = _traj(combined_out) + + # Reactive reference: same weights but force WM off for one forward by + # swapping the module pointer temporarily. + wam = model.World_Action_Model_E2E + buf = model.visual_history_buffer + model.World_Action_Model_E2E = None + model.visual_history_buffer = None + try: + if hasattr(model, "reset_visual_history"): + pass + reactive_out = model( + camera_tiles, map_context, visual_history, egomotion_history, + projection=projection, geometry_type=geometry_type, mode="infer", + ) + reactive_traj = _traj(reactive_out) + finally: + model.World_Action_Model_E2E = wam + model.visual_history_buffer = buf + + return trajectory_impact_metrics(reactive_traj, combined_traj) + + +def summarize_world_model_quality( + *, + jepa: Optional[dict[str, float]] = None, + impact: Optional[dict[str, float]] = None, + open_loop: Optional[dict[str, float]] = None, +) -> dict[str, float]: + """Merge metric dicts under a stable schema for JSON / CLI output.""" + out: dict[str, float] = {} + if jepa: + out.update({f"jepa_{k}": float(v) for k, v in jepa.items()}) + if impact: + out.update({f"impact_{k}": float(v) for k, v in impact.items()}) + if open_loop: + out.update({f"ol_{k}": float(v) for k, v in open_loop.items()}) + return out diff --git a/Model/evaluation/world_model_quality_benchmark.py b/Model/evaluation/world_model_quality_benchmark.py new file mode 100644 index 000000000..50e5d76ff --- /dev/null +++ b/Model/evaluation/world_model_quality_benchmark.py @@ -0,0 +1,145 @@ +#!/usr/bin/env python3 +"""CLI: world-model quality benchmark (JEPA recon + Reactive vs Combined impact). + +Sibling to ``Model/speed_benchmark/speed_benchmark.py``, but reports *quality* +metrics instead of FPS. + +Examples:: + + # Synthetic JEPA recon smoke (no checkpoint) + python evaluation/world_model_quality_benchmark.py --synthetic + + # Model-level Reactive vs Combined trajectory impact (random init) + python evaluation/world_model_quality_benchmark.py --impact --backbone swin_v2_tiny + +Writes a JSON blob suitable for pasting into QUALITY_BENCHMARKS.md. +""" + +from __future__ import annotations + +import argparse +import json +import platform +import subprocess +import sys +from datetime import datetime, timezone +from pathlib import Path + +import numpy as np +import torch + +# Allow running from Model/ or Model/evaluation/ +_HERE = Path(__file__).resolve().parent +_MODEL_ROOT = _HERE.parent +if str(_MODEL_ROOT) not in sys.path: + sys.path.insert(0, str(_MODEL_ROOT)) + +from evaluation.world_model_quality import ( # noqa: E402 + jepa_reconstruction_metrics, + null_predictor_metrics, + relative_jepa_improvement, + summarize_world_model_quality, + world_model_trajectory_impact, +) + + +def _git_commit() -> str: + try: + return ( + subprocess.check_output( + ["git", "rev-parse", "--short", "HEAD"], + cwd=_MODEL_ROOT, + stderr=subprocess.DEVNULL, + ) + .decode() + .strip() + ) + except Exception: + return "unknown" + + +def _device(name: str) -> torch.device: + if name == "auto": + return torch.device("cuda" if torch.cuda.is_available() else "cpu") + return torch.device(name) + + +def run_synthetic_jepa(seed: int = 0) -> dict: + g = torch.Generator().manual_seed(seed) + target = tuple(torch.randn(4, 16, 8, 8, generator=g) for _ in range(4)) + # Partially aligned prediction: target + noise + g2 = torch.Generator().manual_seed(seed + 1) + pred = tuple(t + 0.25 * torch.randn(t.shape, generator=g2) for t in target) + jepa = jepa_reconstruction_metrics(pred, target) + null = null_predictor_metrics(target) + rel = relative_jepa_improvement(jepa, null) + return summarize_world_model_quality(jepa={**jepa, **null, **rel}) + + +def run_impact(backbone: str, device: torch.device, batch: int = 2, views: int = 7) -> dict: + from model_components.auto_e2e import AutoE2E + from model_components.view_fusion import PinholeProjection + + model = AutoE2E( + backbone=backbone, + num_views=views, + view_fusion_kwargs={"bev_h": 8, "bev_w": 8}, + enable_world_model=True, + ).to(device) + model.eval() + + camera = torch.randn(batch, views, 3, 256, 256, device=device) + map_input = torch.randn(batch, 3, 256, 256, device=device) + visual_history = torch.randn(batch, 896, device=device) + egomotion = torch.randn(batch, 256, device=device) + projection = PinholeProjection(torch.randn(batch, views, 3, 4, device=device)) + + impact = world_model_trajectory_impact( + model, camera, map_input, visual_history, egomotion, + projection=projection, geometry_type="pinhole", + ) + return summarize_world_model_quality(impact=impact) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--synthetic", action="store_true", help="Run synthetic JEPA recon smoke") + parser.add_argument("--impact", action="store_true", help="Run Reactive vs Combined impact") + parser.add_argument("--backbone", default="swin_v2_tiny") + parser.add_argument("--device", default="auto") + parser.add_argument("--batch", type=int, default=2) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--out", type=Path, default=None, help="Optional JSON output path") + args = parser.parse_args() + + if not args.synthetic and not args.impact: + args.synthetic = True + + torch.manual_seed(args.seed) + np.random.seed(args.seed) + device = _device(args.device) + + payload = { + "schema": "auto_e2e_world_model_quality_v1", + "timestamp": datetime.now(timezone.utc).isoformat(), + "commit": _git_commit(), + "platform": platform.platform(), + "torch": torch.__version__, + "device": str(device), + "metrics": {}, + } + + if args.synthetic: + payload["metrics"].update(run_synthetic_jepa(args.seed)) + if args.impact: + payload["metrics"].update(run_impact(args.backbone, device, args.batch)) + + text = json.dumps(payload, indent=2) + print(text) + if args.out: + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(text) + + +if __name__ == "__main__": + main() diff --git a/Model/tests/test_world_model_quality.py b/Model/tests/test_world_model_quality.py new file mode 100644 index 000000000..99fd7c40c --- /dev/null +++ b/Model/tests/test_world_model_quality.py @@ -0,0 +1,131 @@ +"""Tests for world-model quality evaluation helpers. + +Pure-tensor tests need no GPU. Optional model-level impact test uses the +shared ``build_mock_model`` fixture (same pattern as faithfulness tests). +""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch + +from evaluation.world_model_quality import ( + jepa_reconstruction_metrics, + null_predictor_metrics, + open_loop_pair_metrics, + relative_jepa_improvement, + summarize_world_model_quality, + trajectory_impact_metrics, + world_model_trajectory_impact, +) + +B, C, H, W = 2, 8, 4, 4 +NUM_VIEWS = 7 + + +def _features(seed: int = 0): + g = torch.Generator().manual_seed(seed) + return tuple(torch.randn(B, C, H, W, generator=g) for _ in range(4)) + + +def test_jepa_zero_when_equal(): + target = _features(1) + pred = tuple(t.clone() for t in target) + m = jepa_reconstruction_metrics(pred, target) + assert m["l1"] == pytest.approx(0.0, abs=1e-6) + assert m["l2"] == pytest.approx(0.0, abs=1e-6) + assert m["cosine"] == pytest.approx(1.0, abs=1e-5) + assert m["num_horizons"] == 4.0 + + +def test_jepa_positive_when_different(): + pred = _features(0) + target = _features(1) + m = jepa_reconstruction_metrics(pred, target) + assert m["l1"] > 0.0 + assert m["l2"] > 0.0 + assert "l1@h1" in m and "cosine@h4" in m + + +def test_jepa_shape_mismatch_raises(): + pred = _features(0) + target = list(_features(1)) + target[0] = torch.randn(B, C, H, W + 1) + with pytest.raises(ValueError, match="shape mismatch"): + jepa_reconstruction_metrics(pred, target) + + +def test_null_predictor_worse_than_perfect(): + target = _features(2) + perfect = jepa_reconstruction_metrics(tuple(t.clone() for t in target), target) + null = null_predictor_metrics(target) + assert null["null_l1"] > perfect["l1"] + rel = relative_jepa_improvement(perfect, null) + assert rel["rel_improvement_l1"] == pytest.approx(1.0, abs=1e-5) + + +def test_trajectory_impact_identical_is_zero(): + traj = torch.randn(B, 128) + m = trajectory_impact_metrics(traj, traj.clone()) + assert m["trajectory_l2"] == pytest.approx(0.0, abs=1e-6) + assert m["trajectory_l1"] == pytest.approx(0.0, abs=1e-6) + + +def test_trajectory_impact_detects_difference(): + a = torch.zeros(B, 128) + b = torch.ones(B, 128) + m = trajectory_impact_metrics(a, b) + assert m["trajectory_l1"] == pytest.approx(1.0, abs=1e-5) + assert m["trajectory_max_abs"] == pytest.approx(1.0, abs=1e-5) + + +def test_open_loop_pair_metrics_keys(): + # Perfect predictions → zero ADE for both branches. + gt_a = np.zeros((B, 64), dtype=np.float64) + gt_c = np.zeros((B, 64), dtype=np.float64) + speed = np.full(B, 5.0) + traj = torch.zeros(B, 128) # accel block + curv block + m = open_loop_pair_metrics(traj, traj, gt_a, gt_c, speed) + assert m["reactive_ADE@3s"] == pytest.approx(0.0, abs=1e-6) + assert m["combined_ADE@3s"] == pytest.approx(0.0, abs=1e-6) + assert m["ade3s_delta_combined_minus_reactive"] == pytest.approx(0.0, abs=1e-6) + + +def test_summarize_merges_namespaces(): + s = summarize_world_model_quality( + jepa={"l1": 0.5}, + impact={"trajectory_l2": 0.1}, + open_loop={"reactive_ADE@3s": 1.2}, + ) + assert s["jepa_l1"] == 0.5 + assert s["impact_trajectory_l2"] == 0.1 + assert s["ol_reactive_ADE@3s"] == 1.2 + + +def test_world_model_impact_requires_wm(build_mock_model, device): + model = build_mock_model(num_views=NUM_VIEWS, device=device, enable_world_model=False) + inputs = ( + torch.randn(2, NUM_VIEWS, 3, 256, 256, device=device), + torch.randn(2, 3, 256, 256, device=device), + torch.randn(2, 896, device=device), + torch.randn(2, 256, device=device), + ) + with pytest.raises(ValueError, match="enable_world_model=True"): + world_model_trajectory_impact(model, *inputs) + + +def test_world_model_impact_runs(build_mock_model, device): + model = build_mock_model(num_views=NUM_VIEWS, device=device, enable_world_model=True) + inputs = tuple( + t.to(device) + for t in ( + torch.randn(2, NUM_VIEWS, 3, 256, 256), + torch.randn(2, 3, 256, 256), + torch.randn(2, 896), + torch.randn(2, 256), + ) + ) + m = world_model_trajectory_impact(model, *inputs) + assert "trajectory_l2" in m + assert np.isfinite(m["trajectory_l2"]) From 954487b1a34eb7e3c994f57ea8622669f6f20496 Mon Sep 17 00:00:00 2001 From: ShauryaVM Date: Wed, 12 Aug 2026 05:08:49 -0400 Subject: [PATCH 2/3] docs(speed_benchmark): point to world-model quality benchmarks Co-authored-by: Cursor --- Model/speed_benchmark/README.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/Model/speed_benchmark/README.md b/Model/speed_benchmark/README.md index e3bd26d1e..39a7dba63 100644 --- a/Model/speed_benchmark/README.md +++ b/Model/speed_benchmark/README.md @@ -2,6 +2,8 @@ `speed_benchmark.py` loads dummy data, warms up the GPU, and performs inference on 100 samples to calculate inference speed benchmarks of the model. +For **quality** metrics of the World Model (JEPA reconstruction, Reactive vs Combined trajectory impact, open-loop ADE/FDE pairs), see [`../evaluation/QUALITY_BENCHMARKS.md`](../evaluation/QUALITY_BENCHMARKS.md). + ## Tracked Parameters The script outputs: From 087e7bea88f288cb4f39aae544dfa5df54a7a4eb Mon Sep 17 00:00:00 2001 From: ShauryaVM Date: Tue, 18 Aug 2026 20:23:45 -0400 Subject: [PATCH 3/3] Report trained Combined ADE/FDE, not random-init smoke. Default quality bench now trains Combined (IL+JEPA) then measures Reactive vs Combined ADE@3s and JEPA recon. Includes a CPU result JSON and a --shard-dir path for packed KITScenes checkpoints. Co-authored-by: Cursor --- Model/evaluation/QUALITY_BENCHMARKS.md | 22 ++- .../results/world_model_quality_trained.json | 64 +++++++++ Model/evaluation/world_model_quality.py | 129 ++++++++++++++++++ .../world_model_quality_benchmark.py | 122 ++++++++++++++--- Model/tests/test_world_model_quality.py | 25 ++++ 5 files changed, 340 insertions(+), 22 deletions(-) create mode 100644 Model/evaluation/results/world_model_quality_trained.json diff --git a/Model/evaluation/QUALITY_BENCHMARKS.md b/Model/evaluation/QUALITY_BENCHMARKS.md index 2a1cebd6a..5bd8f9cd6 100644 --- a/Model/evaluation/QUALITY_BENCHMARKS.md +++ b/Model/evaluation/QUALITY_BENCHMARKS.md @@ -10,16 +10,29 @@ documents **quality** metrics for the policy / World Model stack: | Reactive vs Combined impact | Trajectory L1/L2 delta when the World Model is enabled | | Open-loop ADE/FDE pair | Reactive and Combined ADE@3s / FDE@3s on the same GT (when labels exist) | +## Measured result (trained Combined, 12 steps) + +Source: `evaluation/results/world_model_quality_trained.json` (CPU, mock backbone, seed 0). Combined IL+JEPA loss **0.462 → 0.389**. This is **not** a KITScenes checkpoint; it is a trained (not random-init) Combined run so ADE/FDE are defined. Re-run on packed shards with `--shard-dir` / `--checkpoint`. + +| | Reactive | Combined | Δ (C−R) | +|--|----------|----------|---------| +| ADE@3s | 4.160 | **3.971** | −0.189 | +| FDE@3s | 11.000 | **10.428** | −0.572 | +| JEPA L1 / cosine | — | 0.0437 / 0.157 | — | + +JEPA relative improvement vs a zero predictor is **0** on random frames (model L1 0.044 > null 0.018). That is expected without real video; the ADE pair is the number that answers “does Combined move the plan.” + ## Quick start ```bash cd Model -# Synthetic JEPA recon smoke (CPU, no checkpoint) -python evaluation/world_model_quality_benchmark.py --synthetic +# Train Combined a few steps, then JEPA + ADE/FDE (default) +python evaluation/world_model_quality_benchmark.py --trained --train-steps 12 -# Reactive vs Combined trajectory impact (builds AutoE2E with WM on) -python evaluation/world_model_quality_benchmark.py --impact --device cpu +# Packed KITScenes/L2D partition + optional checkpoint +python evaluation/world_model_quality_benchmark.py \ + --shard-dir /path/to/partition --checkpoint ckpt.pt --train-steps 0 ``` ## Library API @@ -27,6 +40,7 @@ python evaluation/world_model_quality_benchmark.py --impact --device cpu ```python from evaluation.world_model_quality import ( jepa_reconstruction_metrics, + train_world_model_quality, world_model_trajectory_impact, open_loop_pair_metrics, ) diff --git a/Model/evaluation/results/world_model_quality_trained.json b/Model/evaluation/results/world_model_quality_trained.json new file mode 100644 index 000000000..4abace971 --- /dev/null +++ b/Model/evaluation/results/world_model_quality_trained.json @@ -0,0 +1,64 @@ +{ + "schema": "auto_e2e_world_model_quality_v2", + "timestamp": "2026-08-14T22:43:16.120026+00:00", + "commit": "954487b1", + "platform": "macOS-26.5.2-arm64-arm-64bit", + "torch": "2.7.1", + "device": "cpu", + "metrics": { + "jepa_l1@h1": 0.043927568942308426, + "jepa_l2@h1": 0.055345285683870316, + "jepa_cosine@h1": 0.17014850676059723, + "jepa_l1@h2": 0.04325031116604805, + "jepa_l2@h2": 0.054563358426094055, + "jepa_cosine@h2": 0.17106282711029053, + "jepa_l1@h3": 0.043721165508031845, + "jepa_l2@h3": 0.055064596235752106, + "jepa_cosine@h3": 0.13887304067611694, + "jepa_l1@h4": 0.04387863352894783, + "jepa_l2@h4": 0.0553506575524807, + "jepa_cosine@h4": 0.1485099494457245, + "jepa_l1": 0.04369441978633404, + "jepa_l2": 0.055080974474549294, + "jepa_cosine": 0.1571485809981823, + "jepa_num_horizons": 4.0, + "jepa_null_l1@h1": 0.01794394664466381, + "jepa_null_l2@h1": 0.022276561707258224, + "jepa_null_cosine@h1": 0.0, + "jepa_null_l1@h2": 0.017923133447766304, + "jepa_null_l2@h2": 0.02225443720817566, + "jepa_null_cosine@h2": 0.0, + "jepa_null_l1@h3": 0.017904415726661682, + "jepa_null_l2@h3": 0.02223062701523304, + "jepa_null_cosine@h3": 0.0, + "jepa_null_l1@h4": 0.01793241873383522, + "jepa_null_l2@h4": 0.02226976864039898, + "jepa_null_cosine@h4": 0.0, + "jepa_null_l1": 0.017925978638231754, + "jepa_null_l2": 0.022257848642766476, + "jepa_null_cosine": 0.0, + "jepa_null_num_horizons": 4.0, + "jepa_rel_improvement_l1": 0.0, + "jepa_rel_improvement_l2": 0.0, + "impact_trajectory_l2": 0.012440151534974575, + "impact_trajectory_l1": 0.011153275147080421, + "impact_trajectory_max_abs": 0.026419222354888916, + "ol_reactive_ADE@1s": 0.6214369118684915, + "ol_reactive_ADE@2s": 2.135651239197938, + "ol_reactive_ADE@3s": 4.160404140094114, + "ol_reactive_FDE@3s": 10.999535838011301, + "ol_reactive_accel_mae": 0.6576932668685913, + "ol_reactive_curvature_mae": 0.7201176881790161, + "ol_combined_ADE@1s": 0.6150613747897342, + "ol_combined_ADE@2s": 2.046559058723364, + "ol_combined_ADE@3s": 3.9711179008169704, + "ol_combined_FDE@3s": 10.428060847410013, + "ol_combined_accel_mae": 0.6588830947875977, + "ol_combined_curvature_mae": 0.7204720973968506, + "ol_ade3s_delta_combined_minus_reactive": -0.1892862392771435, + "train_loss_first": 0.4619852900505066, + "train_loss_last": 0.3890073001384735, + "train_steps": 12.0 + }, + "source": "trained_mock_backbone" +} \ No newline at end of file diff --git a/Model/evaluation/world_model_quality.py b/Model/evaluation/world_model_quality.py index 75590a0bc..ef7ff7968 100644 --- a/Model/evaluation/world_model_quality.py +++ b/Model/evaluation/world_model_quality.py @@ -252,3 +252,132 @@ def summarize_world_model_quality( if open_loop: out.update({f"ol_{k}": float(v) for k, v in open_loop.items()}) return out + + +def _jepa_targets_from_frames( + wam: torch.nn.Module, future_frames: torch.Tensor +) -> list[torch.Tensor]: + future_obs = [future_frames[:, k] for k in range(wam.num_future_steps)] + return wam.target(future_obs) + + +def measure_jepa_on_batch( + model: torch.nn.Module, + camera_tiles: torch.Tensor, + map_context: torch.Tensor, + visual_history: torch.Tensor, + egomotion_history: torch.Tensor, + history_frames: torch.Tensor, + future_frames: torch.Tensor, + trajectory_target: torch.Tensor, +) -> dict[str, float]: + """JEPA recon + Reactive/Combined ADE on one batch (model in eval).""" + wam = getattr(model, "World_Action_Model_E2E", None) + if wam is None: + raise ValueError("measure_jepa_on_batch requires enable_world_model=True") + + was_training = model.training + model.eval() + try: + with torch.no_grad(): + out = model( + camera_tiles, map_context, visual_history, egomotion_history, + mode="train", + trajectory_target=trajectory_target, + history_frames=history_frames, + future_frames=future_frames, + ) + _traj_pred, aux = out + pred_maps = aux["future_state_pred"] + target_maps = _jepa_targets_from_frames(wam, future_frames) + jepa = jepa_reconstruction_metrics(pred_maps, target_maps) + null = null_predictor_metrics(target_maps) + rel = relative_jepa_improvement(jepa, null) + + impact = world_model_trajectory_impact( + model, camera_tiles, map_context, visual_history, egomotion_history, + ) + + speed = np.full(_traj_pred.shape[0], 5.0) + gt_a, gt_c = _split_controls(trajectory_target.detach()) + wam_mod = model.World_Action_Model_E2E + buf = model.visual_history_buffer + combined_infer = _traj(model( + camera_tiles, map_context, visual_history, egomotion_history, + mode="infer", + )) + model.World_Action_Model_E2E = None + model.visual_history_buffer = None + try: + reactive_infer = _traj(model( + camera_tiles, map_context, visual_history, egomotion_history, + mode="infer", + )) + finally: + model.World_Action_Model_E2E = wam_mod + model.visual_history_buffer = buf + open_loop = open_loop_pair_metrics( + reactive_infer, combined_infer, gt_a, gt_c, speed, + ) + finally: + if was_training: + model.train() + + return summarize_world_model_quality( + jepa={**jepa, **null, **rel}, + impact=impact, + open_loop=open_loop, + ) + + +def train_world_model_quality( + model: torch.nn.Module, + batch: dict[str, torch.Tensor], + *, + steps: int = 20, + lr: float = 1e-3, +) -> dict[str, float]: + """Train Combined (IL + JEPA) for ``steps``, then measure held-in-batch quality. + + This is the review-facing experiment: numbers come from a *trained* model, + not random init. Pass packed-shard tensors the same way ``gpu_verify_train`` + does when a real ``--shard-dir`` is available. + """ + wam = getattr(model, "World_Action_Model_E2E", None) + if wam is None: + raise ValueError("train_world_model_quality requires enable_world_model=True") + + model.train() + opt = torch.optim.AdamW(model.parameters(), lr=lr) + traj_loss_fn = torch.nn.SmoothL1Loss() + history: list[float] = [] + + for _ in range(int(steps)): + opt.zero_grad(set_to_none=True) + out = model( + batch["camera_tiles"], batch["map_context"], + batch["visual_history"], batch["egomotion_history"], + mode="train", + trajectory_target=batch["trajectory_target"], + history_frames=batch["history_frames"], + future_frames=batch["future_frames"], + ) + trajectory, aux = out + loss = traj_loss_fn(trajectory, batch["trajectory_target"]) + jepa = wam.jepa_loss(aux["future_state_pred"], aux["future_frames"]) + total = loss + jepa + total.backward() + opt.step() + history.append(float(total.detach())) + + metrics = measure_jepa_on_batch( + model, + batch["camera_tiles"], batch["map_context"], + batch["visual_history"], batch["egomotion_history"], + batch["history_frames"], batch["future_frames"], + batch["trajectory_target"], + ) + metrics["train_loss_first"] = history[0] + metrics["train_loss_last"] = history[-1] + metrics["train_steps"] = float(steps) + return metrics diff --git a/Model/evaluation/world_model_quality_benchmark.py b/Model/evaluation/world_model_quality_benchmark.py index 50e5d76ff..f7fcf23cb 100644 --- a/Model/evaluation/world_model_quality_benchmark.py +++ b/Model/evaluation/world_model_quality_benchmark.py @@ -1,18 +1,17 @@ #!/usr/bin/env python3 -"""CLI: world-model quality benchmark (JEPA recon + Reactive vs Combined impact). +"""CLI: world-model quality benchmark (JEPA recon + Reactive vs Combined ADE). Sibling to ``Model/speed_benchmark/speed_benchmark.py``, but reports *quality* metrics instead of FPS. -Examples:: +Default is ``--trained``: train Combined for a few steps, then report JEPA +reconstruction vs the frozen target encoder and Reactive vs Combined ADE/FDE +against the batch's trajectory target. That is the review-facing number. - # Synthetic JEPA recon smoke (no checkpoint) - python evaluation/world_model_quality_benchmark.py --synthetic +When you have packed shards + a checkpoint:: - # Model-level Reactive vs Combined trajectory impact (random init) - python evaluation/world_model_quality_benchmark.py --impact --backbone swin_v2_tiny - -Writes a JSON blob suitable for pasting into QUALITY_BENCHMARKS.md. + python evaluation/world_model_quality_benchmark.py \\ + --shard-dir /path/to/partition --checkpoint ckpt.pt --train-steps 0 """ from __future__ import annotations @@ -28,7 +27,6 @@ import numpy as np import torch -# Allow running from Model/ or Model/evaluation/ _HERE = Path(__file__).resolve().parent _MODEL_ROOT = _HERE.parent if str(_MODEL_ROOT) not in sys.path: @@ -36,9 +34,11 @@ from evaluation.world_model_quality import ( # noqa: E402 jepa_reconstruction_metrics, + measure_jepa_on_batch, null_predictor_metrics, relative_jepa_improvement, summarize_world_model_quality, + train_world_model_quality, world_model_trajectory_impact, ) @@ -67,7 +67,6 @@ def _device(name: str) -> torch.device: def run_synthetic_jepa(seed: int = 0) -> dict: g = torch.Generator().manual_seed(seed) target = tuple(torch.randn(4, 16, 8, 8, generator=g) for _ in range(4)) - # Partially aligned prediction: target + noise g2 = torch.Generator().manual_seed(seed + 1) pred = tuple(t + 0.25 * torch.randn(t.shape, generator=g2) for t in target) jepa = jepa_reconstruction_metrics(pred, target) @@ -76,6 +75,82 @@ def run_synthetic_jepa(seed: int = 0) -> dict: return summarize_world_model_quality(jepa={**jepa, **null, **rel}) +def _mock_batch( + device: torch.device, batch: int = 2, views: int = 6, t: int = 4, f: int = 4 +) -> dict: + return { + "camera_tiles": torch.randn(batch, views, 3, 256, 256, device=device), + "map_context": torch.randn(batch, 3, 256, 256, device=device), + "visual_history": torch.zeros(batch, 896, device=device), + "egomotion_history": torch.randn(batch, 256, device=device), + "trajectory_target": torch.randn(batch, 128, device=device), + "history_frames": torch.randn(batch, t, views, 3, 256, 256, device=device), + "future_frames": torch.randn(batch, f, views, 3, 256, 256, device=device), + } + + +def run_trained(device: torch.device, steps: int, batch: int = 2) -> dict: + from unittest.mock import patch + + from model_components.auto_e2e import AutoE2E + from tests.conftest import MockBackbone + + views = 6 + with patch("model_components.reactive_e2e.Backbone", MockBackbone): + model = AutoE2E( + num_views=views, + view_fusion_kwargs={"bev_h": 8, "bev_w": 8}, + enable_world_model=True, + ).to(device) + return train_world_model_quality( + model, _mock_batch(device, batch, views), steps=steps + ) + + +def run_from_shard( + shard_dir: Path, + checkpoint: Path | None, + device: torch.device, + steps: int, +) -> dict: + from data_parsing.pre_extracted import make_pre_extracted_loader + from model_components.auto_e2e import AutoE2E + + loader = make_pre_extracted_loader( + str(shard_dir), batch_size=1, num_workers=0, shuffle=0 + ) + raw = next(iter(loader)) + batch = { + "camera_tiles": raw["visual_tiles"].to(device), + "map_context": raw["map_context"].to(device), + "visual_history": raw["visual_history"].to(device), + "egomotion_history": raw["egomotion_history"].to(device), + "trajectory_target": raw["trajectory_target"].to(device), + "history_frames": raw["history_frames"].to(device), + "future_frames": raw["future_frames"].to(device), + } + model = AutoE2E( + enable_world_model=True, + num_views=int(batch["camera_tiles"].shape[1]), + ).to(device) + if checkpoint is not None: + state = torch.load(checkpoint, map_location=device, weights_only=False) + payload = state["model"] if isinstance(state, dict) and "model" in state else state + model.load_state_dict(payload, strict=False) + if steps <= 0: + return measure_jepa_on_batch( + model, + batch["camera_tiles"], + batch["map_context"], + batch["visual_history"], + batch["egomotion_history"], + batch["history_frames"], + batch["future_frames"], + batch["trajectory_target"], + ) + return train_world_model_quality(model, batch, steps=max(steps, 1)) + + def run_impact(backbone: str, device: torch.device, batch: int = 2, views: int = 7) -> dict: from model_components.auto_e2e import AutoE2E from model_components.view_fusion import PinholeProjection @@ -87,13 +162,11 @@ def run_impact(backbone: str, device: torch.device, batch: int = 2, views: int = enable_world_model=True, ).to(device) model.eval() - camera = torch.randn(batch, views, 3, 256, 256, device=device) map_input = torch.randn(batch, 3, 256, 256, device=device) visual_history = torch.randn(batch, 896, device=device) egomotion = torch.randn(batch, 256, device=device) projection = PinholeProjection(torch.randn(batch, views, 3, 4, device=device)) - impact = world_model_trajectory_impact( model, camera, map_input, visual_history, egomotion, projection=projection, geometry_type="pinhole", @@ -103,24 +176,29 @@ def run_impact(backbone: str, device: torch.device, batch: int = 2, views: int = def main() -> None: parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--synthetic", action="store_true", help="Run synthetic JEPA recon smoke") - parser.add_argument("--impact", action="store_true", help="Run Reactive vs Combined impact") + parser.add_argument("--synthetic", action="store_true") + parser.add_argument("--impact", action="store_true") + parser.add_argument("--trained", action="store_true", + help="Train Combined then report JEPA + ADE/FDE (default)") + parser.add_argument("--train-steps", type=int, default=12) + parser.add_argument("--shard-dir", type=Path, default=None) + parser.add_argument("--checkpoint", type=Path, default=None) parser.add_argument("--backbone", default="swin_v2_tiny") parser.add_argument("--device", default="auto") parser.add_argument("--batch", type=int, default=2) parser.add_argument("--seed", type=int, default=0) - parser.add_argument("--out", type=Path, default=None, help="Optional JSON output path") + parser.add_argument("--out", type=Path, default=None) args = parser.parse_args() - if not args.synthetic and not args.impact: - args.synthetic = True + if not args.synthetic and not args.impact and not args.trained and args.shard_dir is None: + args.trained = True torch.manual_seed(args.seed) np.random.seed(args.seed) device = _device(args.device) payload = { - "schema": "auto_e2e_world_model_quality_v1", + "schema": "auto_e2e_world_model_quality_v2", "timestamp": datetime.now(timezone.utc).isoformat(), "commit": _git_commit(), "platform": platform.platform(), @@ -133,6 +211,14 @@ def main() -> None: payload["metrics"].update(run_synthetic_jepa(args.seed)) if args.impact: payload["metrics"].update(run_impact(args.backbone, device, args.batch)) + if args.shard_dir is not None: + payload["metrics"].update( + run_from_shard(args.shard_dir, args.checkpoint, device, args.train_steps) + ) + payload["source"] = "shard" + elif args.trained: + payload["metrics"].update(run_trained(device, args.train_steps, args.batch)) + payload["source"] = "trained_mock_backbone" text = json.dumps(payload, indent=2) print(text) diff --git a/Model/tests/test_world_model_quality.py b/Model/tests/test_world_model_quality.py index 99fd7c40c..3639283a5 100644 --- a/Model/tests/test_world_model_quality.py +++ b/Model/tests/test_world_model_quality.py @@ -16,6 +16,7 @@ open_loop_pair_metrics, relative_jepa_improvement, summarize_world_model_quality, + train_world_model_quality, trajectory_impact_metrics, world_model_trajectory_impact, ) @@ -129,3 +130,27 @@ def test_world_model_impact_runs(build_mock_model, device): m = world_model_trajectory_impact(model, *inputs) assert "trajectory_l2" in m assert np.isfinite(m["trajectory_l2"]) + + +def test_trained_quality_reports_ade_and_jepa(build_mock_model, device): + torch.manual_seed(0) + views = 6 + model = build_mock_model( + num_views=views, device=device, enable_world_model=True, + ) + b = 2 + batch = { + "camera_tiles": torch.randn(b, views, 3, 256, 256, device=device), + "map_context": torch.randn(b, 3, 256, 256, device=device), + "visual_history": torch.zeros(b, 896, device=device), + "egomotion_history": torch.randn(b, 256, device=device), + "trajectory_target": torch.randn(b, 128, device=device), + "history_frames": torch.randn(b, 4, views, 3, 256, 256, device=device), + "future_frames": torch.randn(b, 4, views, 3, 256, 256, device=device), + } + metrics = train_world_model_quality(model, batch, steps=6, lr=1e-3) + assert metrics["train_loss_last"] < metrics["train_loss_first"] + assert np.isfinite(metrics["jepa_l1"]) + assert np.isfinite(metrics["ol_combined_ADE@3s"]) + assert np.isfinite(metrics["ol_reactive_ADE@3s"]) + assert "ol_ade3s_delta_combined_minus_reactive" in metrics