Fix/ part of #115 compute planner loss

[FLagbusted]

Summary

Wires BasePlanner.compute_planner_loss (#115) into train_il for real, threads
the dataset-specific training policy through it (#124 review), and exposes
planner_mode so a FlowMatching run can actually be launched.

What's in here

• BasePlanner.compute_planner_loss is now an abstractmethod returning
  dict[str, Tensor] with "loss" — FlowMatchingPlanner does real velocity-MSE,
  BezierPlanner does SmoothL1, both share a trajectory_target shape/device
  validator lifted onto BasePlanner.
• training_policy (the DatasetTrainingPolicy object, not pre-extracted
  scalars) threads from train_il through AutoE2E/ReactiveE2E into
  BezierPlanner.compute_planner_loss, which builds a real weighted
  TrajectoryImitationLoss from it. FlowMatchingPlanner accepts the same
  param for signature parity but does not apply it (documented why —
  scaling accel/curvature on a direct regression vs. a velocity target
  aren't obviously the same operation; left as an open question rather
  than guessed at).
• ReactiveE2E.forward() / AutoE2E.forward(): new explicit
  return_planner_loss=False opt-in. Default False so the ~20 existing call
  sites that already pass trajectory_target expecting a plain trajectory
  tensor back are completely unaffected. train_il is the one caller that
  sets it True.
• train_il: removed the external TrajectoryImitationLoss instantiation;
  loss now read from model(..., return_planner_loss=True,
  training_policy=training_policy)'s result["loss"].
• train_il / wf_train_il: exposed planner_mode / planner_kwargs, threaded
  into AutoE2E(...) and saved into checkpoint_config. Found while preparing
  to actually request a training run — neither function had ANY way to
  select flow_matching before this; every run silently trained bezier
  regardless of intent, which would have made the compute_planner_loss fix
  above untestable in practice.
• Rebased onto main post-Add route-conditioned navigation inputs for KITScenes (#149) #161 (route-conditioned navigation inputs) — one
  contained conflict in reactive_e2e.py's forward() signature (both sides
  added params to the same line), resolved by keeping Add route-conditioned navigation inputs for KITScenes (#149) #161's new positional
  route params and adding ours alongside. Everything else, including this
  PR's own dispatch logic, auto-merged cleanly onto the new
  NavigationEncoder/fused_features flow.
• Removed a stale RuntimeWarning in build_planner() that claimed
  flow_matching "is NOT correctly trainable via train_il" — no longer true
  as of this PR.

Why

The actual #115 bug, traced to its source: train_il called
model(mode="train", trajectory_target=target, ...), but trajectory_target
was silently absorbed as an inert kwarg all the way down to
ReactiveE2E.forward(), which called planner.forward() unconditionally
regardless of mode. FlowMatchingPlanner's Euler-from-noise rollout got
SmoothL1-regressed against the target externally instead of ever running
its real velocity-MSE objective.

Not in this PR

FlowMatching + #76's TrajectoryComplianceScorer as a third benchmark row —
evaluate_kitscenes_benchmark_checkpoint doesn't call the scorer at all
today, so that's real work I haven't started. Will follow as its own PR
once a FlowMatching checkpoint actually exists to test it against.

Testing done

test_planner_loss.py (incl. training_policy regression tests — proves the
policy actually changes BezierPlanner's loss, matching the 71% divergence
measured in review), test_auto_e2e.py (incl. tests exercising
return_planner_loss=True end-to-end: loss dict shape, gradient reaches the
Backbone, training_policy changes the number), test_workflow_training_
lifecycle.py (source-inspection tests for both the training_policy wiring
and the new planner_mode wiring), test_reasoning_integration.py. All
passing, ruff clean, verified against main post-#161 rebase.

Checklist

• Conventional Commit PR title
• References FlowMatchingPlanner has no training objective wired — train_il regresses an Euler-from-noise rollout #115 / this thread's discussion
• Existing callers unaffected (return_planner_loss defaults False)
• Tests pass, ruff clean
• Rebased onto current main (post-Add route-conditioned navigation inputs for KITScenes (#149) #161)
• FlowMatching is actually launchable now (planner_mode exposed)

[gcordova10]

Thanks for this. Wiring compute_planner_loss as the planner's training objective — rather than
"the flow-matching loss" — is the right call, and it does two things at once: it closes the #115
failure where train_il was regressing an Euler-from-noise rollout toward the conditional mean, and
it gives the multi-sample scorer in #76 a velocity field actually trained for transport, so the K
samples come out meaningfully diverse instead of clustered around the mean. Making it
@abstractmethod so a missing implementation fails at build time hardens the same bug.

I went through the branch. Two things.

1. BezierPlanner.compute_planner_loss doesn't validate trajectory_target, and a target with
the batch dimension missing trains against the wrong data silently.

FlowMatchingPlanner validates:

self._validate_trajectory_target(trajectory_target, B, bev_features.device)   # :334

but that helper is defined on FlowMatchingPlanner itself (flow_matching_planner.py:159), not on
BasePlanner, so the Bezier path goes straight into the loss:

trajectory = self.forward(bev_features, visual_history, egomotion_history)
imitation_loss = torch.nn.functional.smooth_l1_loss(trajectory, trajectory_target)

On your branch, with two samples whose targets differ:

planner = BezierPlanner().eval()
B, T = 2, 128
bev, vis, ego = torch.zeros(B,256,8,8), torch.zeros(B,896), torch.zeros(B,256)
tgt = torch.stack([torch.full((T,), 1.0), torch.full((T,), 5.0)])   # [B, T]

planner.compute_planner_loss(bev, vis, ego, tgt)["loss"]      # 2.494182
planner.compute_planner_loss(bev, vis, ego, tgt[0])["loss"]   # 0.494356  <- no error

smooth_l1_loss broadcasts sample 0's target across the whole batch: an 80% error, and the only
signal is a UserWarning that disappears in a training log. A [B, 64, 2] target does raise, so it
is specifically the missing-batch-dimension case that slips through.

To be clear about the blast radius, though: this cannot fire today. train_il always passes [B, T] — the
loader emits trajectory_target per sample and the collate adds the batch dimension — and right now
the only callers of compute_planner_loss are the tests. What makes it worth a guard is that the
entry point is about to get its first real callers: the train_il wiring in point 2, and later a
stage-3 objective. Cheap now, and the failure mode is a silent wrong answer rather than a crash.

I would lift the guard to BasePlanner so both planners share one:

# base.py
def _validate_trajectory_target(self, trajectory_target, batch_size, device): ...

# bezier_planner.py — first line of compute_planner_loss
self._validate_trajectory_target(trajectory_target, bev_features.shape[0], bev_features.device)

and add a Bezier counterpart to test_invalid_trajectory_target_shape_raises: TestBezierPlannerLoss
has three tests today and none covers target shape, while the FlowMatching class does
(test_planner_loss.py:100).

2. Nothing calls compute_planner_loss yet, so the #115 bug survives this merge.

#115's proposed fix has four steps. This PR does the first two — implement the objective on both
planners, and make it @abstractmethod. The third one is the one that actually changes behaviour:

In train_il, call the planner's training hook when available instead of
loss_fn(forward_output, target).

That is still open. On this branch and on main, train_il runs:

loss_fn = TrajectoryImitationLoss(loss_type="smooth_l1")
...
pred = model(..., mode="train", trajectory_target=target)
loss = loss_fn(pred, target)          # workflows.py:441 here, :805 on main

compute_planner_loss appears nowhere in workflows.py. So after this lands, flow matching is
still trained by regressing the Euler-from-noise rollout — exactly the failure #115 describes. The
correct objective exists and is tested, but nothing reaches it.

Since the description says Fixes #115, merging would close the issue while its bug is still live.
Two options that both seem fine: add the train_il wiring here, or keep this PR as the planner-side
half and change the keyword to Part of #115 so the issue stays open for step 3.

If you do add the wiring here, it is worth rebasing first: the branch is based on 8ac3240 (7 July,
before the #119 merge) and Platform/pipelines/workflows.py has moved a lot since — 823 insertions,
including gradient accumulation in train_il (3b9f1ac) and the train/val split threading
(f9b89a7). The PR as it stands merges cleanly, since it doesn't touch that file; it is specifically
the train_il change that would land on top of code that has changed underneath.

3. Should trajectory_target be optional? A question rather than a request, and it comes from
the stage-3 side.

The reason I care about this entry point is #123: closed-loop RL needs somewhere to put a
reward-based objective, and your dict[str, Tensor] return is what makes that additive rather than
a refactor — the RL objective returns {"loss": total, "imitation_loss": ..., "reward": ...} and
train_il never needs to know which stage produced it. That part already works.

The one piece that doesn't quite fit is the signature: trajectory_target is a required positional
argument, and an objective computing reward over rollouts may not have a target trajectory at all —
it would pass None by convention or route the rollout through **kwargs. Since the contract is
being set now and there are no callers yet, making it optional may be cheaper than unwinding a
convention later:

def compute_planner_loss(self, bev_features, visual_history, egomotion_history,
                         trajectory_target=None, **kwargs):

I am not asking for it in this PR — the goal right now is a minimum architecture that can train,
and this is not that. It is only that the default is a one-line decision while there are no callers,
and an awkward migration once there are.

---

@riita10069, two of these are really yours to call rather than mine, and I would rather ask than
assume:

• Whether Fixes #115 should stay. You wrote the issue, so you know better than I do whether
  the planner-side half closes it or whether step 3 (train_il) is part of what "fixed" means.
• Where the guards belong, and whether trajectory_target should be optional. This is the
  planner contract you designed. One thing that makes it timely: your
  feat/scale-full-episode-training branch adds _validate_initial_noise to
  FlowMatchingPlanner as well, so with the guard I suggest above there would be three validators
  living on the subclass. That might be an argument for putting the shared ones on BasePlanner,
  or it might not — but it seems better decided by you than settled inside this PR.

Neither blocks merging from my side. The one I would fix before it lands is the first: it is a
silent-wrong-answer path rather than a crash.

[FLagbusted]

@gcordova10 Thanks for the detailed review! I will look into it ,
One question on the first point: since BezierPlanner uses direct imitation rather than the flow-matching objective, do you still think it makes sense to move the validation into BasePlanner and enforce it for both planners, or would you prefer to keep the validation specific to the planners that need it?

[FLagbusted]

Thanks for the thorough read — point 1 is real, confirmed it with your
exact repro before fixing: tgt[0] silently gave loss=0.494356 instead of
raising. Fixed by lifting _validate_trajectory_target to BasePlanner
(BezierPlanner now sets self.trajectory_dim = num_timesteps * num_signals
so the shared check has what it needs) and added the mirrored test to
TestBezierPlannerLoss. 14/14 pass, ruff clean.

On point 2 — agreed this PR shouldn't claim "Fixes #115" as-is. I'd
rather keep this scoped to the planner-side objective and change the
title to "Part of #115" than pull in the train_il wiring here, especially
given workflows.py has moved 823 insertions since this branch was cut —
rebasing that file properly deserves its own PR rather than riding along
with this one. @riita10069 — deferring to you on whether that split works,
since it's your issue to close.

On point 3 (optional trajectory_target) — no objection to trajectory_target=None
for the reasons you gave; also deferring to @riita10069 given the
_validate_initial_noise addition on their branch suggests the validator
question and the optional-target question are related calls worth making
together.

[riita10069]

Thank you for following through on #115 so quickly — and for shaping the API exactly as discussed in #123. The dict[str, Tensor] contract with a required "loss" key, the @abstractmethod enforcement so a missing objective fails at build time instead of silently mis-training, and the velocity-MSE implementation itself (random t, u_t = (1-t)x0 + t·x1, regressing v_theta against x1 - x0, with a regression guard proving the Euler loop never enters the loss path) are all exactly right. The test coverage in test_planner_loss.py is thorough. We genuinely appreciate the care here.

Rather than treating this as a PR to review in isolation, we would like to propose something bigger: let's take this all the way to KITScenes Benchmark numbers, together.

Here is the context. The production model does not use FlowMatchingPlanner yet — AutoE2E defaults to planner_mode="bezier", and all checkpoints trained through our Flyte pipeline so far are Bezier-based. build_planner currently emits a RuntimeWarning for flow_matching because the training loop regresses an Euler-from-noise rollout with SmoothL1, which is not the flow-matching velocity-MSE objective. The decision of which planner to adopt should be made on a fair Bezier vs FlowMatching comparison, and that comparison is impossible today precisely because FlowMatching has no correct training objective wired in. This PR is the piece that makes evaluating FlowMatching possible at all.

So here is what we would like to ask, as one arc of work:

1. Fix the merge regressions in the current diff. Compared with today's main, this branch removes ReasoningCoupling from both planners (the production enable_reasoning=True path depends on it), the zero-init visual_history_proj in BezierPlanner (the Proposal: Feature Reconstruction Loss for FutureState (Auxiliary Self-Supervised Task) #13 strictly-additive World Model fix), and initial_noise support in FlowMatchingPlanner.forward(). I suspect the "sync upstream main" merges resolved conflicts in favor of the branch's older copies — could you redo those so main's features are preserved?

2. Wire compute_planner_loss into train_il so FlowMatching actually trains with the correct objective (and the RuntimeWarning can be retired). One design point here: the current TrajectoryImitationLoss carries dataset-policy-derived temporal_decay and signal_scales weighting, so BezierPlanner.compute_planner_loss needs a way to receive those.

3. Train a FlowMatching checkpoint and put numbers on the board. Main now has a checkpoint-only retrospective evaluator (Model/evaluation/kitscenes_benchmark.py, runnable via the wf_evaluate_kitscenes_benchmark Flyte workflow) reporting 3s/5s ADE/FDE for any checkpoint, so Bezier vs FlowMatching can be compared on equal footing. And since your GoalFlow-style scorer in feat(trajectory_planning): add BEV-only multi-sample trajectory scorer and shifted inference schedule #76 becomes meaningful exactly when it re-ranks samples from a correctly trained velocity field, we would love to see FlowMatching+scorer as a third row in that comparison.

On our side, we are happy to carry the infrastructure half of this: running the training and benchmark evaluations on our GPU cluster (we know GPU access has been a constraint for you), and reviewing the calibration questions from #76 against the actual map renderer (Model/data_parsing/kit_scenes/map.py). If the wiring in (2) turns out to be better split into a follow-up PR, that is fine too — the goal is the benchmark table, not any particular PR boundary.

This PR plus #76 together form a coherent research arc: correct objective → trained FlowMatching checkpoint → multi-sample scoring → benchmark comparison. We would be glad to walk that arc with you. Thank you again for driving this.

[FLagbusted]

Thanks for your time,
And for detailed feedback and for the kind words for my broken work😅.

You're right—I did mess up the regression on point 1. I'm fixing that now so the current mainline behavior is preserved, including restoring ReasoningCoupling, the zero-init visual_history_proj, and initial_noise support.

As for the benchmark work, I'm already working on that. I first tried the k-sample trajectory visualization on a simple front-view camera sample from a random NVIDIA physical dataset just to validate the sampling pipeline. It correctly produces the random k trajectories in top-down (x, y) space. As you suggested, the next step is to move this over to the KITScenes dataset, which was already my plan after validating the pipeline.

I'll also try to keep the comparison as apples-to-apples as possible by evaluating Bezier and FlowMatching using the same KITScenes dataset split, the same checkpoint evaluation pipeline, and the same sampling/scoring budget. That should make any performance differences attributable to the planners themselves rather than differences in the evaluation setup.

As long as my laptop holds up( ;__;). Thanks for offering to handle the training ans benchmark evaluations on your GPU cluster—that'll definitely make this much more practical. Looking forward to seeing how FlowMatching stacks up once everything is wired together. 🙌

[gcordova10]

@FLagbusted — answering your earlier question, though you already landed on the same answer: yes,
BasePlanner is the right home. The reason it applies to Bezier too is that the guard is about the
target contract — a [B, T] tensor whose batch matches the features — not about the objective that
consumes it. Direct imitation and velocity-MSE differ in what they do with the target, not in what
shape it has to be, and it is the direct-imitation path that fails silently, so it needs the check
most. Lifting it plus the mirrored test looks right to me.

One leftover from point 2: the title now says "part of #115", but the body still has Fixes #115.
GitHub still lists #115 under this PR's closing references, so merging as it stands would close the
issue. A one-word edit to the body is all it needs.

On point 2 of the arc, one thing worth settling before the wiring — I ran it rather than guessed,
because it fails the same quiet way as the target bug. Since #135, TrajectoryImitationLoss defaults
to _DEFAULT_SIGNAL_SCALES = (1.0, 1.0) (trajectory_loss.py:19), with the note that production
training must pass the explicit dataset policy rather than inherit the L2D-measured (0.79, 0.12).
Feeding the same prediction and target through the real loss both ways, on a synthetic batch with
realistic per-signal magnitudes:

signal_scales (0.79, 0.12)  ->  loss = 0.133080
signal_scales (1.0, 1.0)    ->  loss = 0.038403     # 71% lower, no error, no warning

DatasetTrainingPolicy already carries those scales and temporal_decay
(Model/training/dataset_policy.py:32). So if compute_planner_loss is wired into train_il
without a route for that policy, Bezier training optimises a different objective than it does today
and nothing reports it — the exact number depends on the data, but the divergence is silent either
way. Passing the policy object into the planner rather than re-deriving the scalars keeps train_il
and the hook computing the same thing.

Your apples-to-apples plan is the right instinct. If it is useful, the piece I could bring is on the
measurement side rather than the wiring: alongside 3s/5s ADE/FDE, the reasoning intervention delta
(#109, in main) answers a different question — whether the reasoning branch is used at all rather
than merely tolerated. I have a positive control for it that forces the coupling open and requires
the delta to register and to grow with coupling strength, so it cannot silently read zero. Cheap to
add as a column if the table is being built anyway, and I would bring it as its own small PR so it
does not enlarge this one.

[FLagbusted]

Confirmed the regression — you were right on all three: ReasoningCoupling, the #13 zero-init, and initial_noise all got silently dropped by the earlier sync merges (git's 3-way merge took the "unchanged" side since the loss was already baked into the merge-base). Redid the merge + restored all three in a follow-up commit; diffed against main to confirm nothing else moved. 66/67 tests pass (1 GPU-only skip), ruff clean. Also fixed the Fixes #115 → Part of #115 body text @gcordova10 caught. Pushing now — will go on the KITScenes checkpoint again,

Really appreciate all the time you gave it

[FLagbusted]

Thank you guys for helping out,
@riita10069 @gcordova10 — parts 1 and 2 of the arc are done.

compute_planner_loss is wired into train_il for real now. Traced the actual call chain (train_il → AutoE2E.forward() → ReactiveE2E.forward() → planner.forward()) and found trajectory_target was reaching all the way down but getting silently absorbed as an inert kwarg at every level — nothing ever used it to trigger a planner's real training objective. Fixed with an explicit opt-in (return_planner_loss=True), not on trajectory_target's mere presence — that alone would have broken ~20 existing call sites across the test suite that already pass it expecting a trajectory tensor back. Caught that the hard way via a failing test before it went anywhere

@gcordova10 — for your signal_scales concern, what i did is : training_policy (the object, not derived scalars) now threads from train_il through AutoE2E/ReactiveE2E into BezierPlanner.compute_planner_loss, which builds a real TrajectoryImitationLoss from it. Added a regression test proving the policy actually changes the loss. FlowMatchingPlanner accepts the same param for signature parity but doesn't apply it yet — scaling (accel, curvature) on a direct regression vs. on a velocity target aren't obviously the same operation, so I left that as an explicit, documented open question rather than guess.

@riita10069 — ready to take you up on the GPU offer for part 3. Whenever your cluster has room: a FlowMatching checkpoint trained through train_il as it stands now, benchmarked via wf_evaluate_kitscenes_benchmark against the existing Bezier baseline. The FlowMatching+#76-scorer third row will follow as its own commit once that checkpoint exists — evaluate_kitscenes_benchmark_checkpoint doesn't call the scorer at all today, so that's real work I haven't started, and part of it (the DAC score) isn't meaningful until #148/#149 settle anyway. Wanted the first two rows moving rather than block on that.

Again thank you guys for help and it really helps me understanding the Architecture more. (;__;)

[riita10069]

I will check it.

[riita10069]

I’m sorry for the disruption. Environment-specific infrastructure identifiers, including instance IDs, were inadvertently committed to the repository history, so we had to rewrite the history and force-push the cleaned branches. This caused this pull request to close automatically.

I restored your commits in replacement PR #172 from the sanitized history, with the original commit authorship preserved:
#172

Before continuing work, please back up any uncommitted changes outside the repository, remove your existing local clone, and clone the repository again. Reusing or pushing from an old clone could reintroduce the removed history.

I apologize for the unexpected closure and the additional review noise.