diff --git a/rlvr/autoresearch/ego_shape_diag/README.md b/rlvr/autoresearch/ego_shape_diag/README.md new file mode 100644 index 000000000..c8f9963fe --- /dev/null +++ b/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 --scenes \ + [--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. diff --git a/rlvr/autoresearch/ego_shape_diag/__init__.py b/rlvr/autoresearch/ego_shape_diag/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/rlvr/autoresearch/ego_shape_diag/_common.py b/rlvr/autoresearch/ego_shape_diag/_common.py new file mode 100644 index 000000000..ddf1d8ae2 --- /dev/null +++ b/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] diff --git a/rlvr/autoresearch/ego_shape_diag/check_tl_gate.py b/rlvr/autoresearch/ego_shape_diag/check_tl_gate.py new file mode 100644 index 000000000..1317c4183 --- /dev/null +++ b/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 --scenes \ + [--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()) diff --git a/rlvr/test_tl_gate.py b/rlvr/test_tl_gate.py new file mode 100644 index 000000000..29bc43969 --- /dev/null +++ b/rlvr/test_tl_gate.py @@ -0,0 +1,163 @@ +"""Tests for the traffic-light acceptance gate. + +The gate logic is tested without a model: traffic-light classification reads the scene, +and the verdict is a pure function of the measured spans. +""" + +import numpy as np +import pytest +import torch +from diffusion_planner.dimensions import ( + TRAFFIC_LIGHT_GREEN, + TRAFFIC_LIGHT_RED, + TRAFFIC_LIGHT_YELLOW, +) + +from rlvr.autoresearch.ego_shape_diag import _common +from rlvr.autoresearch.ego_shape_diag.check_tl_gate import ( + route_tl_class, + untested_halves, + verdict, +) + + +def _scene(*onehot_indices: int, valid: bool = True) -> dict: + """One scene whose route lanes carry the given traffic-light one-hots.""" + route = np.zeros((max(len(onehot_indices), 1), 20, 33), dtype=np.float32) + for lane, index in enumerate(onehot_indices): + route[lane, :, 0] = 1.0 # mark the lane valid + route[lane, 0, index] = 1.0 + if not valid: + route[:] = 0.0 + return {"route_lanes": torch.from_numpy(route)} + + +def test_route_tl_class_reads_each_state(): + assert route_tl_class(_scene(TRAFFIC_LIGHT_GREEN)) == "green" + assert route_tl_class(_scene(TRAFFIC_LIGHT_YELLOW)) == "amber" + assert route_tl_class(_scene(TRAFFIC_LIGHT_RED)) == "red" + + +def test_route_tl_class_unsignalled_is_none(): + # a valid lane carrying no traffic-light one-hot, and an all-zero (absent) route + assert route_tl_class(_scene()) == "none" + assert route_tl_class(_scene(TRAFFIC_LIGHT_GREEN, valid=False)) == "none" + + +def test_route_tl_class_mixed_states_are_ambiguous(): + """route_lanes can hold a perpendicular red while the ego's own approach is green. + + Picking either one would be a guess, and guessing "red" makes this gate fail a normal + plan for running a light the ego never faced. Such scenes are excluded instead. + """ + assert route_tl_class(_scene(TRAFFIC_LIGHT_GREEN, TRAFFIC_LIGHT_RED)) == "ambiguous" + assert route_tl_class(_scene(TRAFFIC_LIGHT_GREEN, TRAFFIC_LIGHT_YELLOW)) == "ambiguous" + # agreeing lanes are not ambiguous + assert route_tl_class(_scene(TRAFFIC_LIGHT_RED, TRAFFIC_LIGHT_RED)) == "red" + + +def test_verdict_passes_when_it_goes_on_green_and_holds_on_red(): + spans = {"green": [16.0, 18.0], "amber": [0.1], "red": [0.04], "none": []} + assert verdict(spans, min_green_m=15.0, max_red_m=2.0) == [] + + +def test_verdict_catches_a_model_that_will_not_go(): + spans = {"green": [1.0, 2.0], "amber": [], "red": [0.04], "none": []} + (failure,) = verdict(spans, min_green_m=15.0, max_red_m=2.0) + assert "does not commit to go" in failure + + +def test_verdict_catches_a_model_that_runs_the_light(): + """The failure an unstratified span gate scores as healthy.""" + spans = {"green": [20.0], "amber": [], "red": [19.0], "none": []} + (failure,) = verdict(spans, min_green_m=15.0, max_red_m=2.0) + assert "runs the light" in failure + + +def test_verdict_refuses_a_vacuous_pass(): + """No signalled scene means the gate proved nothing; it must not report PASS.""" + with pytest.raises(SystemExit, match="no green and no red scenes"): + verdict({"green": [], "amber": [1.0], "red": [], "none": [5.0]}, 15.0, 2.0) + + +def test_verdict_refuses_a_green_only_set(): + """The case that matters: a green-only set cannot detect a red-runner.""" + green_only = {"green": [16.0], "amber": [], "red": [], "none": []} + with pytest.raises(SystemExit, match="no red scenes"): + verdict(green_only, 15.0, 2.0) + # the override grades what is there, and the untested half is named for the caller + assert verdict(green_only, 15.0, 2.0, allow_one_sided=True) == [] + assert untested_halves(green_only) == ["red"] + + +def test_verdict_refuses_a_red_only_set(): + """The mirror case: a red-only set cannot detect a checkpoint that will not move.""" + red_only = {"green": [], "amber": [], "red": [0.04], "none": []} + with pytest.raises(SystemExit, match="no green scenes"): + verdict(red_only, 15.0, 2.0) + assert verdict(red_only, 15.0, 2.0, allow_one_sided=True) == [] + assert untested_halves(red_only) == ["green"] + + +def test_one_sided_override_still_fails_a_bad_checkpoint(): + """The override relaxes coverage, never the thresholds.""" + runs_reds = {"green": [], "amber": [], "red": [19.0], "none": []} + (failure,) = verdict(runs_reds, 15.0, 2.0, allow_one_sided=True) + assert "runs the light" in failure + + +def test_untested_halves_empty_when_both_covered(): + assert untested_halves({"green": [16.0], "amber": [], "red": [0.04], "none": []}) == [] + + +def test_amber_is_reported_but_never_gated(): + holds = {"green": [16.0], "amber": [0.1], "red": [0.04], "none": []} + proceeds = {"green": [16.0], "amber": [17.0], "red": [0.04], "none": []} + assert verdict(holds, 15.0, 2.0) == verdict(proceeds, 15.0, 2.0) == [] + + +def test_load_deployable_model_refuses_a_training_milestone(tmp_path): + """A milestone holds raw AND EMA weights; load_model would silently take the raw.""" + milestone = tmp_path / "milestone.pth" + torch.save({"model": {"w": torch.zeros(1)}, "ema_state_dict": {"w": torch.ones(1)}}, milestone) + with pytest.raises(SystemExit, match="training milestone"): + _common.load_deployable_model(str(milestone), torch.device("cpu")) + + +def test_set_shape_replaces_dimensions_without_touching_the_original(): + scene = {"ego_shape": torch.tensor([2.75, 4.34, 1.84])} + out = _common.set_shape(scene, 4.76, 7.24, 2.43) + assert torch.allclose(out["ego_shape"], torch.tensor([4.76, 7.24, 2.43])) + assert torch.allclose(scene["ego_shape"], torch.tensor([2.75, 4.34, 1.84])) + + +def test_parse_shape_rejects_a_wrong_length(): + assert _common.parse_shape("4.76,7.24,2.43") == (4.76, 7.24, 2.43) + with pytest.raises(SystemExit, match="wheelbase,length,width"): + _common.parse_shape("4.76,7.24") + + +def test_report_flags_a_degenerate_class(capsys): + """All-identical spans mean one situation sampled n times, not n samples.""" + from rlvr.autoresearch.ego_shape_diag.check_tl_gate import report + + # near-identical, as real runs of one situation are - not bit-equal + green = [16.8064 + i * 1e-6 for i in range(23)] + spans = {"green": green, "amber": [], "red": [0.04, 0.04], "none": [], "ambiguous": []} + report(spans, "ckpt.pth", 25, None) + out = capsys.readouterr().out + assert "one situation" in out + assert "median" in out and "min" in out and "max" in out + + +def test_report_exposes_a_split_the_mean_would_hide(capsys): + """18 scenes drive, 5 stand still: the mean clears 15 m while 5 scenes are collapsed.""" + from rlvr.autoresearch.ego_shape_diag.check_tl_gate import report + + spans = {"green": [20.0] * 18 + [1.0] * 5, "amber": [], "red": [], "none": [], "ambiguous": []} + report(spans, "ckpt.pth", 23, None) + out = capsys.readouterr().out + # the mean clears the 15 m threshold, but min exposes the collapsed scenes + assert "mean 15.87" in " ".join(out.split(" ")) or "15.87" in out + assert "min 1.00" in out + assert "all identical" not in out # and this class is not degenerate