feat(eval): traffic-light acceptance gate for planner checkpoints

rlvr/autoresearch/ego_shape_diag/README.md

@@ -0,0 +1,60 @@
+# Traffic-light acceptance gate
+
+A pre-deployment check on a planner checkpoint: **does it plan to go on green and hold on
+red?** Seconds of compute, no closed-loop simulation.
+
+```
+python -m rlvr.autoresearch.ego_shape_diag.check_tl_gate \
+    --model <best_model.pth> --scenes <scene dir> \
+    [--shape wheelbase,length,width] [--goal X,Y]
+```
+
+## Why a stratified gate
+
+Open-loop trajectory error does not detect either failure this gate is for. A model can
+hold a standstill plan at signalised geometry and still score normally on displacement
+metrics; a model that drives through stop lines scores normally too, because the recorded
+future it is compared against is a car that stopped for perfectly ordinary reasons.
+
+Grading such scenes with a single plan-span threshold is worse than useless. Scenes taken
+across a light cycle contain frames where standing still is the *correct* answer, so a
+one-threshold gate marks those as stalls — and, in the same breath, scores a model that
+runs reds as healthy, since running a red produces exactly the long confident plan the
+threshold rewards.
+
+So scenes are grouped by the traffic-light state their own route lanes carry, and each
+group is held to its own standard:
+
+| group | requirement | catches |
+|---|---|---|
+| green | mean plan span ≥ `--min_green_m` | a checkpoint that will not commit to go |
+| red | mean plan span ≤ `--max_red_m` | a checkpoint that drives through the stop |
+| amber | reported, never graded | — |
+
+Amber is deliberately not a criterion: easing and proceeding are both defensible, so it
+is a behaviour difference between checkpoints rather than a pass or a fail. It is printed
+because that difference is worth seeing.
+
+## What makes a result trustworthy
+
+* **Both halves must be exercised.** They catch opposite failures, so a set carrying only
+  one leaves the other undetectable while still printing PASS. A green-only set cannot
+  notice a red-runner. The gate refuses such a set; `--allow_one_sided` overrides that
+  and names the untested half in the result, which is then not a full acceptance.
+* **Point `--model` at a deployable checkpoint.** A training milestone holds both the raw
+  optimizer iterates and the EMA copy, and the loader takes the raw ones — a plausible
+  but wrong verdict. The gate refuses a milestone and prints the conversion.
+* **Scenes need a signalled route.** If no scene carries a traffic-light state the gate
+  fails rather than reporting a vacuous pass.
+
+## Scenes
+
+Use scenes at the geometry you care about, converted with the same converter the runtime
+uses so the tensors match. Two conversion details, both handled here:
+
+* `ego_agent_past` may arrive as float64; the encoder needs float32.
+* a converter may rewrite `goal_pose` onto the ego when the recorded run ends stopped,
+  which makes the scene meaningless — `--goal X,Y` restores the true ego-frame goal.
+
+`--shape` overrides the recorded vehicle dimensions, which is how to check a checkpoint at
+the geometry it will actually be deployed on rather than the one it was recorded with.

rlvr/autoresearch/ego_shape_diag/_common.py

@@ -0,0 +1,98 @@
+"""Shared helpers for the ego-dimension diagnostics. Run all scripts from the repo root."""
+
+from __future__ import annotations
+
+import glob
+import os
+
+import numpy as np
+import torch
+
+from preference_optimization.utils import load_npz_data
+from rlvr.autoresearch.tools.eval_det_avoidance import load_model
+
+
+def load_deployable_model(model_path: str, device):
+    """Load a checkpoint, refusing one whose deployable copy would be skipped.
+
+    ``load_model`` takes ``ckpt["model"]``, and a training milestone holds BOTH that
+    (the raw optimizer iterates) and ``ema_state_dict`` (the deployable copy). Loading
+    the raw weights yields a plausible-looking but wrong verdict with no warning, so
+    fail loudly and say how to convert instead.
+    """
+    ckpt = torch.load(model_path, map_location="cpu", weights_only=False)
+    if isinstance(ckpt, dict) and "ema_state_dict" in ckpt:
+        raise SystemExit(
+            f"{model_path} is a training milestone carrying both raw and EMA weights; "
+            "load_model would silently use the RAW iterates. Convert it first:\n"
+            '  ck = torch.load(src, map_location="cpu", weights_only=False)\n'
+            '  state = {k.replace("module.", ""): v for k, v in ck["ema_state_dict"].items()}\n'
+            '  torch.save({"model": state}, dst)   # copy args.json next to it'
+        )
+    return load_model(model_path, device)
+
+
+def find_scenes(path: str) -> list[str]:
+    """Accept a directory (searched recursively), a glob, or a single .npz."""
+    if os.path.isdir(path):
+        out = sorted(glob.glob(os.path.join(path, "**", "*.npz"), recursive=True))
+    else:
+        out = sorted(glob.glob(path))
+    if not out:
+        raise SystemExit(f"no .npz found under {path!r}")
+    return out
+
+
+def to_f32(d: dict) -> dict:
+    """Some converters emit float64 for ego_agent_past; the encoder requires float32."""
+    return {
+        k: (v.float() if torch.is_tensor(v) and v.dtype == torch.float64 else v)
+        for k, v in d.items()
+    }
+
+
+def load_scene(path: str, device, goal: tuple[float, float] | None = None) -> dict:
+    """Load one scene, optionally overriding the ego-frame goal.
+
+    The override exists because a converter may rewrite ``goal_pose`` onto the ego when
+    the recorded run ends stopped; passing the true goal restores the intended scene.
+    """
+    d = to_f32(load_npz_data(path, device))
+    if goal is not None:
+        g = d["goal_pose"].clone().reshape(-1)
+        g[0], g[1] = goal
+        if g.numel() >= 4:  # [x, y, cos, sin] after heading_to_cos_sin
+            g[2], g[3] = 1.0, 0.0
+        d["goal_pose"] = g.reshape(d["goal_pose"].shape)
+    return d
+
+
+def set_shape(d: dict, wheelbase: float, length: float, width: float) -> dict:
+    """Return a copy of the scene with ``ego_shape`` replaced."""
+    out = dict(d)
+    s = d["ego_shape"].clone().reshape(-1)
+    s[0], s[1], s[2] = wheelbase, length, width
+    out["ego_shape"] = s.reshape(d["ego_shape"].shape)
+    return out
+
+
+def plan_span(model, margs, d: dict, device) -> float:
+    """Straight-line distance from the first to the last predicted ego waypoint.
+
+    This is a coarse "is the model planning to move at all" measure, not a path length:
+    at cruise it lands in the tens of metres, and a value near zero means the model is
+    planning to stand still.
+    """
+    from rlvr.autoresearch.tools.eval_det_avoidance import det_inference_batched
+
+    out = det_inference_batched(model, margs, [d], device)
+    t = out[0] if not isinstance(out, tuple) else out[0][0]
+    t = np.squeeze(t.detach().cpu().numpy())
+    return float(np.hypot(t[-1, 0] - t[0, 0], t[-1, 1] - t[0, 1]))
+
+
+def parse_shape(text: str) -> tuple[float, float, float]:
+    parts = [float(x) for x in text.split(",")]
+    if len(parts) != 3:
+        raise SystemExit(f"expected wheelbase,length,width (got {text!r})")
+    return parts[0], parts[1], parts[2]

rlvr/autoresearch/ego_shape_diag/check_tl_gate.py

@@ -0,0 +1,229 @@
+#!/usr/bin/env python3
+"""ACCEPTANCE GATE, traffic-light stratified: does the model go on green and hold on red?
+
+A stop is only a failure if the light says go. Scenes captured across a light cycle
+therefore cannot be graded with a single plan-span threshold: an unstratified gate marks
+the red frames as stalls and — worse — scores a model that *runs* red lights as healthy.
+Group the scenes by the traffic-light state their own route lanes carry, then require:
+
+    green scenes  ->  mean plan span >= --min_green_m   (the model commits to go)
+    red scenes    ->  mean plan span <= --max_red_m     (the model holds)
+
+Amber is reported but not gated: easing and proceeding are both defensible, so it is a
+behaviour difference between checkpoints rather than a pass or a fail.
+
+    python -m rlvr.autoresearch.ego_shape_diag.check_tl_gate \
+        --model <best_model.pth> --scenes <dir-or-glob> \
+        [--goal X,Y] [--shape wheelbase,length,width] \
+        [--min_green_m 15] [--max_red_m 2]
+
+Seconds of compute and no closed-loop simulation. Worth running on every checkpoint,
+because open-loop trajectory error does not detect this failure: a model can hold a
+standstill plan at signalised geometry while scoring normally on displacement metrics.
+"""
+
+from __future__ import annotations
+
+import argparse
+
+import numpy as np
+import torch
+from diffusion_planner.dimensions import (
+    TRAFFIC_LIGHT_GREEN,
+    TRAFFIC_LIGHT_NO_TRAFFIC_LIGHT,
+    TRAFFIC_LIGHT_RED,
+    TRAFFIC_LIGHT_YELLOW,
+)
+
+from rlvr.autoresearch.ego_shape_diag._common import (
+    find_scenes,
+    load_deployable_model,
+    load_scene,
+    parse_shape,
+    plan_span,
+    set_shape,
+)
+
+_TL_SLICE = slice(TRAFFIC_LIGHT_GREEN, TRAFFIC_LIGHT_NO_TRAFFIC_LIGHT + 1)
+_CLASSES = {
+    TRAFFIC_LIGHT_GREEN - TRAFFIC_LIGHT_GREEN: "green",
+    TRAFFIC_LIGHT_YELLOW - TRAFFIC_LIGHT_GREEN: "amber",
+    TRAFFIC_LIGHT_RED - TRAFFIC_LIGHT_GREEN: "red",
+}
+_ORDER = ("green", "amber", "red", "ambiguous", "none")
+# Two spans closer than this are the same plan, not two measurements worth averaging.
+_SAME_PLAN_M = 0.01
+
+
+def route_tl_class(d: dict) -> str:
+    """The traffic-light state governing the ego, or ``ambiguous`` / ``none``.
+
+    A route tensor does not carry one signal. At an intersection ``route_lanes`` can
+    include a perpendicular approach whose light is correctly red while the ego's own
+    approach is green, so treating "any red lane" as red mislabels ordinary green scenes
+    — and this gate would then fail a perfectly normal plan for running a light the ego
+    never faced.
+
+    Deciding which signal governs the ego needs heading alignment and stop-line
+    proximity. That logic already exists in the C++ frame filters
+    (``detect_red_light_run``) and is deliberately not reimplemented here, because a
+    second, uncross-checked copy of the geometry is how the two drift apart. A route
+    whose signalled lanes disagree is therefore reported as ``ambiguous`` and left out of
+    the verdict: fewer graded scenes, but no invented answer. A route with no signalled
+    lane at all is ``none``, likewise excluded.
+    """
+    rl = d["route_lanes"].detach().cpu().numpy()
+    if rl.ndim > 3:
+        rl = rl.reshape(-1, rl.shape[-2], rl.shape[-1])
+    valid = np.abs(rl).sum(axis=(1, 2)) > 0
+    if not valid.any():
+        return "none"
+    onehot = rl[valid][:, 0, _TL_SLICE]
+    hot = onehot.max(axis=1) > 0
+    if not hot.any():
+        return "none"
+    present = {_CLASSES.get(int(c)) for c in onehot[hot].argmax(axis=1)}
+    present.discard(None)  # one-hot slots outside green/amber/red carry no state
+    if not present:
+        return "none"
+    if len(present) > 1:
+        return "ambiguous"
+    return present.pop()
+
+
+def measure(model, margs, scenes, device, goal, shape) -> dict[str, list[float]]:
+    """Plan span for every scene, bucketed by its route traffic-light state."""
+    spans: dict[str, list[float]] = {k: [] for k in _ORDER}
+    for path in scenes:
+        d = load_scene(path, device, goal)
+        spans[route_tl_class(d)].append(
+            plan_span(model, margs, set_shape(d, *shape) if shape else d, device)
+        )
+    return spans
+
+
+def untested_halves(spans: dict[str, list[float]]) -> list[str]:
+    """The gate halves this scene set cannot exercise, in report order."""
+    return [state for state in ("green", "red") if not spans[state]]
+
+
+def verdict(
+    spans: dict[str, list[float]],
+    min_green_m: float,
+    max_red_m: float,
+    allow_one_sided: bool = False,
+) -> list[str]:
+    """Gate failures, empty when the checkpoint passes.
+
+    Both halves have to be exercised for a PASS to mean anything. They catch opposite
+    failures — the green half catches a checkpoint that will not move, the red half one
+    that drives through a stop — so a set carrying only one of them leaves the other
+    failure undetectable while still printing PASS. A green-only set in particular cannot
+    notice a model that runs red lights, which is the failure an unstratified span gate
+    already scores as healthy.
+    """
+    missing = untested_halves(spans)
+    if missing and not allow_one_sided:
+        raise SystemExit(
+            f"the scene set carries no {' and no '.join(missing)} scenes, so the "
+            f"{'/'.join(missing)} half of this gate would never be tested — yet a PASS would "
+            "claim both. Supply scenes spanning a light cycle, or pass --allow_one_sided to "
+            "grade only the half that is present."
+        )
+    failures = []
+    if spans["green"]:
+        green = float(np.mean(spans["green"]))
+        if green < min_green_m:
+            failures.append(f"green {green:.2f} m < {min_green_m} m (does not commit to go)")
+    if spans["red"]:
+        red = float(np.mean(spans["red"]))
+        if red > max_red_m:
+            failures.append(f"red {red:.2f} m > {max_red_m} m (runs the light)")
+    return failures
+
+
+def report(spans: dict[str, list[float]], model_path: str, n_scenes: int, shape) -> None:
+    """Per class: the mean the gate grades on, and the spread behind it.
+
+    The mean alone is not enough to read a result. A class whose scenes all return the
+    same span is one situation sampled many times, not many independent samples, and its
+    n overstates the evidence; a class that splits into "drives" and "stands still" can
+    average out above the threshold while half its scenes fail. Both are visible in
+    min/median/max and invisible in the mean, so print all four.
+    """
+    print(f"model  : {model_path}")
+    print(f"scenes : {n_scenes}   " + "  ".join(f"{k} {len(spans[k])}" for k in _ORDER))
+    print(f"shape  : {shape if shape else 'as recorded'}")
+    for k in _ORDER:
+        v = sorted(spans[k])
+        if not v:
+            continue
+        # Compare with a tolerance, not for equality: scenes that are the same situation
+        # still differ in the last few decimals, and a spread under a centimetre is not a
+        # difference in behaviour at these thresholds.
+        degenerate = len(v) > 1 and (v[-1] - v[0]) < _SAME_PLAN_M
+        note = (
+            f"  (spread < {_SAME_PLAN_M * 100:.0f} cm - one situation, not {len(v)} samples)"
+            if degenerate
+            else ""
+        )
+        print(
+            f"   {k:<9} n={len(v):<3} mean {float(np.mean(v)):6.2f} m"
+            f"   median {float(np.median(v)):6.2f}   min {v[0]:6.2f}   max {v[-1]:6.2f}{note}"
+        )
+    excluded = len(spans["ambiguous"]) + len(spans["none"])
+    if excluded:
+        print(
+            f"   ({excluded} scene(s) not graded: ambiguous = route lanes disagree, "
+            "none = no signalled lane)"
+        )
+
+
+def _parse_args() -> argparse.Namespace:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model", required=True, help="deployable checkpoint (args.json alongside)")
+    ap.add_argument("--scenes", required=True, help="scene dir, glob, or single .npz")
+    ap.add_argument("--goal", default=None, help="x,y ego-frame goal override")
+    ap.add_argument("--shape", default=None, help="wheelbase,length,width to evaluate at")
+    ap.add_argument("--min_green_m", type=float, default=15.0)
+    ap.add_argument("--max_red_m", type=float, default=2.0)
+    ap.add_argument("--limit", type=int, default=0, help="0 = every scene")
+    ap.add_argument(
+        "--allow_one_sided",
+        action="store_true",
+        help="grade even when the scenes cover only green or only red; the untested half "
+        "is named in the result, which is then not a full acceptance",
+    )
+    return ap.parse_args()
+
+
+def main() -> int:
+    a = _parse_args()
+    goal = tuple(float(x) for x in a.goal.split(",")) if a.goal else None
+    shape = parse_shape(a.shape) if a.shape else None
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model, margs = load_deployable_model(a.model, device)
+    scenes = find_scenes(a.scenes)
+    if a.limit:
+        scenes = scenes[: a.limit]
+
+    spans = measure(model, margs, scenes, device, goal, shape)
+    report(spans, a.model, len(scenes), shape)
+
+    failures = verdict(spans, a.min_green_m, a.max_red_m, a.allow_one_sided)
+    print()
+    for f in failures:
+        print(f"FAIL: {f}")
+    missing = untested_halves(spans)
+    if failures:
+        print("FAIL")
+    elif missing:
+        # Never let a half that was never exercised read as an acceptance.
+        print(f"PASS ({'/'.join(missing)} half NOT tested — not a full acceptance)")
+    else:
+        print("PASS")
+    return 1 if failures else 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())