Fix/ part of #115 compute planner loss

Model/model_components/auto_e2e.py

@@ -99,16 +99,33 @@ def forward(self, camera_tiles, map_context, visual_history,
                 egomotion_history, route_mask=None, map_valid=None,
                 route_valid=None,
                 projection=None, geometry_type=None, image_transform=None,
-                mode="train", trajectory_target=None,
+                mode="train", trajectory_target=None, training_policy=None,
+                return_planner_loss=False,
                 history_frames=None, future_frames=None, **kwargs):
         """
         Run the full autonomous-driving pipeline.
 
         Return contract:
-            * Inference (``mode != "train"``), or train with both branches off →
+            * Inference (``mode != "train"``), or train with both branches off,
+              and ``return_planner_loss`` not set →
               a single trajectory tensor ``[B, num_timesteps * num_signals]``.
+            * ``return_planner_loss=True`` AND ``trajectory_target`` given AND
+              ``mode == "train"`` → the trajectory planner's
+              ``compute_planner_loss()`` result (a ``dict[str, Tensor]`` with
+              at least ``"loss"``, see ``BasePlanner.compute_planner_loss``)
+              in place of the trajectory tensor above — the #115 fix:
+              previously ``trajectory_target`` reached this point but was
+              silently absorbed as an inert kwarg all the way down (it's
+              already passed at ~20 existing call sites across the test
+              suite that expect a trajectory tensor back), so a planner's
+              actual training objective never ran; ``train_il`` regressed
+              whatever forward() happened to return instead.
+              ``return_planner_loss`` defaults to False specifically so
+              none of those existing callers change behavior — only a
+              caller that explicitly opts in gets the new path.
             * Train mode with the World Model and/or the reasoning branch on →
-              ``(trajectory, aux_outputs)`` where ``aux_outputs`` is a dict with
+              ``(result, aux_outputs)`` where ``result`` is whichever of the
+              two things above applies, and ``aux_outputs`` is a dict with
               ``"future_state_pred"`` (World Model) and/or ``"reasoning_pred"``
               (HorizonReasoningPrediction). A dict avoids a positional-tuple
               that grows with every optional branch (#98 Task 4.2).
@@ -131,9 +148,21 @@ def forward(self, camera_tiles, map_context, visual_history,
                 "rectified_pinhole", "ftheta", "pseudo") passed to BEV fusion.
             image_transform: Optional ImageTransform for the model-input frame.
             mode: "train" also returns aux branch outputs for their losses.
+            trajectory_target: optional (B, num_timesteps * num_signals)
+                ground-truth trajectory. By itself, changes nothing — see
+                return_planner_loss.
+            training_policy: optional DatasetTrainingPolicy
+                (Model/training/dataset_policy.py), used only when
+                return_planner_loss=True, forwarded into
+                compute_planner_loss. Pass the object itself, not
+                pre-extracted scalars (#124 review).
+            return_planner_loss: opt-in flag, see "Return contract" above.
+                Default False — train_il is the one caller that should set
+                this True; every other existing caller is unaffected.
 
         Returns:
-            trajectory, or (trajectory, aux_outputs) in train mode with a branch on.
+            result, or (result, aux_outputs) in train mode with a branch on —
+            see "Return contract" above for what result is.
         """
 
         # World Action Model (1 Hz): produce the Encoded Visual History fed to the
@@ -194,21 +223,26 @@ def forward(self, camera_tiles, map_context, visual_history,
 
         # The reasoning branch runs INSIDE ReactiveE2E (after TemporalMemory).
         # In train mode with reasoning on, ReactiveE2E returns
-        # (trajectory, reasoning_pred); otherwise just the trajectory.
+        # (result, reasoning_pred); otherwise just the result — where
+        # `result` is a trajectory tensor if trajectory_target is None, or
+        # a compute_planner_loss dict if it was given (see ReactiveE2E's
+        # own docstring for the full contract).
         reactive_out = self.Reactive_E2E(
             camera_tiles, map_context, visual_history, egomotion_history,
             route_mask=route_mask,
             map_valid=map_valid,
             route_valid=route_valid,
             projection=projection, geometry_type=geometry_type,
             image_transform=image_transform,
-            mode=mode, trajectory_target=trajectory_target, **kwargs,
+            mode=mode, trajectory_target=trajectory_target,
+            training_policy=training_policy,
+            return_planner_loss=return_planner_loss, **kwargs,
         )
         reasoning_pred = None
         if self.enable_reasoning and mode == "train":
-            trajectory, reasoning_pred = reactive_out
+            result, reasoning_pred = reactive_out
         else:
-            trajectory = reactive_out
+            result = reactive_out
 
         # Assemble aux outputs (dict, not a growing positional tuple). The WM
         # keeps its future_frames alongside the prediction so the training loop
@@ -222,7 +256,7 @@ def forward(self, camera_tiles, map_context, visual_history,
                 "future_frames": future_frames,
                 "reasoning_pred": reasoning_pred,
             }
-            return trajectory, aux_outputs
-        return trajectory
+            return result, aux_outputs
+        return result
 
 

Model/model_components/reactive_e2e.py

@@ -118,7 +118,8 @@ def forward(self, camera_tiles, map_context, visual_history,
                 egomotion_history, route_mask=None, map_valid=None,
                 route_valid=None,
                 projection=None, geometry_type=None, image_transform=None,
-                mode="train", **kwargs):
+                mode="train", trajectory_target=None, training_policy=None,
+                return_planner_loss=False, **kwargs):
         """
         Run the reactive end-to-end autonomous-driving pipeline.
 
@@ -136,12 +137,41 @@ def forward(self, camera_tiles, map_context, visual_history,
             geometry_type: Optional explicit geometry label passed to BEV fusion.
             image_transform: Optional ImageTransform for the model-input frame.
             mode: "train" also returns the reasoning prediction (for its loss).
+            trajectory_target: optional (B, num_timesteps * num_signals)
+                ground-truth trajectory. NOTE: passing this alone does NOT
+                change what's returned — trajectory_target is already
+                passed at ~20 existing call sites across the test suite
+                that expect a trajectory tensor back (an external loss_fn
+                consumes it separately, e.g. train_il's own
+                TrajectoryImitationLoss today). Only return_planner_loss
+                (below) opts into the new behavior.
+            training_policy: optional DatasetTrainingPolicy
+                (Model/training/dataset_policy.py), used only when
+                return_planner_loss=True — forwarded into
+                compute_planner_loss. Pass the object itself, not
+                pre-extracted scalars (#124 review: signal_scales/
+                temporal_decay drifting from the policy silently is a
+                real, measured 71% loss difference, not theoretical).
+            return_planner_loss: when True (and mode == "train" and
+                trajectory_target is given), calls the trajectory
+                planner's compute_planner_loss() instead of forward() and
+                returns that dict instead of a trajectory tensor — the
+                #115 fix: FlowMatchingPlanner's Euler rollout is no longer
+                silently regressed against the target externally. Default
+                False so every existing caller's behavior is byte-identical
+                to before this parameter existed; train_il is the one
+                caller that should pass True.
 
         Returns:
             trajectory (B, num_timesteps * num_signals), OR — when the reasoning
             branch is enabled and ``mode == "train"`` — a tuple
             ``(trajectory, reasoning_pred)`` so the training loop can compute the
             reasoning loss. ``reasoning_pred`` is a HorizonReasoningPrediction.
+
+            If return_planner_loss=True (and trajectory_target given, mode
+            == "train"): a loss dict (see BasePlanner.compute_planner_loss)
+            in place of trajectory above, or (loss_dict, reasoning_pred) if
+            reasoning is enabled — same tuple shape as the trajectory case.
         """
         B, V, C, H, W = camera_tiles.shape
 
@@ -236,7 +266,20 @@ def validity_gate(value, valid, *, default, name):
             reasoning_latent = reasoning_pred.reasoning_latent
             reasoning_horizon_tokens = reasoning_pred.horizon_tokens
 
-        # --- Trajectory Prediction ---
+        # --- Training objective path (#115) ---
+        # Explicit opt-in via return_planner_loss — NOT trajectory_target's
+        # mere presence, since trajectory_target is already passed at ~20
+        # existing call sites that expect a trajectory tensor back.
+        if mode == "train" and return_planner_loss and trajectory_target is not None:
+            loss_dict = self.TrajectoryPlanner.compute_planner_loss(
+                fused_features, visual_ctx, ego_ctx, trajectory_target,
+                training_policy=training_policy, **kwargs,
+            )
+            if self.ReasoningHead is not None:
+                return loss_dict, reasoning_pred
+            return loss_dict
+
+        # --- Trajectory Prediction (inference path — unchanged) ---
         trajectory = self.TrajectoryPlanner(
             fused_features, visual_ctx, ego_ctx,
             reasoning_latent=reasoning_latent,
@@ -246,4 +289,4 @@ def validity_gate(value, valid, *, default, name):
 
         if self.ReasoningHead is not None and mode == "train":
             return trajectory, reasoning_pred
-        return trajectory
+        return trajectory

Model/model_components/trajectory_planning/__init__.py

@@ -20,21 +20,6 @@ def build_planner(planner_mode, **kwargs):
             f"Unknown planner_mode {planner_mode!r}. "
             f"Available: {sorted(PLANNER_REGISTRY)}."
         )
-    if planner_mode == "flow_matching":
-        # The training loop (train_il) regresses forward()'s output with SmoothL1,
-        # but forward() Euler-integrates from a fresh noise sample every step — so
-        # that is NOT the flow-matching objective (velocity MSE against x1-x0) and
-        # drives the model to the conditional mean. The proper objective lives in
-        # compute_planner_loss, which is not wired into the train loop yet. Warn
-        # loudly so nobody trains this expecting real flow matching.
-        import warnings
-        warnings.warn(
-            "planner_mode='flow_matching' is NOT correctly trainable via the "
-            "current train_il loop (it L1-regresses an Euler-from-noise rollout, "
-            "not the velocity-MSE flow objective). Use 'bezier' unless/until "
-            "compute_planner_loss is wired in.",
-            RuntimeWarning, stacklevel=2,
-        )
     return PLANNER_REGISTRY[planner_mode](**kwargs)
 
 __all__ = [
@@ -43,4 +28,4 @@ def build_planner(planner_mode, **kwargs):
     "BezierPlanner",
     "PLANNER_REGISTRY",
     "build_planner",
-]
+]

Model/model_components/trajectory_planning/base.py

@@ -15,10 +15,26 @@ class BasePlanner(nn.Module, ABC):
       flow-matching velocity field). A caller can rely on the first return
       being a fully-formed ``[B, num_timesteps * num_signals]`` trajectory.
 
-    * ``compute_planner_loss()`` runs the training objective and returns
-      ``(loss)``. It owns any decoder-specific scratch tensors
-      (noise samples, target velocities, ...) so they never escape into
-      the caller's scope where they could be paired with the wrong target.
+    * ``compute_planner_loss()`` runs the training objective and returns a
+      ``dict[str, Tensor]`` with at least a ``"loss"`` key — the scalar
+      actually used for backprop. Additional keys are diagnostic /
+      loggable sub-terms specific to the decoder (e.g. ``"velocity_mse"``
+      for FlowMatchingPlanner, ``"imitation_loss"`` for BezierPlanner).
+
+      Returning a dict rather than a bare scalar is deliberate (see #123):
+      it shapes ``compute_planner_loss`` as *the planner's training
+      objective* in general, not specifically "the flow-matching loss" or
+      "the imitation loss". ``train_il`` (or any future training loop)
+      only ever reads ``result["loss"]`` and stays agnostic to which
+      planner/stage produced it. A future stage-3 RL objective can swap in
+      behind this same entry point — returning e.g.
+      ``{"loss": total, "imitation_loss": ..., "reward": ...}`` — blending
+      an imitation anchor with RL terms without forcing a signature change
+      on every caller.
+
+      Each planner owns any decoder-specific scratch tensors (noise
+      samples, target velocities, ...) so they never escape into the
+      caller's scope where they could be paired with the wrong target.
 
     This split mirrors Diffusion Policy / Alpamayo / torchcfm: a polymorphic
     ``forward()`` whose output meaning flips by mode is a footgun (e.g. an
@@ -31,3 +47,53 @@ def forward(self, bev_features, visual_history, egomotion_history,
                 **kwargs):
         """Inference: return ``(trajectory)``."""
         raise NotImplementedError
+
+    @abstractmethod
+    def compute_planner_loss(self, bev_features, visual_history,
+                             egomotion_history, trajectory_target,
+                             training_policy=None, **kwargs):
+        """Training objective. Returns ``dict[str, Tensor]`` with a
+        ``"loss"`` key (see class docstring). A missing implementation now
+        fails loudly at planner-build time instead of silently mis-training
+        (the #115 failure mode).
+
+        Args:
+            training_policy: optional ``DatasetTrainingPolicy``
+                (``Model/training/dataset_policy.py``). Pass the object
+                itself, not pre-extracted scalars — a caller that derives
+                ``signal_scales``/``temporal_decay`` separately and passes
+                them alongside the policy risks the two silently drifting
+                apart (see #124 review: ``signal_scales=(1.0, 1.0)`` vs.
+                production ``(0.79, 0.12)`` gave a 71% loss difference with
+                no error). Implementations that don't use per-signal
+                weighting (e.g. a velocity-MSE objective, where scaling
+                (accel, curvature) channels isn't obviously the same
+                operation as it is on a direct trajectory regression) may
+                accept and ignore it — see FlowMatchingPlanner's docstring
+                for why that's left unresolved rather than guessed at.
+        """
+        raise NotImplementedError
+
+    def _validate_trajectory_target(self, trajectory_target, batch_size, device):
+        """Shared shape/device guard for compute_planner_loss implementations.
+
+        Lifted here (rather than left on FlowMatchingPlanner alone) because
+        every subclass's compute_planner_loss needs the same check, and a
+        missing batch dimension is a silent-wrong-answer bug, not a crash:
+        smooth_l1_loss / mse_loss both broadcast a [T] target across a [B, T]
+        prediction without error, training against the wrong sample for the
+        whole batch. Requires ``self.trajectory_dim`` to be set by the
+        subclass __init__ (num_timesteps * num_signals).
+        """
+        expected = (batch_size, self.trajectory_dim)
+        if tuple(trajectory_target.shape) != expected:
+            raise ValueError(
+                f"trajectory_target must have shape {expected} "
+                f"(batch_size, num_timesteps * num_signals), got "
+                f"{tuple(trajectory_target.shape)}."
+            )
+        if trajectory_target.device != device:
+            raise ValueError(
+                f"trajectory_target must be on the same device as bev_features, "
+                f"got {trajectory_target.device} and {device}."
+            )

Model/model_components/trajectory_planning/bezier_planner.py

@@ -4,6 +4,7 @@
 import torch.nn as nn
 
 from .base import BasePlanner
+from ..losses.trajectory_loss import TrajectoryImitationLoss
 from .reasoning_coupling import ReasoningCoupling
 
 
@@ -47,6 +48,8 @@ def __init__(self, embed_dim=256, num_timesteps=64, num_signals=2,
         self.num_controls = num_controls
         self.egomotion_dim = egomotion_dim
         self.visual_history_dim = visual_history_dim
+        # Shared with BasePlanner._validate_trajectory_target.
+        self.trajectory_dim = num_timesteps * num_signals
 
         # Context aggregation: ego state + visual history + global BEV summary.
         self.ego_state_proj = nn.Linear(egomotion_dim, embed_dim)
@@ -163,3 +166,52 @@ def forward(self, bev_features, visual_history, egomotion_history,
         )
         return trajectory
 
+
+    def compute_planner_loss(self, bev_features, visual_history,
+                             egomotion_history, trajectory_target,
+                             training_policy=None, **kwargs):
+        """SmoothL1 imitation objective (#115), dataset-scale-aware (#124).
+
+        Unlike FlowMatchingPlanner, BezierPlanner's forward() output IS a
+        legitimate direct regression target — there's no sampler/ODE step
+        whose intermediate quantities would leak if regressed against.
+
+        Returns a dict for the same reason as FlowMatchingPlanner (see
+        BasePlanner docstring / #123): keeps train_il agnostic to which
+        planner produced the loss, and reserves room for a future combined
+        objective (e.g. an RL term added alongside "imitation_loss") without
+        a signature change.
+
+        Note: as with FlowMatchingPlanner, reasoning coupling is NOT
+        threaded through this path — forward() is called with its
+        reasoning_latent / reasoning_horizon_tokens defaults, so training
+        optimizes the same context the imitation baseline always has.
+
+        training_policy (#124 review): when given, applies the SAME
+        per-signal scaling and temporal decay TrajectoryImitationLoss
+        applies externally today — plain unweighted SmoothL1 (the
+        training_policy=None fallback) is a real, measured 71% loss
+        difference from the production-scale objective for realistic
+        (accel, curvature) magnitudes, not a rounding difference. When
+        None, falls back to unweighted SmoothL1 for backward
+        compatibility with existing direct callers/tests.
+        """
+        self._validate_trajectory_target(
+            trajectory_target, bev_features.shape[0], bev_features.device
+        )
+        trajectory = self.forward(bev_features, visual_history, egomotion_history)
+
+        if training_policy is not None:
+            weighted_loss_fn = TrajectoryImitationLoss(
+                loss_type="smooth_l1",
+                temporal_decay=training_policy.temporal_decay,
+                signal_scales=training_policy.signal_scales,
+                num_timesteps=self.num_timesteps,
+                num_signals=self.num_signals,
+            ).to(trajectory.device)
+            imitation_loss = weighted_loss_fn(trajectory, trajectory_target)
+        else:
+            imitation_loss = torch.nn.functional.smooth_l1_loss(
+                trajectory, trajectory_target)
+
+        return {"loss": imitation_loss, "imitation_loss": imitation_loss}