From be9a48fa0df613903330929f844b97ef814179c6 Mon Sep 17 00:00:00 2001 From: ShauryaVM Date: Wed, 12 Aug 2026 14:54:21 -0400 Subject: [PATCH 1/3] Add fidelity-aware reward prototype for stage-3 RL (#123). Gate base (and optional consequence) rewards by world-model prediction fidelity so the objective only trusts information the reactive policy does not already observe. Co-authored-by: Cursor --- Model/tests/test_fidelity_aware_reward.py | 78 +++++++ Model/training/losses/__init__.py | 14 ++ .../training/losses/fidelity_aware_reward.py | 213 ++++++++++++++++++ 3 files changed, 305 insertions(+) create mode 100644 Model/tests/test_fidelity_aware_reward.py create mode 100644 Model/training/losses/fidelity_aware_reward.py diff --git a/Model/tests/test_fidelity_aware_reward.py b/Model/tests/test_fidelity_aware_reward.py new file mode 100644 index 000000000..491b58252 --- /dev/null +++ b/Model/tests/test_fidelity_aware_reward.py @@ -0,0 +1,78 @@ +"""Unit tests for the fidelity-aware reward prototype (#123).""" + +from __future__ import annotations + +import torch + +from training.losses.fidelity_aware_reward import ( + FIDELITY_AWARE_REWARD_VERSION, + consequence_alignment_reward, + fidelity_aware_reward, + soft_advantage_from_reward, + world_model_fidelity, +) + + +def test_version_pin() -> None: + assert FIDELITY_AWARE_REWARD_VERSION == "wm_fidelity_gate_v1" + + +def test_perfect_wm_fidelity_is_one() -> None: + x = torch.randn(4, 8, 8) + fid = world_model_fidelity(x, x.clone()) + assert torch.allclose(fid, torch.ones_like(fid)) + + +def test_error_lowers_fidelity() -> None: + pred = torch.zeros(2, 4) + target = torch.ones(2, 4) + fid = world_model_fidelity(pred, target, temperature=1.0) + assert torch.all(fid < 1.0) + assert torch.all(fid > 0.0) + + +def test_fidelity_gates_base_reward() -> None: + base = torch.tensor([1.0, 1.0]) + # Sample 0: perfect WM; sample 1: large error → low fidelity. + wm_pred = torch.tensor([[0.0, 0.0], [0.0, 0.0]]) + wm_tgt = torch.tensor([[0.0, 0.0], [10.0, 10.0]]) + out = fidelity_aware_reward( + base, + wm_prediction=wm_pred, + wm_target=wm_tgt, + fidelity_temperature=1.0, + ) + assert out.reward[0].item() == 1.0 + assert out.reward[1].item() < out.reward[0].item() + assert out.metadata["reward_mean"] < 1.0 + + +def test_consequence_term_requires_external_preference() -> None: + base = torch.zeros(2) + wm = torch.zeros(2, 3) + pred_fut = torch.zeros(2, 3) + pref = torch.ones(2, 3) + out = fidelity_aware_reward( + base, + wm_prediction=wm, + wm_target=wm.clone(), + predicted_future=pred_fut, + preferred_future=pref, + consequence_weight=1.0, + consequence_scale=1.0, + ) + # Perfect WM fidelity=1, consequence = -mse = -1 → reward -1. + assert torch.allclose(out.reward, torch.tensor([-1.0, -1.0])) + assert out.consequence_reward is not None + + +def test_consequence_alignment_reward_shapes() -> None: + r = consequence_alignment_reward(torch.zeros(3, 2), torch.ones(3, 2), scale=2.0) + assert r.shape == (3,) + assert torch.all(r < 0) + + +def test_soft_advantage_zero_mean_without_baseline() -> None: + reward = torch.tensor([1.0, 3.0]) + adv = soft_advantage_from_reward(reward) + assert abs(adv.mean().item()) < 1e-6 diff --git a/Model/training/losses/__init__.py b/Model/training/losses/__init__.py index 2160f633e..7c74524c8 100644 --- a/Model/training/losses/__init__.py +++ b/Model/training/losses/__init__.py @@ -14,14 +14,28 @@ RouteConsistencyWeights, ego_points_to_grid, ) +from .fidelity_aware_reward import ( + FIDELITY_AWARE_REWARD_VERSION, + FidelityAwareRewardResult, + consequence_alignment_reward, + fidelity_aware_reward, + soft_advantage_from_reward, + world_model_fidelity, +) __all__ = [ + "FIDELITY_AWARE_REWARD_VERSION", + "FidelityAwareRewardResult", "HorizonReasoningLoss", "ROLLOUT_ALIGNED_LOSS_VERSION", "ROLLOUT_POLICY_VERSION", "RouteConsistencyLoss", "RouteConsistencyWeights", "RolloutAlignedLoss", + "consequence_alignment_reward", "ego_points_to_grid", + "fidelity_aware_reward", "integrate_controls_torch", + "soft_advantage_from_reward", + "world_model_fidelity", ] diff --git a/Model/training/losses/fidelity_aware_reward.py b/Model/training/losses/fidelity_aware_reward.py new file mode 100644 index 000000000..2ce58dbed --- /dev/null +++ b/Model/training/losses/fidelity_aware_reward.py @@ -0,0 +1,213 @@ +"""Fidelity-aware reward prototype for stage-3 closed-loop RL (issue #123). + +Design rule from the issue: every reward term must add information the +reactive policy cannot already infer from its own inputs. A reward built +only from the reasoning band (a re-encoding of planner inputs) is redundant +for the same DPI reason that drove ``reasoning_coupling.alpha → 0``. + +This module therefore shapes rewards with **world-model fidelity** — how +well the WM predicts consequences of the rolled-out action — and optionally +scores those predicted consequences against an external preference signal. +Both are information the camera-only reactive policy does not observe at +decision time. + +This is deliberately a pure function / small helper, not a training-loop +integration. Wire it behind ``compute_planner_loss`` (#115) when that hook +lands; until then it is safe to unit-test and prototype against. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + + +FIDELITY_AWARE_REWARD_VERSION = "wm_fidelity_gate_v1" + + +@dataclass(frozen=True) +class FidelityAwareRewardResult: + """Per-sample reward and diagnostic scalars.""" + + reward: torch.Tensor + fidelity: torch.Tensor + base_reward: torch.Tensor + consequence_reward: torch.Tensor | None + metadata: dict[str, float] + + +def world_model_fidelity( + prediction: torch.Tensor, + target: torch.Tensor, + *, + temperature: float = 1.0, + reduce_dims: tuple[int, ...] | None = None, +) -> torch.Tensor: + """Map WM prediction error to a ``(0, 1]`` fidelity weight. + + ``fidelity = exp(-mse / temperature)``. High fidelity means the world + model is a trustworthy source of consequence information for this sample; + low fidelity down-weights any reward that depends on those predictions. + """ + if prediction.shape != target.shape: + raise ValueError( + f"prediction/target shape mismatch: {tuple(prediction.shape)} vs " + f"{tuple(target.shape)}" + ) + if temperature <= 0: + raise ValueError(f"temperature must be > 0, got {temperature}") + + err = (prediction.float() - target.float()).pow(2) + if reduce_dims is None: + # Reduce all but batch dim when present. + if err.ndim == 0: + mse = err + else: + mse = err.flatten(1).mean(dim=1) + else: + mse = err.mean(dim=reduce_dims) + + return torch.exp(-mse / temperature) + + +def consequence_alignment_reward( + predicted_future: torch.Tensor, + preferred_future: torch.Tensor, + *, + scale: float = 1.0, +) -> torch.Tensor: + """Reward from WM-predicted futures vs an external preference target. + + ``preferred_future`` is intended to carry information outside the + reactive policy's inputs (e.g. HD-map / LiDAR conflict geometry, or a + human preference embedding) — not a re-encoding of BEV/ego already fed + to the planner. + """ + if predicted_future.shape != preferred_future.shape: + raise ValueError( + "predicted_future/preferred_future shape mismatch: " + f"{tuple(predicted_future.shape)} vs {tuple(preferred_future.shape)}" + ) + if scale <= 0: + raise ValueError(f"scale must be > 0, got {scale}") + + # Negative MSE so better alignment → higher reward. Keep per-batch. + if predicted_future.ndim == 0: + mse = (predicted_future.float() - preferred_future.float()).pow(2) + else: + mse = ( + (predicted_future.float() - preferred_future.float()) + .pow(2) + .flatten(1) + .mean(dim=1) + ) + return -scale * mse + + +def fidelity_aware_reward( + base_reward: torch.Tensor, + *, + wm_prediction: torch.Tensor, + wm_target: torch.Tensor, + preferred_future: torch.Tensor | None = None, + predicted_future: torch.Tensor | None = None, + fidelity_temperature: float = 1.0, + consequence_scale: float = 1.0, + consequence_weight: float = 1.0, + min_fidelity: float = 0.0, +) -> FidelityAwareRewardResult: + """Gate ``base_reward`` (and optional consequence term) by WM fidelity. + + Final reward:: + + r = fidelity * (base_reward + consequence_weight * consequence) + + where ``fidelity`` is clipped below by ``min_fidelity`` so a broken WM + cannot zero the entire objective if a floor is desired. + """ + if base_reward.ndim > 1: + raise ValueError( + f"base_reward must be scalar or (B,), got shape {tuple(base_reward.shape)}" + ) + + fidelity = world_model_fidelity( + wm_prediction, + wm_target, + temperature=fidelity_temperature, + ) + if fidelity.shape != base_reward.shape and not ( + fidelity.ndim == 1 and base_reward.ndim == 0 + ): + # Allow broadcasting scalar base over batch fidelity. + if base_reward.ndim == 0 and fidelity.ndim == 1: + base_reward = base_reward.expand_as(fidelity) + elif fidelity.ndim == 0 and base_reward.ndim == 1: + fidelity = fidelity.expand_as(base_reward) + elif fidelity.shape != base_reward.shape: + raise ValueError( + "fidelity/base_reward shape mismatch after reduction: " + f"{tuple(fidelity.shape)} vs {tuple(base_reward.shape)}" + ) + + fidelity = fidelity.clamp(min=float(min_fidelity), max=1.0) + + consequence: torch.Tensor | None = None + total = base_reward.float() + if preferred_future is not None: + if predicted_future is None: + raise ValueError( + "predicted_future is required when preferred_future is provided" + ) + consequence = consequence_alignment_reward( + predicted_future, + preferred_future, + scale=consequence_scale, + ) + if consequence.shape != total.shape: + raise ValueError( + "consequence/base_reward shape mismatch: " + f"{tuple(consequence.shape)} vs {tuple(total.shape)}" + ) + total = total + float(consequence_weight) * consequence + + reward = fidelity * total + metadata = { + "fidelity_mean": float(fidelity.detach().mean().cpu()), + "base_reward_mean": float(base_reward.detach().float().mean().cpu()), + "reward_mean": float(reward.detach().mean().cpu()), + "fidelity_temperature": float(fidelity_temperature), + "consequence_weight": float(consequence_weight), + "min_fidelity": float(min_fidelity), + } + if consequence is not None: + metadata["consequence_reward_mean"] = float( + consequence.detach().mean().cpu() + ) + + return FidelityAwareRewardResult( + reward=reward, + fidelity=fidelity, + base_reward=base_reward.float(), + consequence_reward=consequence, + metadata=metadata, + ) + + +def soft_advantage_from_reward( + reward: torch.Tensor, + *, + baseline: torch.Tensor | None = None, + temperature: float = 1.0, +) -> torch.Tensor: + """Optional REINFORCE-style advantage: ``(r - baseline) / temperature``. + + Kept separate from :func:`fidelity_aware_reward` so callers can plug the + gated reward into whatever policy-gradient estimator they choose once + ``compute_planner_loss`` (#115) is restored. + """ + if temperature <= 0: + raise ValueError(f"temperature must be > 0, got {temperature}") + if baseline is None: + baseline = reward.detach().mean() + return (reward - baseline) / temperature From 94b35d56edea254174d0856c91ff686f7fe6d02e Mon Sep 17 00:00:00 2001 From: ShauryaVM Date: Wed, 19 Aug 2026 01:07:50 -0400 Subject: [PATCH 2/3] Gate only the world-model term; rank expert vs jerky. Co-authored-by: Cursor #123 v1 uses rollout comfort/progress as the base and multiplies fidelity onto the WM consequence alone, so a noisy world model cannot wipe safety or flip the ranking. The offline pair is the tensor AlpaSim (#177) should call later. --- .../evaluation/fidelity_reward_experiment.py | 122 ++++++++++++++++++ .../results/fidelity_reward_experiment.json | 29 +++++ Model/tests/test_fidelity_aware_reward.py | 23 +++- Model/training/losses/__init__.py | 4 + .../training/losses/fidelity_aware_reward.py | 94 ++++++++++---- 5 files changed, 238 insertions(+), 34 deletions(-) create mode 100644 Model/evaluation/fidelity_reward_experiment.py create mode 100644 Model/evaluation/results/fidelity_reward_experiment.json diff --git a/Model/evaluation/fidelity_reward_experiment.py b/Model/evaluation/fidelity_reward_experiment.py new file mode 100644 index 000000000..6d290c5c1 --- /dev/null +++ b/Model/evaluation/fidelity_reward_experiment.py @@ -0,0 +1,122 @@ +"""Offline #123 experiment: expert vs jerky ranking under a WM fidelity gate. + +Handcrafted safety/comfort/progress should prefer the expert plan. A +*misleading* WM consequence term would prefer the jerky plan if left ungated. +When the WM is noise, ``g ≈ 0`` so the ranking must match the handcrafted +base — the gate AlpaSim (#177) should call, without importing that PR. +""" + +from __future__ import annotations + +from typing import Any + +import torch + +from training.losses.fidelity_aware_reward import ( + FIDELITY_AWARE_REWARD_VERSION, + fidelity_aware_reward, + v1_handcrafted_reward, +) + + +def _pair(timesteps: int = 32) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + t = torch.arange(timesteps, dtype=torch.float32) + expert = torch.zeros(1, timesteps, 2) + jerky = torch.stack( + ( + 8.0 * torch.sin(2.0 * torch.pi * t / 2.0), + 0.25 * torch.sin(2.0 * torch.pi * t / 3.0), + ), + dim=-1, + ).unsqueeze(0) + controls = torch.cat([expert, jerky], dim=0) + expert_batch = expert.expand(2, -1, -1).contiguous() + v0 = torch.full((2,), 5.0) + return controls, expert_batch, v0 + + +def run_fidelity_reward_experiment(*, seed: int = 123) -> dict[str, Any]: + torch.manual_seed(seed) + controls, expert, v0 = _pair() + terms = v1_handcrafted_reward(controls, expert, v0) + base = terms["base"] + expert_minus_jerky = float(base[0] - base[1]) + + # Misleading WM preference: consequence wants the jerky sample. + preferred = torch.ones(2, 4) + predicted = torch.stack([torch.ones(4) * 3.0, torch.ones(4)]) + + faithful = fidelity_aware_reward( + base, + wm_prediction=torch.zeros(2, 8), + wm_target=torch.zeros(2, 8), + predicted_future=predicted, + preferred_future=preferred, + consequence_weight=1.0, + consequence_scale=1.0, + ) + noise_pred = 10.0 * torch.randn(2, 8) + noisy = fidelity_aware_reward( + base, + wm_prediction=noise_pred, + wm_target=torch.zeros(2, 8), + predicted_future=predicted, + preferred_future=preferred, + consequence_weight=1.0, + consequence_scale=1.0, + ) + + def _rank(reward: torch.Tensor) -> str: + return "expert" if float(reward[0] - reward[1]) > 0 else "jerky" + + base_winner = _rank(base) + return { + "version": FIDELITY_AWARE_REWARD_VERSION, + "source": "constructed_expert_vs_jerky", + "base": { + "expert": float(base[0]), + "jerky": float(base[1]), + "expert_minus_jerky": expert_minus_jerky, + "winner": base_winner, + "safety_expert": float(terms["safety"][0]), + "safety_jerky": float(terms["safety"][1]), + "progress_expert": float(terms["progress"][0]), + "progress_jerky": float(terms["progress"][1]), + "comfort_expert": float(terms["comfort"][0]), + "comfort_jerky": float(terms["comfort"][1]), + }, + "faithful_wm": { + "expert": float(faithful.reward[0]), + "jerky": float(faithful.reward[1]), + "winner": _rank(faithful.reward), + "fidelity_mean": faithful.metadata["fidelity_mean"], + }, + "noise_wm": { + "expert": float(noisy.reward[0]), + "jerky": float(noisy.reward[1]), + "winner": _rank(noisy.reward), + "fidelity_mean": noisy.metadata["fidelity_mean"], + "ranking_matches_base": _rank(noisy.reward) == base_winner, + }, + } + + +def main() -> None: + import argparse + import json + from pathlib import Path + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--out", type=Path, default=None) + parser.add_argument("--seed", type=int, default=123) + args = parser.parse_args() + report = run_fidelity_reward_experiment(seed=args.seed) + text = json.dumps(report, 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/evaluation/results/fidelity_reward_experiment.json b/Model/evaluation/results/fidelity_reward_experiment.json new file mode 100644 index 000000000..7d0380505 --- /dev/null +++ b/Model/evaluation/results/fidelity_reward_experiment.json @@ -0,0 +1,29 @@ +{ + "version": "v1_safety_comfort_progress_gated_wm", + "source": "constructed_expert_vs_jerky", + "base": { + "expert": -0.0, + "jerky": -0.3004057705402374, + "expert_minus_jerky": 0.3004057705402374, + "winner": "expert", + "safety_expert": -0.0, + "safety_jerky": -0.0, + "progress_expert": -0.0, + "progress_jerky": -0.29754963517189026, + "comfort_expert": -0.0, + "comfort_jerky": -0.005712286103516817 + }, + "faithful_wm": { + "expert": -4.0, + "jerky": -0.3004057705402374, + "winner": "jerky", + "fidelity_mean": 1.0 + }, + "noise_wm": { + "expert": -2.3237379309648531e-07, + "jerky": -0.3004057705402374, + "winner": "expert", + "fidelity_mean": 2.9046724137060664e-08, + "ranking_matches_base": true + } +} \ No newline at end of file diff --git a/Model/tests/test_fidelity_aware_reward.py b/Model/tests/test_fidelity_aware_reward.py index 491b58252..f87f6d22c 100644 --- a/Model/tests/test_fidelity_aware_reward.py +++ b/Model/tests/test_fidelity_aware_reward.py @@ -14,7 +14,7 @@ def test_version_pin() -> None: - assert FIDELITY_AWARE_REWARD_VERSION == "wm_fidelity_gate_v1" + assert FIDELITY_AWARE_REWARD_VERSION == "v1_safety_comfort_progress_gated_wm" def test_perfect_wm_fidelity_is_one() -> None: @@ -31,9 +31,9 @@ def test_error_lowers_fidelity() -> None: assert torch.all(fid > 0.0) -def test_fidelity_gates_base_reward() -> None: +def test_fidelity_does_not_zero_handcrafted_base() -> None: + """#123: a broken WM must not wipe safety/comfort/progress.""" base = torch.tensor([1.0, 1.0]) - # Sample 0: perfect WM; sample 1: large error → low fidelity. wm_pred = torch.tensor([[0.0, 0.0], [0.0, 0.0]]) wm_tgt = torch.tensor([[0.0, 0.0], [10.0, 10.0]]) out = fidelity_aware_reward( @@ -42,9 +42,9 @@ def test_fidelity_gates_base_reward() -> None: wm_target=wm_tgt, fidelity_temperature=1.0, ) - assert out.reward[0].item() == 1.0 - assert out.reward[1].item() < out.reward[0].item() - assert out.metadata["reward_mean"] < 1.0 + assert torch.allclose(out.reward, base) + assert out.fidelity[0].item() == 1.0 + assert out.fidelity[1].item() < out.fidelity[0].item() def test_consequence_term_requires_external_preference() -> None: @@ -76,3 +76,14 @@ def test_soft_advantage_zero_mean_without_baseline() -> None: reward = torch.tensor([1.0, 3.0]) adv = soft_advantage_from_reward(reward) assert abs(adv.mean().item()) < 1e-6 + + +def test_expert_outranks_jerky_on_v1_base() -> None: + from evaluation.fidelity_reward_experiment import run_fidelity_reward_experiment + + report = run_fidelity_reward_experiment() + assert report["base"]["expert_minus_jerky"] > 0 + assert report["noise_wm"]["ranking_matches_base"] is True + assert report["noise_wm"]["fidelity_mean"] < 0.05 + assert report["faithful_wm"]["fidelity_mean"] > 0.99 + diff --git a/Model/training/losses/__init__.py b/Model/training/losses/__init__.py index 7c74524c8..afd50bc56 100644 --- a/Model/training/losses/__init__.py +++ b/Model/training/losses/__init__.py @@ -17,9 +17,11 @@ from .fidelity_aware_reward import ( FIDELITY_AWARE_REWARD_VERSION, FidelityAwareRewardResult, + V1_WEIGHTS, consequence_alignment_reward, fidelity_aware_reward, soft_advantage_from_reward, + v1_handcrafted_reward, world_model_fidelity, ) @@ -32,10 +34,12 @@ "RouteConsistencyLoss", "RouteConsistencyWeights", "RolloutAlignedLoss", + "V1_WEIGHTS", "consequence_alignment_reward", "ego_points_to_grid", "fidelity_aware_reward", "integrate_controls_torch", "soft_advantage_from_reward", + "v1_handcrafted_reward", "world_model_fidelity", ] diff --git a/Model/training/losses/fidelity_aware_reward.py b/Model/training/losses/fidelity_aware_reward.py index 2ce58dbed..086f51a49 100644 --- a/Model/training/losses/fidelity_aware_reward.py +++ b/Model/training/losses/fidelity_aware_reward.py @@ -1,19 +1,19 @@ -"""Fidelity-aware reward prototype for stage-3 closed-loop RL (issue #123). - -Design rule from the issue: every reward term must add information the -reactive policy cannot already infer from its own inputs. A reward built -only from the reasoning band (a re-encoding of planner inputs) is redundant -for the same DPI reason that drove ``reasoning_coupling.alpha → 0``. - -This module therefore shapes rewards with **world-model fidelity** — how -well the WM predicts consequences of the rolled-out action — and optionally -scores those predicted consequences against an external preference signal. -Both are information the camera-only reactive policy does not observe at -decision time. - -This is deliberately a pure function / small helper, not a training-loop -integration. Wire it behind ``compute_planner_loss`` (#115) when that hook -lands; until then it is safe to unit-test and prototype against. +"""Fidelity-aware reward for stage-3 closed-loop RL (issue #123). + +v1 (the issue's concrete formula, minus an imitation KL term):: + + R = w_safe R_safety + w_prog R_progress + w_comf R_comfort + + g * R_wm + +``R_safety / R_progress / R_comfort`` come from the same tensors +``RolloutAlignedLoss`` already uses (unicycle rollout + comfort excess). +``g = exp(-mse / T)`` is world-model fidelity. It gates **only** the WM +consequence term — a broken WM must not zero collision/comfort. + +This is the tensor #177 (AlpaSim closed-loop) should call later. It does +not import that PR. + +The reasoning-band is intentionally absent: #123's DPI trap. """ from __future__ import annotations @@ -22,8 +22,17 @@ import torch +from training.losses.control_rollout import integrate_controls_torch +from training.losses.rollout_aligned_loss import comfort_excess_per_sample -FIDELITY_AWARE_REWARD_VERSION = "wm_fidelity_gate_v1" + +FIDELITY_AWARE_REWARD_VERSION = "v1_safety_comfort_progress_gated_wm" + +V1_WEIGHTS = { + "safety": 1.0, + "progress": 1.0, + "comfort": 0.5, +} @dataclass(frozen=True) @@ -105,6 +114,43 @@ def consequence_alignment_reward( return -scale * mse +def v1_handcrafted_reward( + controls: torch.Tensor, + expert_controls: torch.Tensor, + initial_speed: torch.Tensor, + *, + dt: float = 0.1, + lane_half_width_m: float = 1.75, +) -> dict[str, torch.Tensor]: + """#123 handcrafted terms from the existing rollout/comfort helpers. + + Returns per-sample ``safety``, ``progress``, ``comfort``, and ``base`` + (weighted sum). Higher is better. ``expert_controls`` is the logged plan + used as the progress target and the comfort reference. + """ + pos, _, speeds = integrate_controls_torch(controls, initial_speed, dt=dt) + exp_pos, _, exp_speeds = integrate_controls_torch( + expert_controls, initial_speed, dt=dt, + ) + r_progress = -(pos - exp_pos).pow(2).sum(dim=-1).sqrt().mean(dim=1) + comfort, _, _ = comfort_excess_per_sample( + controls, expert_controls, speeds, exp_speeds, dt=dt, + ) + r_comfort = -comfort + r_safety = -torch.relu(pos[..., 1].abs() - lane_half_width_m).mean(dim=1) + base = ( + V1_WEIGHTS["safety"] * r_safety + + V1_WEIGHTS["progress"] * r_progress + + V1_WEIGHTS["comfort"] * r_comfort + ) + return { + "safety": r_safety, + "progress": r_progress, + "comfort": r_comfort, + "base": base, + } + + def fidelity_aware_reward( base_reward: torch.Tensor, *, @@ -117,15 +163,7 @@ def fidelity_aware_reward( consequence_weight: float = 1.0, min_fidelity: float = 0.0, ) -> FidelityAwareRewardResult: - """Gate ``base_reward`` (and optional consequence term) by WM fidelity. - - Final reward:: - - r = fidelity * (base_reward + consequence_weight * consequence) - - where ``fidelity`` is clipped below by ``min_fidelity`` so a broken WM - cannot zero the entire objective if a floor is desired. - """ + """#123 v1: ``R = base + g * R_wm``. ``g`` does not multiply safety/comfort.""" if base_reward.ndim > 1: raise ValueError( f"base_reward must be scalar or (B,), got shape {tuple(base_reward.shape)}" @@ -169,9 +207,9 @@ def fidelity_aware_reward( "consequence/base_reward shape mismatch: " f"{tuple(consequence.shape)} vs {tuple(total.shape)}" ) - total = total + float(consequence_weight) * consequence + total = total + float(consequence_weight) * fidelity * consequence - reward = fidelity * total + reward = total metadata = { "fidelity_mean": float(fidelity.detach().mean().cpu()), "base_reward_mean": float(base_reward.detach().float().mean().cpu()), From af0fc401f6f073e28e595d3ccd15749c193a42bd Mon Sep 17 00:00:00 2001 From: ShauryaVM Date: Sat, 22 Aug 2026 17:14:19 -0400 Subject: [PATCH 3/3] Fix path-frame v1 terms and derive the WM gate. Safety is cross-track vs the intended path (not ego-y), progress is along-track, comfort is vs physical jerk limits, and g saturates at 0.9. Co-authored-by: Cursor --- .../evaluation/fidelity_reward_experiment.py | 46 +++- .../results/fidelity_reward_experiment.json | 29 +-- Model/tests/test_fidelity_aware_reward.py | 129 +++++++++++- Model/training/losses/__init__.py | 8 + .../training/losses/fidelity_aware_reward.py | 197 +++++++++++++++--- 5 files changed, 356 insertions(+), 53 deletions(-) diff --git a/Model/evaluation/fidelity_reward_experiment.py b/Model/evaluation/fidelity_reward_experiment.py index 6d290c5c1..c9972e353 100644 --- a/Model/evaluation/fidelity_reward_experiment.py +++ b/Model/evaluation/fidelity_reward_experiment.py @@ -8,23 +8,37 @@ from __future__ import annotations +import math from typing import Any import torch +from training.losses.control_rollout import integrate_controls_torch from training.losses.fidelity_aware_reward import ( + EXPERIMENT_NOISE_WM_SIGMA, + EXPERIMENT_WM_DIM, FIDELITY_AWARE_REWARD_VERSION, + FIDELITY_SATURATION, + FIDELITY_TEMPERATURE, fidelity_aware_reward, v1_handcrafted_reward, ) -def _pair(timesteps: int = 32) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: +def _pair( + timesteps: int = 32, + dt: float = 0.1, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: t = torch.arange(timesteps, dtype=torch.float32) + time_s = t * dt expert = torch.zeros(1, timesteps, 2) + # Original ``8.0 * sin(2π t / 2)`` with t = 0,1,2,... is sin(π t) and + # every integer sample is a zero crossing (max |a| ~ 4e-5). Use time + # in seconds (2 s period, as written) plus π/4 so sample instants are + # off the zeros. Curvature already samples off-zero at period 3. jerky = torch.stack( ( - 8.0 * torch.sin(2.0 * torch.pi * t / 2.0), + 8.0 * torch.sin(2.0 * torch.pi * time_s / 2.0 + 0.25 * torch.pi), 0.25 * torch.sin(2.0 * torch.pi * t / 3.0), ), dim=-1, @@ -42,24 +56,37 @@ def run_fidelity_reward_experiment(*, seed: int = 123) -> dict[str, Any]: base = terms["base"] expert_minus_jerky = float(base[0] - base[1]) + pos, _, _ = integrate_controls_torch(controls, v0) + max_abs_accel = float(controls[1, :, 0].abs().max()) + max_abs_ego_y = float(pos[1, :, 1].abs().max()) + max_abs_cross_track = float(terms["cross_track_m"][1]) + # Misleading WM preference: consequence wants the jerky sample. + # Magnitude is 2× the handcrafted gap / g_sat so a saturated faithful + # WM still flips, while g ≈ 0 cannot. Reconstruction fidelity is not + # preference fidelity — that is the point of the gate. preferred = torch.ones(2, 4) - predicted = torch.stack([torch.ones(4) * 3.0, torch.ones(4)]) + gap = max(expert_minus_jerky, 0.05) + mse_needed = 2.0 * gap / FIDELITY_SATURATION + delta = math.sqrt(mse_needed) + predicted = torch.stack( + [torch.full((4,), 1.0 + delta), torch.ones(4)], + ) faithful = fidelity_aware_reward( base, - wm_prediction=torch.zeros(2, 8), - wm_target=torch.zeros(2, 8), + wm_prediction=torch.zeros(2, EXPERIMENT_WM_DIM), + wm_target=torch.zeros(2, EXPERIMENT_WM_DIM), predicted_future=predicted, preferred_future=preferred, consequence_weight=1.0, consequence_scale=1.0, ) - noise_pred = 10.0 * torch.randn(2, 8) + noise_pred = EXPERIMENT_NOISE_WM_SIGMA * torch.randn(2, EXPERIMENT_WM_DIM) noisy = fidelity_aware_reward( base, wm_prediction=noise_pred, - wm_target=torch.zeros(2, 8), + wm_target=torch.zeros(2, EXPERIMENT_WM_DIM), predicted_future=predicted, preferred_future=preferred, consequence_weight=1.0, @@ -73,6 +100,11 @@ def _rank(reward: torch.Tensor) -> str: return { "version": FIDELITY_AWARE_REWARD_VERSION, "source": "constructed_expert_vs_jerky", + "fidelity_temperature": FIDELITY_TEMPERATURE, + "fidelity_saturation": FIDELITY_SATURATION, + "jerky_max_abs_accel": max_abs_accel, + "jerky_max_abs_ego_y": max_abs_ego_y, + "jerky_max_abs_cross_track": max_abs_cross_track, "base": { "expert": float(base[0]), "jerky": float(base[1]), diff --git a/Model/evaluation/results/fidelity_reward_experiment.json b/Model/evaluation/results/fidelity_reward_experiment.json index 7d0380505..27abd8e4c 100644 --- a/Model/evaluation/results/fidelity_reward_experiment.json +++ b/Model/evaluation/results/fidelity_reward_experiment.json @@ -1,29 +1,34 @@ { "version": "v1_safety_comfort_progress_gated_wm", "source": "constructed_expert_vs_jerky", + "fidelity_temperature": 1.0, + "fidelity_saturation": 0.9, + "jerky_max_abs_accel": 7.901507377624512, + "jerky_max_abs_ego_y": 0.9882224798202515, + "jerky_max_abs_cross_track": 0.9882224798202515, "base": { - "expert": -0.0, - "jerky": -0.3004057705402374, - "expert_minus_jerky": 0.3004057705402374, + "expert": 1.0, + "jerky": -2.8394618034362793, + "expert_minus_jerky": 3.8394618034362793, "winner": "expert", "safety_expert": -0.0, "safety_jerky": -0.0, - "progress_expert": -0.0, - "progress_jerky": -0.29754963517189026, + "progress_expert": 1.0, + "progress_jerky": 1.4938430786132812, "comfort_expert": -0.0, - "comfort_jerky": -0.005712286103516817 + "comfort_jerky": -8.666609764099121 }, "faithful_wm": { - "expert": -4.0, - "jerky": -0.3004057705402374, + "expert": -6.678924083709717, + "jerky": -2.8394618034362793, "winner": "jerky", - "fidelity_mean": 1.0 + "fidelity_mean": 0.8999999761581421 }, "noise_wm": { - "expert": -2.3237379309648531e-07, - "jerky": -0.3004057705402374, + "expert": 0.9999995827674866, + "jerky": -2.8394618034362793, "winner": "expert", - "fidelity_mean": 2.9046724137060664e-08, + "fidelity_mean": 2.6142050302269126e-08, "ranking_matches_base": true } } \ No newline at end of file diff --git a/Model/tests/test_fidelity_aware_reward.py b/Model/tests/test_fidelity_aware_reward.py index f87f6d22c..4de698c3f 100644 --- a/Model/tests/test_fidelity_aware_reward.py +++ b/Model/tests/test_fidelity_aware_reward.py @@ -2,13 +2,22 @@ from __future__ import annotations +import math + import torch +from training.losses.control_rollout import integrate_controls_torch from training.losses.fidelity_aware_reward import ( + EXPERIMENT_NOISE_WM_SIGMA, + EXPERIMENT_WM_DIM, + FAITHFUL_NOISE_RATIO, FIDELITY_AWARE_REWARD_VERSION, + FIDELITY_SATURATION, + FIDELITY_TEMPERATURE, consequence_alignment_reward, fidelity_aware_reward, soft_advantage_from_reward, + v1_handcrafted_reward, world_model_fidelity, ) @@ -17,20 +26,45 @@ def test_version_pin() -> None: assert FIDELITY_AWARE_REWARD_VERSION == "v1_safety_comfort_progress_gated_wm" -def test_perfect_wm_fidelity_is_one() -> None: +def test_fidelity_temperature_is_derived_from_experiment_mse() -> None: + assert FAITHFUL_NOISE_RATIO == 10.0 + assert EXPERIMENT_NOISE_WM_SIGMA == 10.0 + assert FIDELITY_TEMPERATURE == ( + EXPERIMENT_NOISE_WM_SIGMA / FAITHFUL_NOISE_RATIO + ) ** 2 + assert FIDELITY_SATURATION == ( + EXPERIMENT_WM_DIM + 1 + ) / (EXPERIMENT_WM_DIM + 2) + assert FIDELITY_SATURATION < 1.0 + + +def test_perfect_wm_fidelity_saturates_below_one() -> None: x = torch.randn(4, 8, 8) fid = world_model_fidelity(x, x.clone()) - assert torch.allclose(fid, torch.ones_like(fid)) + assert torch.allclose(fid, torch.full_like(fid, FIDELITY_SATURATION)) + assert torch.all(fid < 1.0) def test_error_lowers_fidelity() -> None: pred = torch.zeros(2, 4) target = torch.ones(2, 4) - fid = world_model_fidelity(pred, target, temperature=1.0) - assert torch.all(fid < 1.0) + fid = world_model_fidelity(pred, target, temperature=FIDELITY_TEMPERATURE) + assert torch.all(fid < FIDELITY_SATURATION) assert torch.all(fid > 0.0) +def test_good_enough_wm_does_not_hit_saturation_ceiling() -> None: + """mse = T (barely-faithful) must sit at g_sat/e, not g_sat.""" + pred = torch.zeros(2, EXPERIMENT_WM_DIM) + target = torch.ones(2, EXPERIMENT_WM_DIM) + fid = world_model_fidelity(pred, target) + expected = FIDELITY_SATURATION * math.exp(-1.0 / FIDELITY_TEMPERATURE) + assert torch.allclose(fid, torch.full_like(fid, expected), atol=1e-5) + perfect = world_model_fidelity(pred, pred) + assert float(perfect[0]) > float(fid[0]) + assert float(perfect[0]) < 1.0 + + def test_fidelity_does_not_zero_handcrafted_base() -> None: """#123: a broken WM must not wipe safety/comfort/progress.""" base = torch.tensor([1.0, 1.0]) @@ -40,10 +74,9 @@ def test_fidelity_does_not_zero_handcrafted_base() -> None: base, wm_prediction=wm_pred, wm_target=wm_tgt, - fidelity_temperature=1.0, ) assert torch.allclose(out.reward, base) - assert out.fidelity[0].item() == 1.0 + assert abs(out.fidelity[0].item() - FIDELITY_SATURATION) < 1e-6 assert out.fidelity[1].item() < out.fidelity[0].item() @@ -61,8 +94,11 @@ def test_consequence_term_requires_external_preference() -> None: consequence_weight=1.0, consequence_scale=1.0, ) - # Perfect WM fidelity=1, consequence = -mse = -1 → reward -1. - assert torch.allclose(out.reward, torch.tensor([-1.0, -1.0])) + # Saturated WM, consequence = -mse = -1 → reward -g_sat. + assert torch.allclose( + out.reward, + torch.full((2,), -FIDELITY_SATURATION), + ) assert out.consequence_reward is not None @@ -78,12 +114,85 @@ def test_soft_advantage_zero_mean_without_baseline() -> None: assert abs(adv.mean().item()) < 1e-6 +def test_safety_ignores_in_lane_turn() -> None: + """Curvature 0.05 at 32 steps / 5 m/s has |y| > 1.75 and used to cost ~1.07. + + In the intended-path frame the same curve is the lane centerline, so + r_safety must be ~0. This fails on the old ego-y gate. + """ + timesteps = 32 + dt = 0.1 + curve = torch.zeros(1, timesteps, 2) + curve[:, :, 1] = 0.05 + v0 = torch.tensor([5.0]) + terms = v1_handcrafted_reward(curve, curve, v0, dt=dt) + pos, _, _ = integrate_controls_torch(curve, v0, dt=dt) + ego_y_penalty = float(torch.relu(pos[..., 1].abs() - 1.75).mean()) + assert pos[0, :, 1].abs().max().item() > 1.75 + assert ego_y_penalty > 1.0 + assert abs(terms["safety"].item()) < 1e-4 + + +def test_safety_penalises_leaving_lane() -> None: + timesteps = 32 + dt = 0.1 + straight = torch.zeros(1, timesteps, 2) + leave = torch.zeros(1, timesteps, 2) + leave[:, :, 1] = 0.05 + v0 = torch.tensor([5.0]) + terms = v1_handcrafted_reward(leave, straight, v0, dt=dt) + assert terms["safety"].item() < -1.0 + assert terms["cross_track_m"].item() > 1.75 + + +def test_progress_rewards_forward_motion_not_imitation() -> None: + """Stopped-vs-stopped is the old imitation optimum; v1 must prefer coasting.""" + timesteps = 32 + brake = torch.zeros(2, timesteps, 2) + brake[:, :, 0] = -100.0 + coast = torch.zeros(2, timesteps, 2) + controls = torch.stack([brake[0], coast[0]]) + v0 = torch.tensor([5.0, 5.0]) + terms = v1_handcrafted_reward(controls, brake, v0) + assert terms["progress"][1].item() > terms["progress"][0].item() + + +def test_comfort_uses_physical_threshold_not_expert_log() -> None: + """Matching a jerky log used to yield comfort 0; v1 must still penalise.""" + timesteps = 16 + t = torch.arange(timesteps, dtype=torch.float32) + jerky = torch.stack( + ( + 8.0 * torch.sin(math.pi * t + 0.5), + torch.zeros(timesteps), + ), + dim=-1, + ).unsqueeze(0) + v0 = torch.tensor([5.0]) + terms = v1_handcrafted_reward(jerky, jerky, v0) + assert terms["comfort"].item() < 0.0 + + +def test_jerky_experiment_signal_has_nonzero_accel_at_samples() -> None: + from evaluation.fidelity_reward_experiment import _pair + + controls, _, _ = _pair() + jerky_accel = controls[1, :, 0] + assert jerky_accel.abs().max().item() > 1.0 + + def test_expert_outranks_jerky_on_v1_base() -> None: from evaluation.fidelity_reward_experiment import run_fidelity_reward_experiment report = run_fidelity_reward_experiment() assert report["base"]["expert_minus_jerky"] > 0 + assert abs(report["base"]["progress_expert"] - 1.0) < 0.02 + assert report["base"]["comfort_jerky"] < report["base"]["comfort_expert"] + assert report["jerky_max_abs_accel"] > 1.0 assert report["noise_wm"]["ranking_matches_base"] is True assert report["noise_wm"]["fidelity_mean"] < 0.05 - assert report["faithful_wm"]["fidelity_mean"] > 0.99 - + assert abs(report["faithful_wm"]["fidelity_mean"] - FIDELITY_SATURATION) < 1e-6 + assert report["faithful_wm"]["winner"] == "jerky" + assert report["fidelity_temperature"] == FIDELITY_TEMPERATURE + assert report["fidelity_saturation"] == FIDELITY_SATURATION + assert report["faithful_wm"]["fidelity_mean"] < 1.0 diff --git a/Model/training/losses/__init__.py b/Model/training/losses/__init__.py index afd50bc56..9cdb9b120 100644 --- a/Model/training/losses/__init__.py +++ b/Model/training/losses/__init__.py @@ -15,18 +15,25 @@ ego_points_to_grid, ) from .fidelity_aware_reward import ( + EXPERIMENT_NOISE_WM_SIGMA, FIDELITY_AWARE_REWARD_VERSION, + FIDELITY_SATURATION, + FIDELITY_TEMPERATURE, FidelityAwareRewardResult, V1_WEIGHTS, consequence_alignment_reward, fidelity_aware_reward, + path_frame_along_cross, soft_advantage_from_reward, v1_handcrafted_reward, world_model_fidelity, ) __all__ = [ + "EXPERIMENT_NOISE_WM_SIGMA", "FIDELITY_AWARE_REWARD_VERSION", + "FIDELITY_SATURATION", + "FIDELITY_TEMPERATURE", "FidelityAwareRewardResult", "HorizonReasoningLoss", "ROLLOUT_ALIGNED_LOSS_VERSION", @@ -39,6 +46,7 @@ "ego_points_to_grid", "fidelity_aware_reward", "integrate_controls_torch", + "path_frame_along_cross", "soft_advantage_from_reward", "v1_handcrafted_reward", "world_model_fidelity", diff --git a/Model/training/losses/fidelity_aware_reward.py b/Model/training/losses/fidelity_aware_reward.py index 086f51a49..449c0b3b3 100644 --- a/Model/training/losses/fidelity_aware_reward.py +++ b/Model/training/losses/fidelity_aware_reward.py @@ -5,10 +5,22 @@ R = w_safe R_safety + w_prog R_progress + w_comf R_comfort + g * R_wm -``R_safety / R_progress / R_comfort`` come from the same tensors -``RolloutAlignedLoss`` already uses (unicycle rollout + comfort excess). -``g = exp(-mse / T)`` is world-model fidelity. It gates **only** the WM -consequence term — a broken WM must not zero collision/comfort. +Handcrafted terms are **not** imitation: + +* ``R_safety`` is mean lane-departure in the intended-path Frenet frame + (cross-track vs the path heading along the trajectory), not ``|y|`` in + the t=0 ego frame. Turning in a ~3.5 m lane is free; leaving it is not. +* ``R_progress`` is along-track displacement of the predicted trajectory + along that same path, normalized by the coast horizon. Matching the + expert's positions is not the objective; going forward along the path is. +* ``R_comfort`` is jerk / lateral-accel excess vs the physical thresholds + already used by ``RolloutAlignedLoss``, not vs the expert log. + +``g = g_sat * exp(-mse / T)`` gates **only** the WM consequence term — a +broken WM must not zero collision/comfort. ``T`` and ``g_sat`` are derived +below from the offline experiment's WM residual scale; ``g_sat < 1`` so a +"good enough" reconstruction does not fully trust a (possibly misleading) +consequence. This is the tensor #177 (AlpaSim closed-loop) should call later. It does not import that PR. @@ -23,7 +35,6 @@ import torch from training.losses.control_rollout import integrate_controls_torch -from training.losses.rollout_aligned_loss import comfort_excess_per_sample FIDELITY_AWARE_REWARD_VERSION = "v1_safety_comfort_progress_gated_wm" @@ -34,6 +45,38 @@ "comfort": 0.5, } +# Physical comfort limits — same numbers as RolloutAlignedLoss. +JERK_THRESHOLD_MPS3 = 4.13 +LATERAL_ACCEL_THRESHOLD_MPS2 = 4.89 + +# --- fidelity temperature / saturation -------------------------------------- +# Offline experiment WM residual is an (B, 8) tensor. +# faithful arm: mse = 0 +# noise arm: pred ~ N(0, σ_noise²) with σ_noise = 10 → E[mse] = 100 +# A barely-faithful WM is defined as 10× smaller residual than that noise +# (σ_good = 1). T is its expected per-element MSE, so: +# g(barely-faithful) = g_sat / e +# g(noise) = g_sat * exp(-100) ≈ 0 +# and the exponential still moves in the operating range instead of sitting +# at the ceiling for every "pretty good" WM. +EXPERIMENT_WM_DIM = 8 +EXPERIMENT_NOISE_WM_SIGMA = 10.0 +FAITHFUL_NOISE_RATIO = 10.0 +FIDELITY_TEMPERATURE = ( + EXPERIMENT_NOISE_WM_SIGMA / FAITHFUL_NOISE_RATIO +) ** 2 # 1.0 + +# Even mse = 0 does not fully trust R_wm: reconstruction fidelity is not +# consequence-target fidelity (the experiment's misleading preference). +# Rule of succession on the experiment's 8-d residual: after D exact +# matches, P(trust) = (D+1)/(D+2) = 0.9, not 1. +FIDELITY_SATURATION = (EXPERIMENT_WM_DIM + 1) / (EXPERIMENT_WM_DIM + 2) + +# Prefix/suffix used to treat the intended path as an infinite lane so +# along-track is not capped at the expert's last point (that cap would +# re-introduce imitation). +_LANE_EXTENSION_M = 1.0e3 + @dataclass(frozen=True) class FidelityAwareRewardResult: @@ -50,14 +93,17 @@ def world_model_fidelity( prediction: torch.Tensor, target: torch.Tensor, *, - temperature: float = 1.0, + temperature: float = FIDELITY_TEMPERATURE, + saturation: float = FIDELITY_SATURATION, reduce_dims: tuple[int, ...] | None = None, ) -> torch.Tensor: - """Map WM prediction error to a ``(0, 1]`` fidelity weight. + """Map WM prediction error to a ``(0, g_sat]`` fidelity weight. - ``fidelity = exp(-mse / temperature)``. High fidelity means the world - model is a trustworthy source of consequence information for this sample; - low fidelity down-weights any reward that depends on those predictions. + ``fidelity = saturation * exp(-mse / temperature)``. High fidelity + means the world model is a trustworthy source of consequence + information for this sample; low fidelity down-weights any reward + that depends on those predictions. ``saturation < 1`` keeps a + "good enough" WM from fully deferring to a misspecified consequence. """ if prediction.shape != target.shape: raise ValueError( @@ -66,6 +112,10 @@ def world_model_fidelity( ) if temperature <= 0: raise ValueError(f"temperature must be > 0, got {temperature}") + if saturation <= 0 or saturation > 1: + raise ValueError( + f"saturation must be in (0, 1], got {saturation}" + ) err = (prediction.float() - target.float()).pow(2) if reduce_dims is None: @@ -77,7 +127,7 @@ def world_model_fidelity( else: mse = err.mean(dim=reduce_dims) - return torch.exp(-mse / temperature) + return float(saturation) * torch.exp(-mse / temperature) def consequence_alignment_reward( @@ -114,30 +164,124 @@ def consequence_alignment_reward( return -scale * mse +def _polyline_along_cross( + pos: torch.Tensor, + ref: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Nearest-point Frenet coords of ``pos`` vs polyline ``ref``. + + Both ``(B, T, 2)``. Along-track is arc length from ``ref[:, 0]``; + cross-track is signed left-positive distance to the closest segment. + """ + seg = ref[:, 1:, :] - ref[:, :-1, :] + seg_len = torch.linalg.norm(seg, dim=-1) + seg_len_sq = seg_len.square().clamp(min=1e-12) + zeros = torch.zeros( + ref.shape[0], 1, dtype=ref.dtype, device=ref.device, + ) + arc_start = torch.cat( + [zeros, torch.cumsum(seg_len, dim=-1)[:, :-1]], dim=-1, + ) + + start = ref[:, None, :-1, :] + rel = pos[:, :, None, :] - start + t_proj = (rel * seg[:, None, :, :]).sum(-1) / seg_len_sq[:, None, :] + t_clamped = t_proj.clamp(0.0, 1.0) + closest = start + t_clamped.unsqueeze(-1) * seg[:, None, :, :] + offset = pos[:, :, None, :] - closest + dist2 = (offset * offset).sum(-1) + idx = dist2.argmin(dim=-1) + + batch = torch.arange(pos.shape[0], device=pos.device)[:, None] + time = torch.arange(pos.shape[1], device=pos.device)[None, :] + t_star = t_clamped[batch, time, idx] + along = arc_start[batch, idx] + t_star * seg_len[batch, idx] + seg_n = seg[batch, idx] + off_n = offset[batch, time, idx] + denom = seg_len[batch, idx].clamp(min=1e-6) + cross = ( + seg_n[..., 0] * off_n[..., 1] - seg_n[..., 1] * off_n[..., 0] + ) / denom + return along, cross + + +def path_frame_along_cross( + pos: torch.Tensor, + ref_pos: torch.Tensor, + ref_heading: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Along-track / cross-track of ``pos`` in the intended-path frame. + + The path is the polyline ``ref_pos``, extended along ``ref_heading`` + at both ends so progress is not capped at the last logged point (that + cap would reward matching the expert's horizon, i.e. imitation). + + This is **not** ``|y|`` in the t=0 ego frame: a curved in-lane path + has large ego-y and near-zero cross-track. + """ + u0 = torch.stack( + [ref_heading[:, 0].cos(), ref_heading[:, 0].sin()], dim=-1, + ) + u1 = torch.stack( + [ref_heading[:, -1].cos(), ref_heading[:, -1].sin()], dim=-1, + ) + origin = torch.zeros( + ref_pos.shape[0], 1, 2, dtype=ref_pos.dtype, device=ref_pos.device, + ) + extend = pos.new_tensor(_LANE_EXTENSION_M) + prefix = origin - extend * u0[:, None, :] + suffix = ref_pos[:, -1:, :] + extend * u1[:, None, :] + ref_ext = torch.cat([prefix, origin, ref_pos, suffix], dim=1) + along_ext, cross = _polyline_along_cross(pos, ref_ext) + return along_ext - _LANE_EXTENSION_M, cross + + +def _absolute_comfort_penalty( + controls: torch.Tensor, + speeds: torch.Tensor, + *, + dt: float, +) -> torch.Tensor: + """Mean jerk / lateral-accel excess vs physical thresholds, not the log. + + ``comfort_excess_per_sample`` charges ``relu(peak_pred - peak_target)``, + which is zero when the expert is equally jerky. v1 must not do that. + """ + accel = controls[:, :, 0] + curvature = controls[:, :, 1] + jerk = (accel[:, 1:] - accel[:, :-1]) / dt + lateral = speeds.square() * curvature + jerk_excess = torch.relu(jerk.abs() - JERK_THRESHOLD_MPS3).mean(dim=1) + lateral_excess = torch.relu( + lateral.abs() - LATERAL_ACCEL_THRESHOLD_MPS2, + ).mean(dim=1) + return 0.5 * (jerk_excess + lateral_excess) + + def v1_handcrafted_reward( controls: torch.Tensor, - expert_controls: torch.Tensor, + intended_controls: torch.Tensor, initial_speed: torch.Tensor, *, dt: float = 0.1, lane_half_width_m: float = 1.75, ) -> dict[str, torch.Tensor]: - """#123 handcrafted terms from the existing rollout/comfort helpers. + """#123 handcrafted terms in the intended-path frame. + ``intended_controls`` is the logged plan used as **path geometry** + (centerline heading along the trajectory), not an imitation target. Returns per-sample ``safety``, ``progress``, ``comfort``, and ``base`` - (weighted sum). Higher is better. ``expert_controls`` is the logged plan - used as the progress target and the comfort reference. + (weighted sum). Higher is better. """ pos, _, speeds = integrate_controls_torch(controls, initial_speed, dt=dt) - exp_pos, _, exp_speeds = integrate_controls_torch( - expert_controls, initial_speed, dt=dt, - ) - r_progress = -(pos - exp_pos).pow(2).sum(dim=-1).sqrt().mean(dim=1) - comfort, _, _ = comfort_excess_per_sample( - controls, expert_controls, speeds, exp_speeds, dt=dt, + ref_pos, ref_heading, _ = integrate_controls_torch( + intended_controls, initial_speed, dt=dt, ) - r_comfort = -comfort - r_safety = -torch.relu(pos[..., 1].abs() - lane_half_width_m).mean(dim=1) + along, cross = path_frame_along_cross(pos, ref_pos, ref_heading) + horizon_m = initial_speed.clamp(min=1e-3) * dt * pos.shape[1] + r_progress = along[:, -1] / horizon_m + r_safety = -torch.relu(cross.abs() - lane_half_width_m).mean(dim=1) + r_comfort = -_absolute_comfort_penalty(controls, speeds, dt=dt) base = ( V1_WEIGHTS["safety"] * r_safety + V1_WEIGHTS["progress"] * r_progress @@ -148,6 +292,8 @@ def v1_handcrafted_reward( "progress": r_progress, "comfort": r_comfort, "base": base, + "along_track_m": along[:, -1], + "cross_track_m": cross.abs().amax(dim=1), } @@ -158,7 +304,8 @@ def fidelity_aware_reward( wm_target: torch.Tensor, preferred_future: torch.Tensor | None = None, predicted_future: torch.Tensor | None = None, - fidelity_temperature: float = 1.0, + fidelity_temperature: float = FIDELITY_TEMPERATURE, + fidelity_saturation: float = FIDELITY_SATURATION, consequence_scale: float = 1.0, consequence_weight: float = 1.0, min_fidelity: float = 0.0, @@ -173,6 +320,7 @@ def fidelity_aware_reward( wm_prediction, wm_target, temperature=fidelity_temperature, + saturation=fidelity_saturation, ) if fidelity.shape != base_reward.shape and not ( fidelity.ndim == 1 and base_reward.ndim == 0 @@ -215,6 +363,7 @@ def fidelity_aware_reward( "base_reward_mean": float(base_reward.detach().float().mean().cpu()), "reward_mean": float(reward.detach().mean().cpu()), "fidelity_temperature": float(fidelity_temperature), + "fidelity_saturation": float(fidelity_saturation), "consequence_weight": float(consequence_weight), "min_fidelity": float(min_fidelity), }