Conversation
Adds decoder_type=plantf which replaces the diffusion decoder (DiT + DPM-Solver) with a one-shot multi-modal regression head ported from planTF (Cheng et al., ICRA 2024), sharing the encoder and the input/output contract. Training uses winner-takes-all regression over num_modes ego candidates plus mode classification, reusing DP's lat/lon/heading losses, penalty losses and turn-indicator heads. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Isamu Yamashita <isamu.yamashita@tier4.jp>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Isamu Yamashita <isamu.yamashita@tier4.jp>
The export pipeline assumed the diffusion decoder (decoder.dit), so checkpoint exports failed with decoder_type=plantf. The planTF head is one-shot, so it exports full.onnx, encoder.onnx and a single-call decoder.onnx (encoding -> prediction, probability, turn_indicator_logit) with no separate turn-indicator graph. Mode selection now uses gather so the batch axis stays dynamic under tracing, and torch2onnx.py validates the planTF split graphs against torch. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Isamu Yamashita <isamu.yamashita@tier4.jp>
The planTF head ignores sampled_trajectories / ego_current_state / delay and the legacy exporter prunes unused graph inputs, so the exported full.onnx had 14 inputs while the Autoware node feeds all 17 by name (both ONNX Runtime and TensorRT reject names absent from the graph). PlanTFFullONNXWrapper now anchors the three inputs with a zero-valued scalar-slice residual, keeping the node unchanged and the outputs bit-identical. Documented the node ONNX compatibility investigation in the planTF integration plan. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Isamu Yamashita <isamu.yamashita@tier4.jp>
Under winner-takes-all training, rarely-winning modes receive almost no regression gradient; with Xavier output weights they keep emitting white-noise trajectories, which the argmax(pi) mode selection can pick at inference (randomly jagged outputs). Zero-init (same convention as Decoder.dit.final_layer) makes every mode start at the normalized-space mean, so an undertrained mode degrades into a smooth prior trajectory instead of noise. Analysis and further mitigations in docs/plantf_dead_mode_improvement.md. Requires retraining to take effect. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Isamu Yamashita <isamu.yamashita@tier4.jp>
Retraining with modes=1 showed the planTF head collapsing to a near-stationary oscillation around the origin while the averaged validation loss stayed unremarkable. Three changes to detect and counter that: - compute_trajectory_progress_metrics reports endpoint progress, path length ratio, FDE, mean second difference and speed MAE for the selected ego prediction, each additionally binned by GT endpoint distance (stop / slow / move) so stop-heavy data cannot mask a moving-scene failure. Logged as valid_traj/* for every decoder type. - An endpoint FDE loss in metres (coeff_endpoint_fde_loss, planTF only) adds a direct forward-progress signal that the 80-step pointwise decomposition dilutes. - The longitudinal velocity down-weighting is now opt-in for the planTF head (plantf_use_lon_velocity_weight, default off): dividing the lon error by up to |v| suppresses exactly the progress signal the regression head needs. Analysis and the recommended ego_history_dropout_rate=0 ablation are documented in docs/plantf_dead_mode_improvement.md. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Isamu Yamashita <isamu.yamashita@tier4.jp>
The endpoint FDE loss was computed in metres (endpoint_diff * std). With ego xy std = 20 that made it ~60x the normalized-space per-timestep position loss, so at coeff=1.0 it dominated the total and the model collapsed each path to a straight terminal-homing line (the endpoint term constrains only the endpoint, not the shape), losing the gentle curves GT and the diffusion head produce. Compute it as a squared error in the normalized state space instead, matching ego_planning_loss, so the per-timestep losses still shape the trajectory. See docs/plantf_dead_mode_improvement.md. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Isamu Yamashita <isamu.yamashita@tier4.jp>
The endpoint loss is not needed for forward progress: once the lon velocity down-weighting is off, the per-timestep position loss teaches progress at every step including the endpoint, so the endpoint term only double-weights a point the position loss already covers. Because it constrains the endpoint but not the path shape, it also adds a straight-line bias (the metre-scale version made the model ignore curves entirely). Default coeff to 0; the term is still computed for logging and can be opted back in. See docs/plantf_dead_mode_improvement.md. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Isamu Yamashita <isamu.yamashita@tier4.jp>
A/B evaluation at modes=6 showed the planTF-specific loss shaping made results worse (ADE 6.80 vs 5.20 without it, and a straight-line prediction that cannot follow gentle curves). The dango collapse it was meant to fix was actually caused by modes=1 (single-mode L2 = mean regression) and was already resolved by modes>=2 with WTA, so the lon velocity down-weighting was not at fault. Restore plantf_use_lon_velocity_weight=True (endpoint loss already defaulted off), making the planTF loss identical to the diffusion head. The progress metrics are kept for diagnosis. See docs/plantf_dead_mode_improvement.md. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Isamu Yamashita <isamu.yamashita@tier4.jp>
The planTF head regressed 80 absolute waypoints independently per timestep, which A/B pinned as the cause of the residual straight-line + comb-jitter failure (second_diff_m_move ~0.5 vs GT ~0.007, fde_m_move ~20m) that survived reverting every loss change. It now regresses per-step displacement when use_velocity_representation is set, and _decode integrates it with a cumulative sum into absolute waypoints before mapping to the normalized state space. The cumsum gives the temporal continuity a per-timestep absolute regression lacks (adjacent waypoints differ by exactly the predicted displacement, killing the comb jitter) and keeps forward progress from collapsing to a stall. Because the integration is confined to _decode, WTA selection, loss_func, mode metrics, prediction and ONNX export are all unchanged (1-epoch smoke exports all three ONNX graphs, no nan loss). The NotImplementedError guard is removed. Still applies the absolute lat/lon/heading loss; switching to the HDP hybrid loss is the next step if the cumsum gradient proves unstable. See docs/plantf_dead_mode_improvement.md. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Isamu Yamashita <isamu.yamashita@tier4.jp>
min_fde_k=3m vs top1 fde=20m confirmed the residual planTF error is a
mode-selection failure, not a regression failure: a mode that reaches
the curve endpoint exists but argmax(pi) never picks it. plantf_mode_usage
only shows the selected-mode distribution; whether collapsing to modes=1
is safe depends on the oracle (best-ADE) mode distribution, which was not
logged. Add plantf_oracle_usage_{k}: concentrated => data is effectively
uni/bi-modal and modes=1 suffices, spread => genuinely multi-modal and one
mode would average the trajectories away, requiring multi-mode output.
Computable by re-running validation on the existing model (no retrain).
See docs/plantf_dead_mode_improvement.md.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Isamu Yamashita <isamu.yamashita@tier4.jp>
Iterating on settings (modes, velocity representation, endpoint weight, alpha) only changed the shape of the collapse — dango, straight-line, reverse-direction, comb jitter — because the loss itself had diverged from the original planTF: compute_plantf_training_loss grafted the DP diffusion-decoder loss (lat/lon/heading L2 decomposition + longitudinal velocity down-weighting + timestep weighting + endpoint term) onto the one-shot regression head, which destabilized it. Original planTF (jchengai/planTF) uses a plain smooth L1 on the WTA best mode over all channels, uniform across timesteps, plus a cross-entropy on the detached best mode — no velocity normalization, no timestep weighting, no lat/lon split. Restore exactly that for ego and neighbor regression. The turn-indicator and road-border/collision penalties are kept for the DP/Autoware interface. Recommended config is num_modes>=2, use_velocity_representation=False (original is absolute-coordinate), and the default learning_rate 1e-4 (1e-5 underfits in 20 epochs). Verified: unit test asserts the ego loss equals an independently computed smooth_l1 of the WTA best mode; 1-epoch smoke completes and exports all three ONNX graphs. See docs/plantf_dead_mode_improvement.md. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Isamu Yamashita <isamu.yamashita@tier4.jp>
The planTF head kept diverging (speed_mae_stop 80, then 20 with a lower lr) and producing start-point-scattered trajectories at settings where the diffusion head is stable. The difference is structural: the diffusion decoder pins the current state into its prediction sequence, so its output is anchored to the current motion; the planTF head regresses 80 absolute waypoints from the ego token alone and never sees the current position/velocity, so the regression is unanchored and scatters/diverges. Feed ego_current_state[:, 4:10] (vx, vy, ax, ay, steering, yaw_rate, already normalized) into PlanTFTrajectoryHead: project to hidden and add to the ego token before the mode branch. This anchors the prediction to the current motion the way the diffusion head's pinned state does, so the head can output a stop when the current speed is 0 and the correct forward extent/curve otherwise. Controlled by plantf_use_ego_state_in_head (default True). The full ONNX graph gets this via model inputs automatically; the split decoder graph still needs ego_current_state added as a second input (follow-up). Verified by a unit test (ego state changes the prediction when enabled, no-op when disabled) and a 1-epoch smoke that exports all three ONNX graphs. See docs/plantf_dead_mode_improvement.md. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Isamu Yamashita <isamu.yamashita@tier4.jp>
…a options (A2/C1) Three planTF-alignment changes from the roadmap toward closing the gap with original planTF. All except B1 default OFF (no behavior change unless enabled). - B1: AdamW weight decay (config weight_decay, default 1e-4) applied only to matmul weights; biases and LayerNorm/BatchNorm/GroupNorm/Embedding are excluded via param groups. Previously torch AdamW's default 0.01 hit all params including norms/biases. Affects both decoder types. - A2 (plantf_ego_state_token, default False): replace the ego encoder token with an embedding of the current motion state (vx,vy,ax,ay,steer,yaw_rate) plus state_dropout (default 0.75), matching original planTF's use_ego_history=false path. Anchors the prediction to the current motion the way the diffusion decoder's pinned state does. - C1 (plantf_input_delta, default False): feed agent (ego+neighbor) history as consecutive-frame xy deltas. Unit tests (17) cover A2/C1 forward contract and A2+C1 backward; the param groups are verified to cover all parameters. See docs/plantf_original_comparison_and_roadmap.md. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Isamu Yamashita <isamu.yamashita@tier4.jp>
Two fixes from the ablation analysis (docs/plantf_original_comparison_and_roadmap.md). - lr schedule: the misnamed CosineAnnealingWarmUpRestarts held the lr CONSTANT after warmup (MultiplicativeLR lambda=1.0). Every model, diffusion included, peaked during warmup and then collapsed, and the checkpoint_trajectory viz showed comb/straight-line predictions past a few epochs. Replace the post-warmup phase with a real CosineAnnealingLR decay to eta_min=1e-6, matching original planTF. Also removed the redundant final-10-epochs manual lr step-down in train.py (the cosine now owns the decay). - planTF smoothness_loss: the head regresses 80 absolute waypoints independently per timestep, producing comb jitter (roughness ~2 vs GT ~0.007) even when ADE/FDE look fine. The diffusion head follows curves smoothly at the same settings, so this is structural. Penalize the xy second difference of the best mode (coeff_smoothness_loss, default 0). Both affect training only. Tests pass (17); cosine schedule verified to warm up then decay. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Isamu Yamashita <isamu.yamashita@tier4.jp>
Add --enable_checkpoint_viz / --enable_onnx_export toggles (both default True to preserve behavior). Disabling viz makes matplotlib/visualize_input a lazy import inside render_checkpoint_trajectory_figure, so training runs on environments without matplotlib. Disabling ONNX skips export at each checkpoint save. Lets other teams run large-scale experiments without the optional dependencies. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Isamu Yamashita <isamu.yamashita@tier4.jp>
…ker, NLL/tail losses For the upcoming comprehensive experiments. All default to the current behavior and are independently combinable (docs/plantf_head_development_notes.md §9): - plantf_head_type=cross_attn: K learnable mode queries cross-attend to ALL encoder tokens (map/agents/route) instead of reshaping the single pooled ego token, giving the 80-point head the scene context the bottleneck lacks. - plantf_route_rerank (+topk): at inference, pick the ego mode by route adherence among the top-k pi modes instead of argmax(pi). Recovers oracle-ish modes for the smooth velocity-rep head at zero training cost. Validation forward path only; the ONNX deploy graph (no route_lanes) keeps argmax. - plantf_use_laplace_nll: head predicts a per-point log-scale; ego regression uses Laplace NLL instead of smooth-L1 (planTF's probabilistic regression; calibrates tail uncertainty and mode confidence). - plantf_tail_weight / plantf_smoothness_tail_weight: weight the ego regression / curvature penalty toward the late horizon where divergence is worst. 9 new tests cover head shapes (incl. predict_scale), cross-attn forward+loss, NLL/tail losses changing the loss, and route re-rank fallback + identity-norm run. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Isamu Yamashita <isamu.yamashita@tier4.jp>
…e input The mini dataset's goal_pose is in global/map frame while the Autoware C++ node feeds it correctly ego-transformed, so goal_pose is a train/deploy-mismatched input (a likely contributor to the Autoware "way-off/dango" behavior). With plantf_mask_goal_pose=True the encoder zeroes goal_pose, so the model plans from route_lanes (which already encodes the goal) instead. Default False (unchanged). Bakes into the ONNX graph, so Autoware needs no change to deploy it. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Isamu Yamashita <isamu.yamashita@tier4.jp>
The flat mlp head regresses 80 waypoints as independent linear read-outs (Linear(hidden, T*C)) — no output-side temporal inductive bias; coherence is only bolted on via velocity-cumsum + smoothness penalty. Add two combinable head types that put temporal structure in the architecture, both keeping the mlp-style single-ego-token mode formation (deploy robustness) and the exact [B, K, T, 4] output contract (velocity integration, WTA, zero-init unchanged): - `--plantf_head_type basis` (+ `--plantf_basis_control_points`, default 8): regress Bezier control points expanded to T via a fixed Bernstein basis buffer. Structurally C-infinity smooth (partition of unity + convex hull), a candidate replacement for the smoothness penalty. ONNX = one extra matmul. - `--plantf_head_type gru`: a GRU unrolls the waypoints (non-autoregressive: fed a learned per-step temporal embedding + mode context, so it exports as a single ONNX GRU op with no feedback loop). Experimental — RNN heads are more deploy-fragile (batch!=1 GRU export caveat; Autoware runs batch=1). Both verified: full 17-input ONNX export with exact PyTorch parity (max|d|=0). 34 tests pass (added basis/gru shape, zero-init, contract, smoothness, velocity+backward, and bezier partition-of-unity/endpoint tests). Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: Isamu Yamashita <isamu.yamashita@tier4.jp>
Different converter versions store neighbor_agents_future as 3 channels
(x, y, heading) or 4 (x, y, cos, sin). train_epoch applies heading_to_cos_sin
AFTER the DataLoader collate, so a batch mixing a 3-channel dataset (mini) with
a 4-channel dataset (odaiba) fails to collate ("Trying to resize storage that is
not resizable"). Convert to 4 channels in __getitem__, before collation, so
differently-generated datasets can be combined for training.
Padded (all-zero) rows stay all-zero, so downstream masking is unchanged, and
already-4-channel data is returned untouched — the result matches what
train_epoch would have produced for 3-channel data, so single-dataset behavior
is preserved. Verified: a mixed odaiba(4ch)+mini(3ch) batch now collates to
[B, 320, 80, 4]. 34 plantf tests still pass.
Co-Authored-By: Claude <noreply@anthropic.com>
Signed-off-by: Isamu Yamashita <isamu.yamashita@tier4.jp>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What changed
PlantF versus the DP head
Current production baseline
The current baseline uses a two-GPU, mode-one MLP head:
decoder_type=plantf
plantf_head_type=mlp
num_modes=1
plantf_relative_xy=True
plantf_use_ego_state_in_head=True
use_velocity_representation=False
plantf_tail_weight=1.0
coeff_smoothness_loss=1.0
plantf_smoothness_tail_weight=1.0
freeze_encoder_epochs=2
global_batch_size=96
head_lr=5e-4
encoder_lr=5e-5
use_ema=True
The complete production command, normalization requirements, checkpoint behavior, and validation checklist are documented in:
Validation
diagnostics.
Deployment notes