Skip to content

Fix/ part of #115 compute planner loss - #172

Open
riita10069 wants to merge 4 commits into
mainfrom
restore/pr-124-compute-planner-loss
Open

Fix/ part of #115 compute planner loss#172
riita10069 wants to merge 4 commits into
mainfrom
restore/pr-124-compute-planner-loss

Conversation

@riita10069

@riita10069 riita10069 commented Jul 31, 2026

Copy link
Copy Markdown
Collaborator

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

FLagbusted and others added 4 commits July 27, 2026 20:27
BasePlanner.compute_planner_loss is now @AbstractMethod, returning
dict[str, Tensor] with a required 'loss' key rather than a bare scalar
(per @gcordova10's #123-motivated ask) — train_il only reads result['loss']
and stays agnostic to which planner/stage produced it; a future stage-3 RL
objective can return {'loss': total, 'imitation_loss': ..., 'reward': ...}
behind the same entry point with no signature change.

FlowMatchingPlanner: proper flow-matching velocity-MSE objective
(construct_training_data -> _v_theta -> MSE against target velocity).
Replaces the previous train_il path, which regressed forward()'s
Euler-from-noise rollout against a fixed target — pushing the network
toward the conditional mean rather than training the velocity field.

BezierPlanner: SmoothL1 on forward()'s output (a legitimate direct
regression target, unlike FlowMatchingPlanner's ODE rollout).

13 new tests, including a regression guard proving forward() (inference)
and compute_planner_loss (training) are fully decoupled — same seed, same
output, before and after exercising the loss path.

Fixes #115

Signed-off-by: FLagbusted <Justthefourofus@proton.me>
Signed-off-by: FLagbusted <justthefourofus@proton.me>
Signed-off-by: Flagbusted <justthefourofus@proton.me>
…al_noise lost in sync

Flagged by @riita10069. Both planners now match upstream main exactly
except for compute_planner_loss (#115) and the trajectory_dim attribute
BasePlanner._validate_trajectory_target needs.

Full suite re-run clean: test_planner_loss.py, test_trajectory_planning.py
(incl. initial_noise tests), test_reasoning_coupling.py, test_bezier_planner.py
— 66 passed, 1 GPU-only skip. ruff clean.

Signed-off-by: FLagbusted <justthefourofus@proton.me>
See prior commits on this branch and the PR thread for full context.
train_il now passes return_planner_loss=True + training_policy=training_policy
into model(...) instead of computing loss externally via
TrajectoryImitationLoss. return_planner_loss is a new, explicit opt-in
flag (default False) on AutoE2E.forward()/ReactiveE2E.forward() -- every
other existing caller (~20 call sites across the test suite) is unaffected.

Signed-off-by: FLagbusted <justthefourofus@proton.me>
FLagbusted added a commit to FLagbusted/auto_fsd that referenced this pull request Aug 12, 2026
…warefoundation#168

Implements exploratory single-GPU training of the Reactive branch on
contributor machines without the full cluster pipeline. Addresses
riita10069's request in autowarefoundation#168 for source code rather than one-off results.

Relates-to: autowarefoundation#168, autowarefoundation#172
Signed-off-by: FLagbusted <justthefourofus@proton.me>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants