diff --git a/Model/evaluation/confidence_calibration.py b/Model/evaluation/confidence_calibration.py new file mode 100644 index 000000000..98358a889 --- /dev/null +++ b/Model/evaluation/confidence_calibration.py @@ -0,0 +1,58 @@ +"""Calibration metrics for reasoning-band confidence (#110).""" + +from __future__ import annotations + +from typing import Dict + +import torch + + +def expected_calibration_error( + confidence: torch.Tensor, + correctness: torch.Tensor, + *, + n_bins: int = 15, +) -> Dict[str, float]: + """Compute ECE between predicted confidence and binary correctness. + + Args: + confidence: probabilities in ``[0, 1]``, any shape (flattened). + correctness: same shape, values in ``{0, 1}`` (or soft ``[0, 1]``). + n_bins: number of equal-width confidence bins. + + Returns: + Dict with ``ece``, ``n``, and per-bin ``bin_confidence`` / ``bin_accuracy`` + lists (empty bins omitted from the lists but counted in normalization). + """ + conf = confidence.detach().float().reshape(-1).clamp(0.0, 1.0) + corr = correctness.detach().float().reshape(-1) + if conf.numel() != corr.numel(): + raise ValueError("confidence and correctness must have the same number of elements") + if conf.numel() == 0: + return {"ece": 0.0, "n": 0, "bin_confidence": [], "bin_accuracy": []} + + bin_edges = torch.linspace(0.0, 1.0, n_bins + 1, device=conf.device) + ece = conf.new_zeros(()) + bin_conf: list[float] = [] + bin_acc: list[float] = [] + n = conf.numel() + for i in range(n_bins): + lo, hi = bin_edges[i], bin_edges[i + 1] + if i == n_bins - 1: + mask = (conf >= lo) & (conf <= hi) + else: + mask = (conf >= lo) & (conf < hi) + count = int(mask.sum().item()) + if count == 0: + continue + avg_conf = conf[mask].mean() + avg_acc = corr[mask].mean() + ece = ece + (count / n) * (avg_conf - avg_acc).abs() + bin_conf.append(float(avg_conf)) + bin_acc.append(float(avg_acc)) + return { + "ece": float(ece), + "n": n, + "bin_confidence": bin_conf, + "bin_accuracy": bin_acc, + } diff --git a/Model/evaluation/confidence_coupling_ab.py b/Model/evaluation/confidence_coupling_ab.py new file mode 100644 index 000000000..4f914cf24 --- /dev/null +++ b/Model/evaluation/confidence_coupling_ab.py @@ -0,0 +1,177 @@ +"""A/B: confidence-scaled coupling vs unscaled (always-1) (#110 / PR #193). + +Trains two Combined-style mock models on the same batch: + * ``scaled`` — current PR (reasoning confidence multiplies the residual) + * ``unscaled`` — pre-#110 (confidence forced to 1, residual always full) + +Reports open-loop ADE/FDE vs the imitation target after the same step count, +plus a post-hoc sweep of confidence 0 vs 1 on the scaled model (the safety +loop: low confidence should sit closer to the unmodulated plan). +""" + +from __future__ import annotations + +from typing import Any, Callable + +import numpy as np +import torch + +from evaluation.metrics import compute_open_loop_metrics +from training.losses.horizon_reasoning_loss import HorizonReasoningLoss +from data_processing.reasoning_label_generation.mock_teacher import MockTeacher +from data_processing.reasoning_label_generation.targets import ( + collate_reasoning_targets, + record_to_target_tensors, +) +from data_processing.reasoning_label_generation.teacher_client import TeacherRequest + + +def _split_controls(traj: torch.Tensor) -> tuple[np.ndarray, np.ndarray]: + 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] % 2 == 0: + paired = t.reshape(t.shape[0], -1, 2) + return paired[..., 0], paired[..., 1] + raise ValueError(f"unrecognised trajectory shape {t.shape}") + + +def ade_fde_vs_target(pred: torch.Tensor, target: torch.Tensor, speed: float = 5.0) -> dict[str, float]: + pa, pc = _split_controls(pred) + ta, tc = _split_controls(target) + v0 = np.full(pa.shape[0], speed) + return compute_open_loop_metrics(pa, pc, ta, tc, v0) + + +def _patch_forced_confidence(model: torch.nn.Module, value: torch.Tensor | None) -> Callable[[], None]: + planner = model.Reactive_E2E.TrajectoryPlanner + orig = planner.forward + + def wrapped(*args: Any, **kwargs: Any): + kwargs["reasoning_confidence"] = value + return orig(*args, **kwargs) + + planner.forward = wrapped # type: ignore[method-assign] + + def restore() -> None: + planner.forward = orig # type: ignore[method-assign] + + return restore + + +def _batch(device: torch.device, b: int = 2, v: int = 6) -> dict[str, torch.Tensor]: + return { + "visual": torch.randn(b, v, 3, 256, 256, device=device), + "map_input": torch.randn(b, 3, 256, 256, device=device), + "vis_hist": torch.zeros(b, 896, device=device), + "ego": torch.randn(b, 256, device=device), + "target": torch.randn(b, 128, device=device), + } + + +def _reasoning_targets(b: int): + teacher = MockTeacher() + per = [record_to_target_tensors(teacher.label(TeacherRequest(f"s{i}", "l2d"))) for i in range(b)] + return collate_reasoning_targets(per) + + +def _train(model: torch.nn.Module, batch: dict[str, torch.Tensor], steps: int, lr: float) -> list[float]: + model.train() + opt = torch.optim.AdamW(model.parameters(), lr=lr) + traj_loss_fn = torch.nn.SmoothL1Loss() + reason_loss_fn = HorizonReasoningLoss() + tb = _reasoning_targets(batch["visual"].shape[0]) + history: list[float] = [] + for _ in range(steps): + opt.zero_grad(set_to_none=True) + traj, aux_or_pred = model( + batch["visual"], batch["map_input"], batch["vis_hist"], batch["ego"], + mode="train", trajectory_target=batch["target"], + ) + # ReactiveE2E returns (traj, reasoning_pred) when reasoning is on. + if isinstance(aux_or_pred, dict): + pred = aux_or_pred["reasoning_pred"] + else: + pred = aux_or_pred + loss = traj_loss_fn(traj, batch["target"]) + terms = reason_loss_fn(pred, tb.targets, source_weights=tb.source_weights, + confidence_targets=tb.confidence_targets) + total = loss + 0.5 * terms["total"] + total.backward() + opt.step() + history.append(float(total.detach())) + return history + + +@torch.no_grad() +def _eval_ade(model: torch.nn.Module, batch: dict[str, torch.Tensor]) -> dict[str, float]: + model.eval() + out = model( + batch["visual"], batch["map_input"], batch["vis_hist"], batch["ego"], + mode="infer", + ) + traj = out[0] if isinstance(out, tuple) else out + return ade_fde_vs_target(traj, batch["target"]) + + +def run_confidence_coupling_ab( + build_model, + *, + device: torch.device | None = None, + steps: int = 10, + lr: float = 1e-3, + seed: int = 0, +) -> dict[str, Any]: + """Train scaled vs unscaled coupling; return ADE/FDE and confidence sweep.""" + device = device or torch.device("cpu") + torch.manual_seed(seed) + batch = _batch(device) + b = batch["visual"].shape[0] + + scaled = build_model( + num_views=6, device=device, + enable_reasoning=True, reasoning_mode="pooled_latent", + ) + hist_s = _train(scaled, batch, steps, lr) + metrics_s = _eval_ade(scaled, batch) + + torch.manual_seed(seed) + batch_u = _batch(device) + unscaled = build_model( + num_views=6, device=device, + enable_reasoning=True, reasoning_mode="pooled_latent", + ) + ones = torch.ones(b, device=device) + restore = _patch_forced_confidence(unscaled, ones) + try: + hist_u = _train(unscaled, batch_u, steps, lr) + metrics_u = _eval_ade(unscaled, batch_u) + finally: + restore() + + # Safety-loop sweep on the scaled model: conf=0 vs conf=1. + zeros = torch.zeros(b, device=device) + r0 = _patch_forced_confidence(scaled, zeros) + try: + ade_c0 = _eval_ade(scaled, batch) + finally: + r0() + r1 = _patch_forced_confidence(scaled, ones) + try: + ade_c1 = _eval_ade(scaled, batch) + finally: + r1() + + return { + "train_steps": steps, + "scaled": {**metrics_s, "loss_first": hist_s[0], "loss_last": hist_s[-1]}, + "unscaled": {**metrics_u, "loss_first": hist_u[0], "loss_last": hist_u[-1]}, + "scaled_conf0": ade_c0, + "scaled_conf1": ade_c1, + "ade3s_delta_scaled_minus_unscaled": ( + metrics_s["ADE@3s"] - metrics_u["ADE@3s"] + ), + "ade3s_delta_conf1_minus_conf0": ( + ade_c1["ADE@3s"] - ade_c0["ADE@3s"] + ), + } diff --git a/Model/evaluation/results/confidence_coupling_ab.json b/Model/evaluation/results/confidence_coupling_ab.json new file mode 100644 index 000000000..63a4cba21 --- /dev/null +++ b/Model/evaluation/results/confidence_coupling_ab.json @@ -0,0 +1,41 @@ +{ + "train_steps": 12, + "scaled": { + "ADE@1s": 1.023074085267849, + "ADE@2s": 1.584022108793444, + "ADE@3s": 2.3926888108850024, + "FDE@3s": 5.20447055848771, + "accel_mae": 0.8210745453834534, + "curvature_mae": 0.6899887919425964, + "loss_first": 4.032444000244141, + "loss_last": 0.9783587455749512 + }, + "unscaled": { + "ADE@1s": 1.0230650247206414, + "ADE@2s": 1.5840292634637652, + "ADE@3s": 2.392811604041399, + "FDE@3s": 5.2049962177061, + "accel_mae": 0.8210744857788086, + "curvature_mae": 0.6899884343147278, + "loss_first": 4.032444000244141, + "loss_last": 0.9783588647842407 + }, + "scaled_conf0": { + "ADE@1s": 1.0230626708395694, + "ADE@2s": 1.5839559284150535, + "ADE@3s": 2.3927409543209825, + "FDE@3s": 5.205003511023582, + "accel_mae": 0.821074366569519, + "curvature_mae": 0.6899888515472412 + }, + "scaled_conf1": { + "ADE@1s": 1.02307659993299, + "ADE@2s": 1.5840455097031116, + "ADE@3s": 2.3926663229586778, + "FDE@3s": 5.20426478376657, + "accel_mae": 0.8210746049880981, + "curvature_mae": 0.689988911151886 + }, + "ade3s_delta_scaled_minus_unscaled": -0.00012279315639673882, + "ade3s_delta_conf1_minus_conf0": -7.463136230478895e-05 +} \ No newline at end of file diff --git a/Model/model_components/reactive_e2e.py b/Model/model_components/reactive_e2e.py index f6f4ba528..2e87178f8 100644 --- a/Model/model_components/reactive_e2e.py +++ b/Model/model_components/reactive_e2e.py @@ -229,18 +229,27 @@ def validity_gate(value, valid, *, default, name): reasoning_pred = None reasoning_latent = None reasoning_horizon_tokens = None + reasoning_confidence = None if self.ReasoningHead is not None: reasoning_pred = self.ReasoningHead( visual_ctx, ego_ctx, ) reasoning_latent = reasoning_pred.reasoning_latent reasoning_horizon_tokens = reasoning_pred.horizon_tokens + # Pool per-horizon confidence logits → [B] probabilities for the gate (#110). + from model_components.trajectory_planning.reasoning_coupling import ( + pool_reasoning_confidence, + ) + reasoning_confidence = pool_reasoning_confidence( + reasoning_pred.confidence_logits, from_logits=True + ) # --- Trajectory Prediction --- trajectory = self.TrajectoryPlanner( fused_features, visual_ctx, ego_ctx, reasoning_latent=reasoning_latent, reasoning_horizon_tokens=reasoning_horizon_tokens, + reasoning_confidence=reasoning_confidence, **kwargs, ) diff --git a/Model/model_components/trajectory_planning/bezier_planner.py b/Model/model_components/trajectory_planning/bezier_planner.py index 9c1b29d2e..285f7cc83 100644 --- a/Model/model_components/trajectory_planning/bezier_planner.py +++ b/Model/model_components/trajectory_planning/bezier_planner.py @@ -106,6 +106,7 @@ def _bernstein_basis(num_points, num_controls): def forward(self, bev_features, visual_history, egomotion_history, reasoning_latent=None, reasoning_horizon_tokens=None, + reasoning_confidence=None, **kwargs): """ Args: @@ -116,6 +117,8 @@ def forward(self, bev_features, visual_history, egomotion_history, (used by reasoning_mode="pooled_latent"). reasoning_horizon_tokens: optional [B, 5, embed_dim] per-horizon reasoning tokens (used by reasoning_mode="horizon_cross_attention"). + reasoning_confidence: optional [B] (or [B,5] / logits) confidence + that scales the reasoning residual (#110). Returns: trajectory: [B, num_timesteps * num_signals] @@ -146,6 +149,7 @@ def forward(self, bev_features, visual_history, egomotion_history, context, reasoning_latent=reasoning_latent, horizon_tokens=reasoning_horizon_tokens, + confidence=reasoning_confidence, ) bezier_feature = self.context_mlp(context) # [B, C] diff --git a/Model/model_components/trajectory_planning/flow_matching_planner.py b/Model/model_components/trajectory_planning/flow_matching_planner.py index c69849fec..c9af26075 100644 --- a/Model/model_components/trajectory_planning/flow_matching_planner.py +++ b/Model/model_components/trajectory_planning/flow_matching_planner.py @@ -223,7 +223,7 @@ def _sinusoidal_time_embedding(self, t): return torch.cat([torch.sin(args), torch.cos(args)], dim=-1) def _modulation_conditioning(self, visual_history, egomotion_history, - reasoning_latent=None): + reasoning_latent=None, reasoning_confidence=None): """Conditioning vector fed into AdaLN — excludes BEV (cross-attn) and time (added per-step in the Euler loop). @@ -238,7 +238,9 @@ def _modulation_conditioning(self, visual_history, egomotion_history, ) # Only the pooled path uses mod_cond; horizon tokens are routed to _v_theta. latent = reasoning_latent if self.reasoning_mode == "pooled_latent" else None - return self.reasoning_coupling(base, reasoning_latent=latent) + return self.reasoning_coupling( + base, reasoning_latent=latent, confidence=reasoning_confidence + ) def _sample_timesteps(self, batch_size, device, dtype): @@ -290,7 +292,8 @@ def _project_bev(self, bev_features): bev_seq = bev_features.flatten(2).transpose(1, 2) return self.bev_kv_proj(bev_seq) - def _v_theta(self, u_t, t, bev_seq, mod_cond, horizon_tokens=None): + def _v_theta(self, u_t, t, bev_seq, mod_cond, horizon_tokens=None, + reasoning_confidence=None): """Conditional velocity network with BEV cross-attention + AdaLN. Args: @@ -304,6 +307,7 @@ def _v_theta(self, u_t, t, bev_seq, mod_cond, horizon_tokens=None): "horizon_cross_attention" mode each action query attends these (zero-init, no-op until trained) so timestep-t actions can look at the now/1s/2s/3s/4s reasoning; ignored in other modes. + reasoning_confidence: optional confidence scale for the residual. Returns: velocity: [B, trajectory_dim] @@ -320,7 +324,10 @@ def _v_theta(self, u_t, t, bev_seq, mod_cond, horizon_tokens=None): # a single pooled vector would lose. if self.reasoning_mode == "horizon_cross_attention" and horizon_tokens is not None: queries = self.reasoning_coupling( - queries, horizon_tokens=horizon_tokens, query=queries + queries, + horizon_tokens=horizon_tokens, + query=queries, + confidence=reasoning_confidence, ) attended, _ = self.cross_attn(queries, bev_seq, bev_seq) # [B, T, C] @@ -336,7 +343,7 @@ def _v_theta(self, u_t, t, bev_seq, mod_cond, horizon_tokens=None): def forward(self, bev_features, visual_history, egomotion_history, generator=None, initial_noise=None, reasoning_latent=None, - reasoning_horizon_tokens=None, **kwargs): + reasoning_horizon_tokens=None, reasoning_confidence=None, **kwargs): """Inference: Euler-integrate ``dx/dt = v_theta(x, t, ...)`` over [0, 1]. Args: @@ -353,6 +360,8 @@ def forward(self, bev_features, visual_history, egomotion_history, (reasoning_mode="pooled_latent"). reasoning_horizon_tokens: optional [B, 5, embed_dim] per-horizon reasoning tokens (reasoning_mode="horizon_cross_attention"). + reasoning_confidence: optional confidence that scales the reasoning + residual (#110). **kwargs: ignored. Accepts extra inputs other planners or callers might pass so call sites can stay planner-agnostic. @@ -363,6 +372,7 @@ def forward(self, bev_features, visual_history, egomotion_history, mod_cond = self._modulation_conditioning( visual_history, egomotion_history, reasoning_latent=reasoning_latent, + reasoning_confidence=reasoning_confidence, ) # bev_seq is computed once and reused across every Euler step. bev_seq = self._project_bev(bev_features) @@ -384,7 +394,10 @@ def forward(self, bev_features, visual_history, egomotion_history, t_val = step * dt t = torch.full((B,), t_val, device=bev_features.device, dtype=bev_features.dtype) - v = self._v_theta(x, t, bev_seq, mod_cond, - horizon_tokens=reasoning_horizon_tokens) + v = self._v_theta( + x, t, bev_seq, mod_cond, + horizon_tokens=reasoning_horizon_tokens, + reasoning_confidence=reasoning_confidence, + ) x = x + dt * v return x diff --git a/Model/model_components/trajectory_planning/reasoning_coupling.py b/Model/model_components/trajectory_planning/reasoning_coupling.py index 3efa6c157..b26b6f72b 100644 --- a/Model/model_components/trajectory_planning/reasoning_coupling.py +++ b/Model/model_components/trajectory_planning/reasoning_coupling.py @@ -1,21 +1,25 @@ -"""Zero-init reasoning→planner coupling (issue #98, R7). +"""Zero-init reasoning→planner coupling (issue #98, R7; confidence loop #110). Injects the reasoning branch's output into a planner conditioning vector behind a zero-initialised gate, so at initialisation the coupling is a strict no-op and the reactive baseline is byte-identical up to numerical tolerance. Training moves the gate away from zero only where reasoning helps the trajectory. +Confidence (#110): when ``confidence`` (probabilities in ``[0, 1]``) is provided, +the residual is scaled by it so **low confidence → weaker reasoning modulation** +(more conservative / closer to the reactive baseline). At init ``alpha=0``, so the +gate remains a strict no-op for any confidence value. + Three modes (the required ablation surface A/B/C): * ``none`` — coupling disabled; the planner is unchanged. - * ``pooled_latent`` — add ``alpha * reason_proj(reasoning_latent)``. + * ``pooled_latent`` — add ``alpha * conf * reason_proj(reasoning_latent)``. * ``horizon_cross_attention`` — a query attends the 5 horizon tokens, then - ``alpha * reason_proj(attended)`` is added — preserving *when* a hazard + ``alpha * conf * reason_proj(attended)`` is added — preserving *when* a hazard matters. ``alpha`` is a learned scalar initialised to 0 (the repo's ResidualMapFusion / -#108 ZeroInitGate pattern); ``reason_proj``'s final layer is also zero-init as a -belt-and-braces guarantee that the residual is exactly 0 at init regardless of -the attention output. +#108 ZeroInitGate pattern); ``reason_proj`` keeps normal init so alpha receives +gradient at init while the residual stays exactly 0. """ from __future__ import annotations @@ -28,6 +32,33 @@ REASONING_MODES = ("none", "pooled_latent", "horizon_cross_attention") +def pool_reasoning_confidence( + confidence: torch.Tensor, + *, + from_logits: bool = False, +) -> torch.Tensor: + """Reduce per-horizon confidence to a per-batch scale ``[B]`` in ``[0, 1]``. + + Accepts ``[B]``, ``[B, 1]``, ``[B, H]``, or ``[B, H, 1]``. When + ``from_logits`` is True, applies ``sigmoid`` first. + """ + x = confidence + if from_logits: + x = torch.sigmoid(x) + if x.dim() == 3: + x = x.squeeze(-1) # [B, H] + if x.dim() == 2: + if x.shape[-1] == 1: + x = x.squeeze(-1) + else: + x = x.mean(dim=-1) # pool horizons + if x.dim() != 1: + raise ValueError( + f"confidence must reduce to [B]; got shape {tuple(confidence.shape)}" + ) + return x.clamp(0.0, 1.0) + + class ReasoningCoupling(nn.Module): """Add a zero-init reasoning residual to a planner conditioning vector. @@ -39,11 +70,10 @@ class ReasoningCoupling(nn.Module): Forward: coupling(context[B,D], reasoning_latent=None, horizon_tokens=None, - query=None) -> context'[B,D] + query=None, confidence=None) -> context'[B,D] With ``mode="none"`` (or missing reasoning inputs) it returns ``context`` - unchanged. In ``horizon_cross_attention`` mode the attention query is - ``query`` if given (e.g. the flow-matching action tokens), else the - context vector itself. + unchanged. ``confidence`` is an optional ``[B]`` (or broadcastable) + probability that scales the residual; low confidence → conservative. """ def __init__(self, embed_dim: int = 256, mode: str = "none", num_heads: int = 4) -> None: @@ -76,12 +106,23 @@ def __init__(self, embed_dim: int = 256, mode: str = "none", num_heads: int = 4) embed_dim, num_heads, dropout=0.0, batch_first=True ) + def _scale_residual(self, delta: torch.Tensor, confidence: Optional[torch.Tensor]) -> torch.Tensor: + """Apply optional confidence scale; detach so planner loss does not train conf.""" + if confidence is None: + return delta + scale = pool_reasoning_confidence(confidence) + # Broadcast [B] onto delta's trailing dims. + while scale.dim() < delta.dim(): + scale = scale.unsqueeze(-1) + return delta * scale.detach() + def forward( self, context: torch.Tensor, reasoning_latent: Optional[torch.Tensor] = None, horizon_tokens: Optional[torch.Tensor] = None, query: Optional[torch.Tensor] = None, + confidence: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Return the reasoning-conditioned context (unchanged if mode='none').""" if self.mode == "none": @@ -91,6 +132,7 @@ def forward( if reasoning_latent is None: return context # no reasoning available this step → no-op delta = self.reason_proj(reasoning_latent) # [B, D] + delta = self._scale_residual(delta, confidence) return context + self.alpha * delta # horizon_cross_attention @@ -99,6 +141,7 @@ def forward( q = query if query is not None else context.unsqueeze(1) # [B, Tq, D] attended, _ = self.cross_attn(q, horizon_tokens, horizon_tokens) # [B, Tq, D] delta = self.reason_proj(attended) # [B, Tq, D] + delta = self._scale_residual(delta, confidence) gated = self.alpha * delta # Broadcast back onto the caller's context shape: a single-vector context # gets the squeezed residual; a per-token query context keeps its tokens. diff --git a/Model/tests/test_reasoning_coupling.py b/Model/tests/test_reasoning_coupling.py index 0b62ee9d7..6f0087b57 100644 --- a/Model/tests/test_reasoning_coupling.py +++ b/Model/tests/test_reasoning_coupling.py @@ -11,6 +11,7 @@ from __future__ import annotations +import numpy as np import pytest import torch @@ -158,3 +159,62 @@ def test_flow_matching_is_horizon_aware_not_pooled(): b = planner(bev, vis, ego, generator=g2, reasoning_horizon_tokens=tokens_b) assert not torch.allclose(a, b, atol=1e-5), \ "zeroing one horizon left the trajectory unchanged — timing info is lost" + + +def test_confidence_scales_residual_and_preserves_zero_init(): + """#110: low confidence → weaker residual; any confidence is still no-op at init.""" + from model_components.trajectory_planning.reasoning_coupling import ( + pool_reasoning_confidence, + ) + + torch.manual_seed(0) + c = ReasoningCoupling(EMBED, mode="pooled_latent") + ctx = torch.randn(B, EMBED) + latent = torch.randn(B, EMBED) + high = torch.ones(B) + low = torch.full((B,), 0.1) + + # Strict no-op at init for any confidence. + with torch.no_grad(): + out_high = c(ctx, reasoning_latent=latent, confidence=high) + out_low = c(ctx, reasoning_latent=latent, confidence=low) + assert torch.allclose(out_high, ctx, atol=1e-6) + assert torch.allclose(out_low, ctx, atol=1e-6) + + # Open the gate: high confidence should move context more than low. + with torch.no_grad(): + c.alpha.fill_(1.0) + c.reason_proj[-1].weight.normal_() + moved_high = c(ctx, reasoning_latent=latent, confidence=high) + moved_low = c(ctx, reasoning_latent=latent, confidence=low) + dist_high = (moved_high - ctx).norm() + dist_low = (moved_low - ctx).norm() + assert dist_high > dist_low + + pooled = pool_reasoning_confidence(torch.tensor([[0.2, 0.4, 0.6, 0.8, 1.0]])) + assert pooled.shape == (1,) + assert float(pooled) == pytest.approx(0.6) + + +def test_confidence_coupling_ab_reports_ade(build_mock_model, device): + from evaluation.confidence_coupling_ab import run_confidence_coupling_ab + + report = run_confidence_coupling_ab( + build_mock_model, device=device, steps=6, seed=0, + ) + assert report["scaled"]["loss_last"] < report["scaled"]["loss_first"] + assert np.isfinite(report["scaled"]["ADE@3s"]) + assert np.isfinite(report["unscaled"]["ADE@3s"]) + assert np.isfinite(report["scaled_conf0"]["ADE@3s"]) + assert np.isfinite(report["scaled_conf1"]["ADE@3s"]) + + +def test_expected_calibration_error_perfect_and_bad(): + from evaluation.confidence_calibration import expected_calibration_error + + conf = torch.tensor([0.1, 0.2, 0.8, 0.9]) + perfect = expected_calibration_error(conf, conf) + assert perfect["ece"] == pytest.approx(0.0, abs=1e-6) + + bad = expected_calibration_error(conf, 1.0 - conf) + assert bad["ece"] > 0.3