diff --git a/diffusion_planner/diffusion_planner/model/diffusion_planner.py b/diffusion_planner/diffusion_planner/model/diffusion_planner.py index f7abb42eb..df0be0f41 100644 --- a/diffusion_planner/diffusion_planner/model/diffusion_planner.py +++ b/diffusion_planner/diffusion_planner/model/diffusion_planner.py @@ -2,13 +2,24 @@ from diffusion_planner.model.module.decoder import Decoder from diffusion_planner.model.module.encoder import Encoder +from diffusion_planner.model.module.plantf_decoder import PlanTFDecoder + + +def build_decoder(config): + # getattr: args.json from runs older than decoder_type defaults to diffusion + decoder_type = getattr(config, "decoder_type", "diffusion") + if decoder_type == "diffusion": + return Decoder(config) + if decoder_type == "plantf": + return PlanTFDecoder(config) + raise ValueError(f"Unknown decoder type: {decoder_type}") class Diffusion_Planner(nn.Module): def __init__(self, config): super().__init__() self.encoder = Encoder(config) - self.decoder = Decoder(config) + self.decoder = build_decoder(config) @property def sde(self): diff --git a/diffusion_planner/diffusion_planner/model/module/encoder.py b/diffusion_planner/diffusion_planner/model/module/encoder.py index 893f44519..07a9c7263 100644 --- a/diffusion_planner/diffusion_planner/model/module/encoder.py +++ b/diffusion_planner/diffusion_planner/model/module/encoder.py @@ -48,6 +48,27 @@ def __init__(self, config): self.ego_history_dropout_rate = config.ego_history_dropout_rate self.use_turn_indicators = config.use_turn_indicators + # A2: current-ego-state token (ego_current_state[:, 4:10] = + # vx, vy, ax, ay, steer, yaw_rate). C1: consecutive-frame xy deltas. + # See docs/plantf_original_comparison_and_roadmap.md. + self.plantf_ego_state_token = getattr(config, "plantf_ego_state_token", False) + self.plantf_ego_state_dropout = getattr(config, "plantf_ego_state_dropout", 0.75) + self.plantf_input_delta = getattr(config, "plantf_input_delta", False) + # Zero out the goal_pose input entirely. The mini dataset's goal_pose is in + # global/map frame (a data bug) while Autoware feeds it ego-frame, so it is + # a train/deploy-mismatched input. Masking it makes the model plan from the + # route (route_lanes still encodes the goal) instead. Bakes into the ONNX, + # so Autoware needs no change. See docs/plantf_experiment_log.md. + self.plantf_mask_goal_pose = getattr(config, "plantf_mask_goal_pose", False) + self._ego_state_slice = slice(4, 10) + if self.plantf_ego_state_token: + ego_state_dim = self._ego_state_slice.stop - self._ego_state_slice.start + self.ego_state_emb = nn.Sequential( + nn.Linear(ego_state_dim, config.hidden_dim), + nn.ReLU(inplace=True), + nn.Linear(config.hidden_dim, config.hidden_dim), + ) + ego_num = 1 goal_pose_num = 1 ego_shape_num = 1 @@ -102,7 +123,7 @@ def __init__(self, config): drop_path_rate=config.encoder_drop_path_rate, hidden_dim=config.hidden_dim, depth=config.encoder_mixer_depth, - num_types=POLYGON_TYPE_NUM, + point_dim=2 + POLYGON_TYPE_NUM, ) self.line_string_encoder = LineEncoder( config.line_string_len, @@ -110,7 +131,7 @@ def __init__(self, config): drop_path_rate=config.encoder_drop_path_rate, hidden_dim=config.hidden_dim, depth=config.encoder_mixer_depth, - num_types=LINE_STRING_TYPE_NUM, + point_dim=2 + LINE_STRING_TYPE_NUM, ) self.goal_pose_encoder = GoalPoseEncoder( drop_path_rate=config.encoder_drop_path_rate, @@ -164,11 +185,24 @@ def _basic_init(m): nn.init.normal_(self.lane_encoder.speed_limit_emb.weight, std=0.02) nn.init.normal_(self.lane_encoder.attribute_emb.weight, std=0.02) + def _to_xy_delta(self, x: torch.Tensor, time_dim: int) -> torch.Tensor: + """C1: replace absolute xy (channels 0,1) with consecutive-frame deltas + along ``time_dim`` (first timestep -> zero delta); other channels + unchanged. Padding rows are all-zero so their deltas stay zero.""" + xy = x[..., :2] + delta = xy - torch.roll(xy, shifts=1, dims=time_dim) + idx = [slice(None)] * x.dim() + idx[time_dim] = slice(0, 1) + delta[tuple(idx)] = 0.0 + return torch.cat([delta, x[..., 2:]], dim=-1) + def forward(self, inputs): # ego agent ego = inputs["ego_agent_past"].clone() # (B, T=INPUT_T + 1, D=4) if not self.use_ego_history: ego = torch.zeros_like(ego) + if self.plantf_input_delta: + ego = self._to_xy_delta(ego, time_dim=1) # C1 ego = torch.cat( # [torch.zeros_like(ego[:, :-6]), ego[:, -6:]], [ego[:, :6], torch.zeros_like(ego[:, 6:])], @@ -177,6 +211,8 @@ def forward(self, inputs): # agents neighbors = inputs["neighbor_agents_past"].clone() # (B, N=32, T=21, D=11) + if self.plantf_input_delta: + neighbors = self._to_xy_delta(neighbors, time_dim=2) # C1 neighbors = torch.cat( [torch.zeros_like(neighbors[:, :, :-6]), neighbors[:, :, -6:]], dim=2, @@ -203,6 +239,8 @@ def forward(self, inputs): # goal pose goal_pose = inputs["goal_pose"] # (B, D=4) + if self.plantf_mask_goal_pose: + goal_pose = torch.zeros_like(goal_pose) # ignore goal_pose; plan from route # ego shape ego_shape = inputs["ego_shape"] # (B, D=3) @@ -222,6 +260,19 @@ def forward(self, inputs): encoding_ego, p=self.ego_history_dropout_rate, training=self.training ) + if self.plantf_ego_state_token: + # A2: replace the ego token with an embedding of the current motion + # state (normalized; observation_normalizer runs before the encoder), + # anchoring the prediction to the current motion. Keep ego_mask/ego_pos + # from the history encoder (ego is always valid at the current pose). + ego_state = inputs["ego_current_state"][:, self._ego_state_slice] # (B, 6) + if self.training and self.plantf_ego_state_dropout > 0: + keep = (torch.rand_like(ego_state) >= self.plantf_ego_state_dropout).to( + ego_state.dtype + ) + ego_state = ego_state * keep + encoding_ego = self.ego_state_emb(ego_state).unsqueeze(1) # (B, 1, hidden) + encoding_neighbors, neighbors_mask, neighbor_pos = self.neighbor_encoder(neighbors) encoding_static, static_mask, static_pos = self.static_encoder(static) encoding_lanes, lanes_mask, lane_pos = self.lane_encoder( @@ -616,7 +667,7 @@ def forward(self, x, speed_limit, has_speed_limit): class LineEncoder(nn.Module): - def __init__(self, line_len, class_type, drop_path_rate, hidden_dim, depth, num_types): + def __init__(self, line_len, class_type, drop_path_rate, hidden_dim, depth, point_dim=2): super().__init__() self._class_type = class_type tokens_mlp_dim = 64 @@ -624,11 +675,13 @@ def __init__(self, line_len, class_type, drop_path_rate, hidden_dim, depth, num_ self._line_len = line_len - # type one-hot is fed via an embedding, not through the Mlp/Mixer - self.type_emb = nn.Linear(num_types, channels_mlp_dim) - + # DP checkpoint compatibility: keep the map-type one-hot channels in + # the MLP input. The production DP checkpoint was trained with + # ``point_dim + 2`` features (coordinates, type one-hot, dx, dy): + # polygon=5 and line_string=6. Moving type to a separate embedding + # changes these projection shapes and prevents encoder warm-start. self.channel_pre_project = Mlp( - in_features=4, # x, y, dx, dy + in_features=point_dim + 2, hidden_features=channels_mlp_dim, out_features=channels_mlp_dim, act_layer=nn.GELU, @@ -661,50 +714,43 @@ def forward(self, x): """ B, P, V, D = x.shape - # Mask before splitting: an element is empty only if every coord/type is 0. - mask_p = torch.sum(torch.ne(x, 0), dim=(-2, -1)) == 0 - valid_indices = ~mask_p.view(-1) - - # Split coordinates and type one-hot. The type is constant within an element, - # so a single representative point is enough. - coords = x[..., :2] # (B, P, V, 2) - type_one_hot = x[:, :, 0, 2:] # (B, P, num_types) - - diff_x = coords[:, :, 1:, 0] - coords[:, :, :-1, 0] # (B, P, V-1) - diff_y = coords[:, :, 1:, 1] - coords[:, :, :-1, 1] # (B, P, V-1) + diff_x = x[:, :, 1:, 0] - x[:, :, :-1, 0] # (B, P, V-1) + diff_y = x[:, :, 1:, 1] - x[:, :, :-1, 1] # (B, P, V-1) diff_x = torch.cat([diff_x, torch.zeros_like(diff_x[:, :, :1])], dim=2) # (B, P, V) diff_x = diff_x.view(B, P, V, 1) diff_y = torch.cat([diff_y, torch.zeros_like(diff_y[:, :, :1])], dim=2) # (B, P, V) diff_y = diff_y.view(B, P, V, 1) - feat = torch.concat([coords, diff_x, diff_y], dim=-1) # (B, P, V, 4): x, y, dx, dy + x = torch.concat([x, diff_x, diff_y], dim=-1) # (B, P, V, D+2) - pos = feat[:, :, int(self._line_len / 2), :4].clone() # x, y, x'-x, y'-y + pos = x[:, :, int(self._line_len / 2), :4].clone() # legacy DP layout heading = torch.atan2(pos[..., 3], pos[..., 2]) pos = torch.stack( [pos[..., 0], pos[..., 1], torch.cos(heading), torch.sin(heading)], dim=-1 ) pos = add_class_type(pos, self._class_type) - feat = feat.view(B * P, V, -1) - - feat = self.channel_pre_project(feat) - feat = feat.permute(0, 2, 1) - feat = self.token_pre_project(feat) - feat = feat.permute(0, 2, 1) - for block in self.blocks: - feat = block(feat) + B, P, V, _ = x.shape + mask_v = torch.sum(torch.ne(x[..., :4], 0), dim=-1).to(x.device) == 0 + mask_p = torch.sum(~mask_v, dim=-1) == 0 + valid_indices = ~mask_p.view(-1) - feat = torch.mean(feat, dim=1) + x = x.view(B * P, V, -1) + x = torch.where(valid_indices.view(-1, 1, 1), x, torch.zeros_like(x)) - # Inject type information via embedding instead of through the Mlp/Mixer. - feat = feat + self.type_emb(type_one_hot.view(B * P, -1)) + x = self.channel_pre_project(x) + x = x.permute(0, 2, 1) + x = self.token_pre_project(x) + x = x.permute(0, 2, 1) + for block in self.blocks: + x = block(x) - feat = self.emb_project(self.norm(feat)) + x = torch.mean(x, dim=1) + x = self.emb_project(self.norm(x)) # Apply mask to zero out invalid positions - feat = feat * valid_indices.float().unsqueeze(-1) + x = x * valid_indices.float().unsqueeze(-1) - return feat.view(B, P, -1), mask_p.reshape(B, -1), pos.view(B, P, -1) + return x.view(B, P, -1), mask_p.reshape(B, -1), pos.view(B, P, -1) class GoalPoseEncoder(nn.Module): diff --git a/diffusion_planner/diffusion_planner/model/module/plantf_decoder.py b/diffusion_planner/diffusion_planner/model/module/plantf_decoder.py new file mode 100644 index 000000000..6ade8b0e1 --- /dev/null +++ b/diffusion_planner/diffusion_planner/model/module/plantf_decoder.py @@ -0,0 +1,1008 @@ +"""PlanTF-style regression decoder head on top of the Diffusion-Planner encoder. + +Ported from planTF (https://github.com/jchengai/planTF, Cheng et al., ICRA 2024) +and adapted to the Diffusion-Planner interface: + +- consumes the shared encoder output ``[B, token_num, hidden_dim]`` and the raw + ``inputs`` dict, like :class:`diffusion_planner.model.module.decoder.Decoder` +- outputs (x, y, cos, sin) trajectories in the ``StateNormalizer`` space during + training and denormalized ``prediction`` ``[B, P, T, 4]`` at inference, so + validation / simulation / ROS consumers work unchanged +- diffusion-only inputs (``sampled_trajectories``, ``diffusion_time``, + ``delay``) are accepted but ignored + +``compute_plantf_training_loss`` is the planTF counterpart of +``decoder.compute_training_loss`` (same convention as ``grpo_utils`` mirroring +it): winner-takes-all regression over ``num_modes`` ego candidates + mode +classification, with DP's lat/lon/heading decomposition, velocity/timestep +weighting and penalty losses reused as-is. +""" + +import math +from argparse import Namespace + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from diffusion_planner.dimensions import TURN_INDICATOR_OUTPUT_DIM +from diffusion_planner.loss import ( + compute_ego_edge_points, + compute_neighbor_collision_penalty, + compute_road_border_penalty, + make_turn_indicator_gt, + velocity_to_waypoints, +) +from diffusion_planner.model.module.turn_indicator import TurnIndicatorNetwork +from diffusion_planner.utils.normalizer import StateNormalizer + + +def bezier_basis(num_control_points: int, num_steps: int) -> torch.Tensor: + """Bernstein/Bezier basis matrix ``[num_control_points, num_steps]``. + + Column ``t`` holds the Bernstein weights ``B_{i,n}(t)`` (n = num_control_points-1) + at ``t = τ/(num_steps-1)`` for control point ``i``. A trajectory is the + basis-weighted control points: ``traj[t] = Σ_i ctrl_i · basis[i, t]``. + + Why: the default head is a flat ``Linear(hidden, T*C)`` — the ``T`` output + steps are independent linear read-outs with no temporal inductive bias, so + temporal coherence is only bolted on afterwards (velocity-rep cumsum + + smoothness penalty). A Bezier expansion makes time an explicit axis: the + partition-of-unity (``Σ_i basis[i,t]=1``) and convex-hull property bound the + curve by its control points, so it is *structurally* smooth (C∞) instead of + smooth-by-penalty. See docs/plantf_head_development_notes.md §9. + """ + n = num_control_points - 1 + t = torch.linspace(0.0, 1.0, num_steps) # [T] + i = torch.arange(num_control_points) # [n+1] + coeff = torch.tensor([math.comb(n, int(k)) for k in i], dtype=torch.float32) # [n+1] + # B_i(t) = C(n,i) t^i (1-t)^(n-i); torch defines 0**0 = 1 so the endpoints + # (t=0 for i=0, t=1 for i=n) evaluate to 1 as required. + t_pow = t[None, :] ** i[:, None].float() # [n+1, T] + tm_pow = (1.0 - t)[None, :] ** (n - i)[:, None].float() # [n+1, T] + return coeff[:, None] * t_pow * tm_pow # [n+1, T] + + +class PlanTFTrajectoryHead(nn.Module): + """Multi-modal ego trajectory head (planTF ``TrajectoryDecoder``). + + When ``num_control_points > 0`` the head regresses that many Bezier control + points per mode/channel instead of ``future_steps`` free waypoints, then + expands them through a fixed :func:`bezier_basis` matrix. This injects a + temporal inductive bias (structural smoothness) the flat head lacks while + keeping the exact same ``[B, K, future_steps, out_channels]`` output contract + (velocity integration, WTA, zero-init, ONNX unchanged). The basis matrix is a + constant buffer, so ONNX export is just an extra matmul — no recurrence. + """ + + def __init__( + self, + embed_dim, + num_modes, + future_steps, + out_channels=4, + ego_state_dim=0, + predict_scale=False, + num_control_points=0, + ): + super().__init__() + + self.embed_dim = embed_dim + self.num_modes = num_modes + self.future_steps = future_steps + self.out_channels = out_channels + self.ego_state_dim = ego_state_dim + self.predict_scale = predict_scale + # Temporal basis: 0 disables (flat per-step head); >0 regresses that many + # Bezier control points and expands them to future_steps via the buffer. + self.num_control_points = num_control_points + if num_control_points > 0: + self.register_buffer( + "basis", bezier_basis(num_control_points, future_steps), persistent=False + ) + self._out_steps = num_control_points + else: + self._out_steps = future_steps + + # Inject the current ego motion state (vx, vy, ax, ay, steering, yaw_rate) + # into the ego token before branching into modes. Unlike the diffusion + # decoder, the planTF head never sees the current state otherwise, so its + # absolute-waypoint regression is not anchored to "where/how fast am I + # now" — the cause of the start-point scatter / stop-jitter / divergence + # that the diffusion head (which pins the current state) does not show. + # See docs/plantf_dead_mode_improvement.md. + if ego_state_dim > 0: + self.ego_state_proj = nn.Sequential( + nn.Linear(ego_state_dim, embed_dim), + nn.ReLU(inplace=True), + nn.Linear(embed_dim, embed_dim), + ) + + self.multimodal_proj = nn.Linear(embed_dim, num_modes * embed_dim) + + hidden = 2 * embed_dim + self.loc = nn.Sequential( + nn.Linear(embed_dim, hidden), + nn.LayerNorm(hidden), + nn.ReLU(inplace=True), + nn.Linear(hidden, self._out_steps * out_channels), + ) + self.pi = nn.Sequential( + nn.Linear(embed_dim, hidden), + nn.LayerNorm(hidden), + nn.ReLU(inplace=True), + nn.Linear(hidden, 1), + ) + # Optional per-point log-scale head for the Laplace NLL loss (planTF's + # probabilistic regression). Only used when predict_scale is True. + if predict_scale: + self.scale = nn.Sequential( + nn.Linear(embed_dim, hidden), + nn.LayerNorm(hidden), + nn.ReLU(inplace=True), + nn.Linear(hidden, self._out_steps * out_channels), + ) + + def _expand(self, flat): + """[B, K, _out_steps*C] -> [B, K, future_steps, C], applying the Bezier + basis when enabled. Zero-init of the final Linear makes control points + zero -> trajectory zero, identical to the flat head's zero-init prior.""" + if self.num_control_points > 0: + ctrl = flat.view(-1, self.num_modes, self.num_control_points, self.out_channels) + # traj[t] = Σ_c ctrl[c] · basis[c, t] + return torch.einsum("bkcd,ct->bktd", ctrl, self.basis) + return flat.view(-1, self.num_modes, self.future_steps, self.out_channels) + + def forward(self, x, ego_state=None): + """ + Args: + x: [B, embed_dim] ego token from the encoder. + ego_state: [B, ego_state_dim] normalized current ego motion state, + or None when ego_state_dim == 0. + + Returns: + loc: [B, num_modes, future_steps, out_channels] mode trajectories. + pi: [B, num_modes] mode logits. + (predict_scale only) log_scale: [B, num_modes, future_steps, out_channels] + """ + if self.ego_state_dim > 0 and ego_state is not None: + x = x + self.ego_state_proj(ego_state) + x = self.multimodal_proj(x).view(-1, self.num_modes, self.embed_dim) + loc = self._expand(self.loc(x)) + pi = self.pi(x).squeeze(-1) + if self.predict_scale: + log_scale = self._expand(self.scale(x)) + return loc, pi, log_scale + return loc, pi + + +class PlanTFCrossAttnHead(nn.Module): + """K learnable mode queries that cross-attend to the full encoder memory. + + Unlike :class:`PlanTFTrajectoryHead` (which reshapes the single pooled + + dropout ego token into K modes), each mode query attends over ALL encoder + tokens (ego / neighbors / map / route), so the 80-point regression gets the + scene context that the single-token bottleneck loses — a candidate fix for + the tail divergence. See docs/plantf_head_development_notes.md §9. Same + (loc, pi) output contract as PlanTFTrajectoryHead, so the rest of the decoder + (velocity integration, WTA, zero-init, ONNX) is unchanged. + """ + + def __init__( + self, + embed_dim, + num_modes, + future_steps, + out_channels=4, + ego_state_dim=0, + num_heads=8, + depth=2, + predict_scale=False, + ): + super().__init__() + self.embed_dim = embed_dim + self.num_modes = num_modes + self.future_steps = future_steps + self.out_channels = out_channels + self.ego_state_dim = ego_state_dim + self.predict_scale = predict_scale + + if ego_state_dim > 0: + self.ego_state_proj = nn.Sequential( + nn.Linear(ego_state_dim, embed_dim), + nn.ReLU(inplace=True), + nn.Linear(embed_dim, embed_dim), + ) + + self.mode_queries = nn.Parameter(torch.randn(num_modes, embed_dim) * 0.02) + self.attn_layers = nn.ModuleList( + [nn.MultiheadAttention(embed_dim, num_heads, batch_first=True) for _ in range(depth)] + ) + self.attn_norms = nn.ModuleList([nn.LayerNorm(embed_dim) for _ in range(depth)]) + self.ffns = nn.ModuleList( + [ + nn.Sequential( + nn.Linear(embed_dim, 2 * embed_dim), + nn.ReLU(inplace=True), + nn.Linear(2 * embed_dim, embed_dim), + ) + for _ in range(depth) + ] + ) + self.ffn_norms = nn.ModuleList([nn.LayerNorm(embed_dim) for _ in range(depth)]) + + hidden = 2 * embed_dim + self.loc = nn.Sequential( + nn.Linear(embed_dim, hidden), + nn.LayerNorm(hidden), + nn.ReLU(inplace=True), + nn.Linear(hidden, future_steps * out_channels), + ) + self.pi = nn.Sequential( + nn.Linear(embed_dim, hidden), + nn.LayerNorm(hidden), + nn.ReLU(inplace=True), + nn.Linear(hidden, 1), + ) + if predict_scale: + self.scale = nn.Sequential( + nn.Linear(embed_dim, hidden), + nn.LayerNorm(hidden), + nn.ReLU(inplace=True), + nn.Linear(hidden, future_steps * out_channels), + ) + + def forward(self, memory, memory_valid, ego_state=None): + """ + Args: + memory: [B, N, embed_dim] all encoder tokens. + memory_valid: [B, N] bool, True where the token is valid. + ego_state: [B, ego_state_dim] normalized current ego motion state. + + Returns: + loc: [B, num_modes, future_steps, out_channels] + pi: [B, num_modes] + """ + B = memory.shape[0] + q = self.mode_queries.unsqueeze(0).expand(B, -1, -1) # [B, K, D] + if self.ego_state_dim > 0 and ego_state is not None: + q = q + self.ego_state_proj(ego_state).unsqueeze(1) + key_padding_mask = ~memory_valid # True => ignore this key + for attn, an, ffn, fn in zip(self.attn_layers, self.attn_norms, self.ffns, self.ffn_norms): + a, _ = attn(q, memory, memory, key_padding_mask=key_padding_mask, need_weights=False) + q = an(q + a) + q = fn(q + ffn(q)) + loc = self.loc(q).view(B, self.num_modes, self.future_steps, self.out_channels) + pi = self.pi(q).squeeze(-1) + if self.predict_scale: + log_scale = self.scale(q).view(B, self.num_modes, self.future_steps, self.out_channels) + return loc, pi, log_scale + return loc, pi + + +class PlanTFGRUHead(nn.Module): + """Recurrent trajectory head: a GRU unrolls the ``future_steps`` waypoints. + + Same mlp-style mode formation as :class:`PlanTFTrajectoryHead` (reshape the + single ego token into K modes), but instead of a flat ``Linear(hidden, T*C)`` + the per-step waypoints come from unrolling a GRU. The recurrence couples + adjacent steps in the *architecture* (each step's hidden state carries the + previous), giving the output an explicit temporal inductive bias the flat + head lacks. This is NON-autoregressive — the GRU is fed a per-step input + (learned temporal embedding + the mode context) rather than its own previous + xy — so it exports as a single ONNX GRU op (no output-feedback loop). Same + (loc, pi[, log_scale]) contract, so velocity integration / WTA / zero-init + are unchanged. Note: RNN heads are more deploy-fragile than mlp/basis; this + is an experimental toggle, not the deploy default. + """ + + def __init__( + self, + embed_dim, + num_modes, + future_steps, + out_channels=4, + ego_state_dim=0, + predict_scale=False, + gru_hidden=None, + ): + super().__init__() + self.embed_dim = embed_dim + self.num_modes = num_modes + self.future_steps = future_steps + self.out_channels = out_channels + self.ego_state_dim = ego_state_dim + self.predict_scale = predict_scale + self.gru_hidden = gru_hidden or embed_dim + + if ego_state_dim > 0: + self.ego_state_proj = nn.Sequential( + nn.Linear(ego_state_dim, embed_dim), + nn.ReLU(inplace=True), + nn.Linear(embed_dim, embed_dim), + ) + + self.multimodal_proj = nn.Linear(embed_dim, num_modes * embed_dim) + # Per-step temporal embedding: breaks the symmetry of feeding the same + # context every step and gives the GRU an explicit time signal. + self.step_emb = nn.Parameter(torch.randn(future_steps, embed_dim) * 0.02) + self.h0_proj = nn.Linear(embed_dim, self.gru_hidden) + self.gru = nn.GRU(embed_dim, self.gru_hidden, batch_first=True) + + hidden = 2 * embed_dim + # loc/scale kept as Sequential ending in Linear so the shared zero-init + # (PlanTFDecoder indexes head[-1]) works uniformly across head types. + self.loc = nn.Sequential(nn.Linear(self.gru_hidden, out_channels)) + self.pi = nn.Sequential( + nn.Linear(embed_dim, hidden), + nn.LayerNorm(hidden), + nn.ReLU(inplace=True), + nn.Linear(hidden, 1), + ) + if predict_scale: + self.scale = nn.Sequential(nn.Linear(self.gru_hidden, out_channels)) + + def _unroll(self, modes): + """[B, K, D] mode embeddings -> [B, K, T, gru_hidden] GRU outputs.""" + B = modes.shape[0] + m = modes.reshape(B * self.num_modes, self.embed_dim) # [BK, D] + h0 = torch.tanh(self.h0_proj(m)).unsqueeze(0) # [1, BK, H] + inp = self.step_emb.unsqueeze(0) + m.unsqueeze(1) # [BK, T, D] + out, _ = self.gru(inp, h0) # [BK, T, H] + return out.view(B, self.num_modes, self.future_steps, self.gru_hidden) + + def forward(self, x, ego_state=None): + """ + Args: + x: [B, embed_dim] ego token from the encoder. + ego_state: [B, ego_state_dim] normalized current ego motion state. + + Returns: + loc: [B, num_modes, future_steps, out_channels] + pi: [B, num_modes] + (predict_scale only) log_scale: [B, num_modes, future_steps, out_channels] + """ + if self.ego_state_dim > 0 and ego_state is not None: + x = x + self.ego_state_proj(ego_state) + modes = self.multimodal_proj(x).view(-1, self.num_modes, self.embed_dim) + out = self._unroll(modes) + loc = self.loc(out) # [B, K, T, C] + pi = self.pi(modes).squeeze(-1) + if self.predict_scale: + log_scale = self.scale(out) + return loc, pi, log_scale + return loc, pi + + +class PlanTFDecoder(nn.Module): + """One-shot regression decoder with the same forward contract as ``Decoder``.""" + + def __init__(self, config): + super().__init__() + + self._predicted_neighbor_num = config.predicted_neighbor_num + self._future_len = config.future_len + self._num_modes = getattr(config, "num_modes", 6) + # When True the heads regress per-step displacement instead of absolute + # waypoints; _decode integrates it (cumsum) into absolute waypoints, which + # gives the temporal continuity a per-timestep absolute regression lacks + # (fixes the "comb" jitter and stalled forward progress of the planTF head). + self._use_velocity = getattr(config, "use_velocity_representation", False) + # Original planTF trains xy targets relative to each agent's current + # position. This is particularly important for neighbors: unlike ego, + # their current xy is not always the origin of the ego frame. The head + # predicts the relative component in StateNormalizer units and we add + # the current state (also in StateNormalizer units) before exposing the + # usual absolute ego-frame prediction. It is opt-in because a legacy + # absolute-xy checkpoint cannot be reinterpreted safely. + self._relative_xy = getattr(config, "plantf_relative_xy", False) + # Feed the current ego motion state (ego_current_state[:, 4:10] = + # vx, vy, ax, ay, steering, yaw_rate) into the trajectory head so the + # absolute-waypoint regression is anchored to the current motion, the + # way the diffusion decoder is via its pinned current state. + self._use_ego_state = getattr(config, "plantf_use_ego_state_in_head", True) + self._ego_state_slice = slice(4, 10) + ego_state_dim = self._ego_state_slice.stop - self._ego_state_slice.start + + # Optional toggles (default = original behavior) for combined ablation. + # head_type "cross_attn": mode queries cross-attend to all encoder tokens + # instead of reshaping the single ego token (docs §9). route_rerank: + # at inference, pick the mode by route adherence among the top-k pi + # modes instead of argmax(pi) (only in `forward`, not the ONNX deploy + # graph, which lacks route_lanes). + self._head_type = getattr(config, "plantf_head_type", "mlp") + # head_type "basis": the mlp head, but the trajectory is a Bezier curve of + # `plantf_basis_control_points` control points expanded over time (a + # temporal inductive bias the flat per-step head lacks). Uses the same + # single-ego-token mode formation as "mlp" (keeps deploy robustness). + self._basis_control_points = ( + getattr(config, "plantf_basis_control_points", 8) if self._head_type == "basis" else 0 + ) + self._route_rerank = getattr(config, "plantf_route_rerank", False) + self._route_rerank_topk = getattr(config, "plantf_route_rerank_topk", 3) + self._observation_normalizer = getattr(config, "observation_normalizer", None) + # Laplace NLL loss: the head additionally regresses a per-point log-scale. + self._predict_scale = getattr(config, "plantf_use_laplace_nll", False) + + hidden_dim = config.hidden_dim + head_ego_state_dim = ego_state_dim if self._use_ego_state else 0 + if self._head_type == "cross_attn": + self.trajectory_head = PlanTFCrossAttnHead( + embed_dim=hidden_dim, + num_modes=self._num_modes, + future_steps=self._future_len, + out_channels=4, # x, y, cos, sin + ego_state_dim=head_ego_state_dim, + num_heads=config.num_heads, + predict_scale=self._predict_scale, + ) + elif self._head_type == "gru": + # Recurrent head: GRU unrolls the waypoints (temporal recurrence). + # Experimental — more deploy-fragile than mlp/basis (RNN ONNX op). + self.trajectory_head = PlanTFGRUHead( + embed_dim=hidden_dim, + num_modes=self._num_modes, + future_steps=self._future_len, + out_channels=4, # x, y, cos, sin + ego_state_dim=head_ego_state_dim, + predict_scale=self._predict_scale, + ) + else: + self.trajectory_head = PlanTFTrajectoryHead( + embed_dim=hidden_dim, + num_modes=self._num_modes, + future_steps=self._future_len, + out_channels=4, # x, y, cos, sin + ego_state_dim=head_ego_state_dim, + predict_scale=self._predict_scale, + num_control_points=self._basis_control_points, + ) + # planTF's agent_predictor, widened from (x, y) to DP's (x, y, cos, sin) + self.neighbor_predictor = nn.Sequential( + nn.Linear(hidden_dim, 2 * hidden_dim), + nn.LayerNorm(2 * hidden_dim), + nn.ReLU(inplace=True), + nn.Linear(2 * hidden_dim, self._future_len * 4), + ) + self.turn_indicator_predictor = nn.Linear( + 2 * (self._future_len // 10) + hidden_dim, TURN_INDICATOR_OUTPUT_DIM + ) + self._state_normalizer: StateNormalizer = config.state_normalizer + + def _basic_init(m): + if isinstance(m, nn.Linear): + torch.nn.init.xavier_uniform_(m.weight) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + elif isinstance(m, nn.Embedding): + nn.init.normal_(m.weight, mean=0.0, std=0.02) + + self.apply(_basic_init) + + # Zero-out the output layers (same convention as Decoder.dit.final_layer). + # 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 makes every mode + # start at the normalized-space mean — a smooth prior trajectory — so an + # undertrained mode degrades gracefully instead of into noise. + zero_init_heads = [ + self.trajectory_head.loc, + self.trajectory_head.pi, + self.neighbor_predictor, + ] + # Zero-init the log-scale head too -> scale starts at exp(0)=1, a neutral + # Laplace width, so early training is not dominated by scale noise. + if self._predict_scale: + zero_init_heads.append(self.trajectory_head.scale) + for head in zero_init_heads: + nn.init.constant_(head[-1].weight, 0) + nn.init.constant_(head[-1].bias, 0) + + self.independent_turn_indicator_predictor = TurnIndicatorNetwork( + hidden_dim=config.hidden_dim // 2, + num_heads=config.num_heads // 2, + mixer_depth=config.encoder_mixer_depth // 2, + fusion_depth=config.encoder_fusion_depth // 2, + drop_path_rate=config.encoder_drop_path_rate, + ) + + def _compute_turn_indicator(self, ego_trajectory, encoding_pooled): + turn_indicator_input = torch.cat([ego_trajectory, encoding_pooled], dim=-1) + return self.turn_indicator_predictor(turn_indicator_input) + + def _integrate_velocity(self, velocity: torch.Tensor, ego: bool) -> torch.Tensor: + """Integrate per-step displacement into normalized absolute waypoints. + + ``velocity_to_waypoints`` cumsums the xy displacement (ego-centric metres, + current pose = origin) and passes the heading channel through. The result + is then mapped into the StateNormalizer space so every downstream consumer + (WTA selection, smooth-L1 loss, mode metrics, prediction, ONNX) is unchanged. + """ + wp = velocity_to_waypoints(velocity) # [..., T, 4] relative xy, heading passthrough + idx = 0 if ego else 1 # ego uses row 0; neighbors share the same neighbor stats + mean = self._state_normalizer.mean[idx].to(wp.device) # [1, 4] + std = self._state_normalizer.std[idx].to(wp.device) # [1, 4] + if self._relative_xy: + # xy is a displacement, not an absolute waypoint. Do not subtract + # the StateNormalizer mean here: the current absolute state is + # added later in _anchor_relative_xy. + xy = wp[..., :2] / std[..., :2] + heading = (wp[..., 2:] - mean[..., 2:]) / std[..., 2:] + return torch.cat([xy, heading], dim=-1) + return (wp - mean) / std + + def _anchor_relative_xy(self, trajectory, neighbor_prediction, inputs): + """Convert planTF-style relative xy head outputs to absolute xy. + + ``trajectory`` and ``neighbor_prediction`` are already in + StateNormalizer units. The observation normalizer is inverted first + because the current neighbor state in ``inputs`` is observation- + normalized, whereas StateNormalizer is the coordinate system of the + decoder targets. Heading channels intentionally remain unchanged; + original planTF only predicts neighbor xy. + """ + if not self._relative_xy: + return trajectory, neighbor_prediction + if inputs is None: + raise ValueError( + "plantf_relative_xy requires raw inputs to anchor ego and neighbor current positions" + ) + + if self._observation_normalizer is None: + ego_current = inputs["ego_current_state"][:, :4] + neighbor_current = inputs["neighbor_agents_past"][:, : self._predicted_neighbor_num, -1, :4] + else: + # The full DP input carries a 10-channel ego_current_state, while + # the split ONNX decoder reconstructs only its 4-channel pose from + # sampled_trajectories. Inverting the complete observation dict + # therefore fails for the split graph. Invert only the pose + # channels needed for the anchor, and preserve zero-padded poses. + ego_current = inputs["ego_current_state"][:, :4] + ego_norm = self._observation_normalizer._normalization_dict[ + "ego_current_state" + ] + ego_mask = torch.sum(torch.ne(ego_current, 0), dim=-1) == 0 + ego_current = ( + ego_current * ego_norm["std"][:4].to(ego_current.device) + + ego_norm["mean"][:4].to(ego_current.device) + ) + # Keep padded ego poses zero without boolean-index assignment. + # The latter is exported as a shape-sensitive Scatter/Where pattern + # that TensorRT cannot compile reliably. + ego_current = torch.where( + ego_mask.unsqueeze(-1), torch.zeros_like(ego_current), ego_current + ) + neighbor_current = self._observation_normalizer.inverse( + {"neighbor_agents_past": inputs["neighbor_agents_past"]} + )["neighbor_agents_past"][:, : self._predicted_neighbor_num, -1, :4] + B = trajectory.shape[0] + Pn = self._predicted_neighbor_num + neighbor_current = neighbor_current[:, :Pn] + current = torch.cat([ego_current[:, None], neighbor_current], dim=1) + mean = self._state_normalizer.mean[:, 0].to(current.device) + std = self._state_normalizer.std[:, 0].to(current.device) + current_norm_xy = (current[..., :2] - mean[None, :, :2]) / std[None, :, :2] + + trajectory = trajectory.clone() + neighbor_prediction = neighbor_prediction.clone() + trajectory[..., :2] += current_norm_xy[:, :1, None, :] + neighbor_prediction[..., :2] += current_norm_xy[:, 1:, None, :] + return trajectory, neighbor_prediction + + def _ego_state_feat(self, inputs): + """Normalized current ego motion state fed to the trajectory head, or + None when the feature is disabled. inputs are already observation- + normalized when this runs (train_epoch / node apply the normalizer).""" + if not self._use_ego_state: + return None + return inputs["ego_current_state"][:, self._ego_state_slice] + + def _decode(self, encoding, ego_state=None, inputs=None): + """Run both heads on the encoder tokens. + + Returns (trajectory [B, K, T, 4], probability [B, K], + neighbor_prediction [B, Pn, T, 4]), all in normalized space. + """ + B = encoding.shape[0] + Pn = self._predicted_neighbor_num + if self._head_type == "cross_attn": + memory_valid = torch.any(encoding != 0, dim=-1) # [B, N] + head_out = self.trajectory_head(encoding, memory_valid, ego_state) + else: + head_out = self.trajectory_head(encoding[:, 0], ego_state) + if self._predict_scale: + trajectory, probability, log_scale = head_out + else: + trajectory, probability = head_out + log_scale = None + neighbor_prediction = self.neighbor_predictor(encoding[:, 1 : 1 + Pn]).view( + B, Pn, self._future_len, 4 + ) + if self._use_velocity: + trajectory = self._integrate_velocity(trajectory, ego=True) + neighbor_prediction = self._integrate_velocity(neighbor_prediction, ego=False) + trajectory, neighbor_prediction = self._anchor_relative_xy( + trajectory, neighbor_prediction, inputs + ) + if self._predict_scale: + return trajectory, probability, neighbor_prediction, log_scale + return trajectory, probability, neighbor_prediction + + def _best_mode_trajectory(self, trajectory, probability, inputs=None): + """Select one ego mode. By default argmax(pi); with route_rerank enabled + (and route_lanes available), pick the mode that best follows the route + among the top-k pi modes. gather keeps the batch axis dynamic for ONNX.""" + if ( + self._route_rerank + and inputs is not None + and "route_lanes" in inputs + and self._observation_normalizer is not None + ): + index = self._route_gated_index(trajectory, probability, inputs) + else: + index = probability.argmax(dim=-1) + index = index.view(-1, 1, 1, 1).expand(-1, 1, self._future_len, 4) + return trajectory.gather(1, index).squeeze(1) + + def _route_gated_index(self, trajectory, probability, inputs): + """Pick, among the top-k pi modes, the one whose waypoints stay closest + to the ego-frame route centreline (mean over time of min distance to any + route point). Both trajectory and route are de-normalized to metres. + Returns [B] mode indices. See docs/plantf_head_development_notes.md §9.""" + B, K, _, _ = trajectory.shape + dev = trajectory.device + std0 = self._state_normalizer.std[0].to(dev) # [1, 4] + mean0 = self._state_normalizer.mean[0].to(dev) + xy = (trajectory * std0 + mean0)[..., :2] # [B, K, T, 2] metres + + route = self._observation_normalizer.inverse({"route_lanes": inputs["route_lanes"]})[ + "route_lanes" + ].to(dev) # [B, S, P, C] metres (zeroed where invalid) + C = route.shape[-1] + route_flat = route.reshape(B, -1, C) # [B, M, C] + route_xy = route_flat[..., :2] # [B, M, 2] + valid = torch.any(route_flat != 0, dim=-1) # [B, M] + + d = torch.cdist(xy.reshape(B, K * xy.shape[2], 2), route_xy) # [B, K*T, M] + d = d.reshape(B, K, xy.shape[2], -1) + d = d + torch.where(valid[:, None, None, :], 0.0, 1e6) # ignore invalid points + route_cost = d.min(dim=-1).values.mean(dim=-1) # [B, K] + + k = min(self._route_rerank_topk, K) + topk = probability.topk(k, dim=-1).indices # [B, k] + sel = route_cost.gather(1, topk).argmin(dim=-1) # [B] + return topk.gather(1, sel[:, None]).squeeze(1) # [B] + + @staticmethod + def _pool_encoding(encoding): + # Pool only valid encoder tokens. The encoder zero-fills masked tokens. + encoding_valid = torch.any(encoding != 0, dim=-1) # [B, N] + encoding_count = encoding_valid.sum(dim=1).clamp_min(1).unsqueeze(-1) + return (encoding * encoding_valid.unsqueeze(-1)).sum(dim=1) / encoding_count + + def _subsampled_ego_xy(self, ego_trajectory): + """Every-10th-step xy, matching gt_trajectories[:, 0, 1::10, :2] used in + training (index 1 on the current-state-prepended axis = future step 0).""" + return ego_trajectory[:, ::10, :2].reshape(-1, 2 * (self._future_len // 10)) + + def forward_deploy( + self, + encoding, + sampled_trajectories=None, + diffusion_time=None, + neighbor_agents_past=None, + ): + """One-shot deployment path for the split ONNX export. + + Unlike the diffusion decoder there is no external denoising loop, so a + single call maps the encoder output to the final prediction. The + independent turn-indicator head is excluded (it re-encodes raw map + inputs itself and is a training-time auxiliary, not a deploy output). + + The split decoder deliberately uses the same four inputs as the DP + decoder graph. ``sampled_trajectories[:, 0, 0]`` supplies the ego + current state for relative-xy reconstruction; ``neighbor_agents_past`` + supplies current neighbor states. The diffusion-time input is kept in + the signature for the shared DP deployment contract, but is otherwise + irrelevant to this one-shot head. + + Returns: + prediction: [B, 1 + Pn, T, 4] best-mode ego + neighbors, denormalized. + probability: [B, K] mode logits. + turn_indicator_logit: [B, TURN_INDICATOR_OUTPUT_DIM] + """ + if self._relative_xy: + if sampled_trajectories is None or neighbor_agents_past is None: + raise ValueError( + "plantf_relative_xy split decoder requires sampled_trajectories and " + "neighbor_agents_past to reconstruct current agent positions" + ) + B = encoding.shape[0] + sampled = sampled_trajectories.reshape( + B, 1 + self._predicted_neighbor_num, 1 + self._future_len, 4 + ) + deploy_inputs = { + "ego_current_state": sampled[:, 0, 0], + "neighbor_agents_past": neighbor_agents_past, + } + else: + deploy_inputs = None + # The original planTF head does not consume ego_current_state directly + # on its split graph; state conditioning remains part of the encoder. + decode_out = self._decode(encoding, inputs=deploy_inputs) + # Deploy path ignores the log-scale (used only in the training NLL loss). + trajectory, probability, neighbor_prediction = decode_out[:3] + best_trajectory = self._best_mode_trajectory(trajectory, probability) + turn_indicator_logit = self._compute_turn_indicator( + self._subsampled_ego_xy(best_trajectory), self._pool_encoding(encoding) + ) + prediction = torch.cat([best_trajectory[:, None], neighbor_prediction], dim=1) + prediction = self._state_normalizer.inverse(prediction) + return prediction, probability, turn_indicator_logit + + def forward(self, encoding, inputs): + """ + Args: + encoding: [B, token_num, D] encoder output (token 0 = ego, + tokens 1 .. 1+predicted_neighbor_num = neighbors). + inputs: Dict. Training additionally requires "gt_trajectories" + [B, P, 1 + future_len, 4] (normalized, current state prepended), + injected by ``compute_plantf_training_loss``. + + Returns: + Dict: + [both] "trajectory": [B, K, T, 4] ego mode trajectories (normalized) + [both] "probability": [B, K] mode logits + [both] "neighbor_prediction": [B, Pn, T, 4] (normalized) + [both] "turn_indicator_logit" / "independent_turn_indicator_logit" + [inference-only] "prediction": [B, 1 + Pn, T, 4] best-mode ego + + neighbors, denormalized — same contract as ``Decoder``. + """ + B = encoding.shape[0] + Pn = self._predicted_neighbor_num + + decode_out = self._decode(encoding, self._ego_state_feat(inputs), inputs) + if self._predict_scale: + trajectory, probability, neighbor_prediction, log_scale = decode_out + else: + trajectory, probability, neighbor_prediction = decode_out + log_scale = None + encoding_pooled = self._pool_encoding(encoding) + + outputs = { + "trajectory": trajectory, + "probability": probability, + "neighbor_prediction": neighbor_prediction, + } + if log_scale is not None: + outputs["scale"] = log_scale # [B, K, T, 4] per-point log-scale for Laplace NLL + + if self.training: + gt_trajectories = inputs["gt_trajectories"].reshape(B, 1 + Pn, 1 + self._future_len, 4) + ego_trajectory = gt_trajectories[:, 0, 1::10, :2].reshape( + B, 2 * (self._future_len // 10) + ) + outputs["turn_indicator_logit"] = self._compute_turn_indicator( + ego_trajectory, encoding_pooled + ) + outputs["independent_turn_indicator_logit"] = self.independent_turn_indicator_predictor( + gt_trajectories[:, 0, 1:], inputs + ) + return outputs + + best_trajectory = self._best_mode_trajectory(trajectory, probability, inputs) + + outputs["turn_indicator_logit"] = self._compute_turn_indicator( + self._subsampled_ego_xy(best_trajectory), encoding_pooled + ) + outputs["independent_turn_indicator_logit"] = self.independent_turn_indicator_predictor( + best_trajectory, inputs + ) + + prediction = torch.cat([best_trajectory[:, None], neighbor_prediction], dim=1) + outputs["prediction"] = self._state_normalizer.inverse(prediction) + return outputs + + +def compute_plantf_training_loss( + model: nn.Module, + inputs: dict[str, torch.Tensor], + futures: tuple[torch.Tensor, torch.Tensor, torch.Tensor], + args: Namespace, +): + """PlanTF-head counterpart of ``decoder.compute_training_loss``. + + Returns the same loss keys plus ``mode_cls_loss`` so ``train_epoch`` can + combine them with the shared coefficients. + """ + norm = args.state_normalizer + + ego_future, neighbors_future, neighbor_future_mask = futures + neighbors_future_valid = ~neighbor_future_mask # [B, Pn, T] + + B, Pn, T, _ = neighbors_future.shape + ego_current, neighbors_current = ( + inputs["ego_current_state"][:, :4], + inputs["neighbor_agents_past"][:, :Pn, -1, :4], + ) + neighbor_current_mask = torch.sum(torch.ne(neighbors_current[..., :4], 0), dim=-1) == 0 + neighbor_mask = torch.concat( + (neighbor_current_mask.unsqueeze(-1), neighbor_future_mask), dim=-1 + ) + + gt_future = torch.cat( + [ego_future[:, None, :, :], neighbors_future[..., :]], dim=1 + ) # [B, P, T, 4] + current_states = torch.cat([ego_current[:, None], neighbors_current], dim=1) # [B, P, 4] + + all_gt = torch.cat([current_states[:, :, None, :], norm(gt_future)], dim=2) + all_gt[:, 1:][neighbor_mask] = 0.0 + + merged_inputs = {**inputs, "gt_trajectories": all_gt} + _, decoder_output = model(merged_inputs) + + trajectory = decoder_output["trajectory"] # [B, K, T, 4], normalized + probability = decoder_output["probability"] # [B, K] + neighbor_prediction = decoder_output["neighbor_prediction"] # [B, Pn, T, 4] + + gt_norm = all_gt[:, :, 1:, :] # [B, P, T, 4] + ego_gt = gt_norm[:, 0] # [B, T, 4] + + # Winner-takes-all mode selection on xy ADE in meters (the normalized x/y + # scales differ, so rescale by std; the mean shift cancels in the diff). + with torch.no_grad(): + ego_std_xy = norm.std[0][..., :2].to(trajectory.device) # [1, 2] + xy_diff = (trajectory[..., :2] - ego_gt[:, None, :, :2]) * ego_std_xy + ade = torch.norm(xy_diff, dim=-1).mean(-1) # [B, K] + best_mode = ade.argmin(dim=-1) # [B,] + best_trajectory = trajectory[torch.arange(B, device=trajectory.device), best_mode] + + prediction = torch.cat([best_trajectory[:, None], neighbor_prediction], dim=1) # [B, P, T, 4] + + # Original planTF ego/neighbor loss (jchengai/planTF): smooth L1 on the + # winner-takes-all best mode over all channels (x, y, cos, sin), uniform + # across timesteps, plus a cross-entropy on the DETACHED best mode. The DP + # diffusion-decoder tuning previously grafted on here — lat/lon/heading L2 + # decomposition, longitudinal velocity down-weighting, timestep weighting and + # the endpoint term — is intentionally dropped: it destabilized this one-shot + # regression head (团子 / 直進 / 反対方向 / 櫛状; see + # docs/plantf_dead_mode_improvement.md). turn-indicator and penalty losses + # below are kept for the DP/Autoware interface. + loss = {} + + ego_pred = prediction[:, 0, : args.ego_prediction_horizon] # [B, T', 4] + ego_tgt = gt_norm[:, 0, : args.ego_prediction_horizon] + # Ego regression loss (opt-in variants, docs/plantf_head_development_notes.md §9): + # - plantf_use_laplace_nll (#5): Laplace NLL |y-mu|*exp(-s)+s with the head's + # per-point log-scale s (planTF's probabilistic regression; calibrates the + # tail uncertainty and the mode confidence). + # - plantf_tail_weight (#6 variant): weight the per-timestep loss toward the + # tail (w_t = 1 + tail_weight * t/(T-1)). + # - both 0/off -> original uniform smooth-L1. + tail_weight = getattr(args, "plantf_tail_weight", 0.0) + tp = ego_pred.shape[1] + time_w = ( + 1.0 + tail_weight * torch.arange(tp, device=ego_pred.device) / max(tp - 1, 1) + if tail_weight > 0 + else None + ) + if getattr(args, "plantf_use_laplace_nll", False) and "scale" in decoder_output: + log_scale_best = decoder_output["scale"][ + torch.arange(B, device=trajectory.device), best_mode + ][:, : args.ego_prediction_horizon] # [B, T', 4] + # Clamp the log-scale: a near-perfect mode drives log_scale -> -inf and + # exp(-log_scale) -> inf, blowing up the NLL/gradient (standard Laplace + # NLL instability). [-6, 6] bounds exp(-s) to ~[0.0025, 400]. + log_scale_best = log_scale_best.clamp(-6.0, 6.0) + nll_t = ((ego_pred - ego_tgt).abs() * torch.exp(-log_scale_best) + log_scale_best).mean(-1) + loss["ego_planning_loss"] = ( + (nll_t * time_w).sum() / (time_w.sum() * B) if time_w is not None else nll_t.mean() + ) + elif time_w is not None: + per_t = F.smooth_l1_loss(ego_pred, ego_tgt, reduction="none").mean(dim=-1) # [B, T'] + loss["ego_planning_loss"] = (per_t * time_w).sum() / (time_w.sum() * B) + else: + loss["ego_planning_loss"] = F.smooth_l1_loss(ego_pred, ego_tgt) + + neighbor_pred = prediction[:, 1:] # [B, Pn, T, 4] + neighbor_tgt = gt_norm[:, 1:] + if bool(neighbors_future_valid.any()): + loss["neighbor_prediction_loss"] = F.smooth_l1_loss( + neighbor_pred[neighbors_future_valid], neighbor_tgt[neighbors_future_valid] + ) + else: + loss["neighbor_prediction_loss"] = torch.tensor(0.0, device=prediction.device) + + # With one mode, cross entropy is identically zero and its logits cannot + # affect the trajectory. Do not compute or log this PlantF-only no-op; + # train_epoch already treats mode_cls_loss as optional. Keep the original + # winner-mode classification objective unchanged for multimodal runs. + if probability.shape[-1] > 1: + loss["mode_cls_loss"] = F.cross_entropy(probability, best_mode.detach()) + + # Smoothness penalty: the planTF head regresses 80 absolute waypoints + # independently per timestep, which produces "comb" jitter (large second + # difference) even when ADE/FDE look fine. Penalizing the xy second + # difference of the best mode directly suppresses that jitter. Computed in + # the normalized space (same scale as ego_planning_loss). Off by default + # (coeff_smoothness_loss). See docs/plantf_original_comparison_and_roadmap.md. + best_xy = best_trajectory[:, :, :2] + second_diff = best_xy[:, 2:] - 2.0 * best_xy[:, 1:-1] + best_xy[:, :-2] + sq = (second_diff**2).sum(dim=-1) # [B, T-2] + # Loss improvement (#6): weight the curvature penalty toward the tail, where + # the divergence is worst. plantf_smoothness_tail_weight=0 -> uniform mean. + sm_tail = getattr(args, "plantf_smoothness_tail_weight", 0.0) + if sm_tail > 0: + ns = sq.shape[1] + ws = 1.0 + sm_tail * torch.arange(ns, device=sq.device) / max(ns - 1, 1) + loss["smoothness_loss"] = (sq * ws).sum() / (ws.sum() * sq.shape[0]) + else: + loss["smoothness_loss"] = sq.mean() + + # Compute ego edge points for penalty losses (best mode only) + need_ego_edge = args.coeff_road_border_loss > 0 or args.coeff_neighbor_collision_loss > 0 + if need_ego_edge: + ego_pred_world = best_trajectory * norm.std[0].to(trajectory.device) + norm.mean[0].to( + trajectory.device + ) # [B, T, 4] + ego_edge_points = compute_ego_edge_points( + ego_pred_world, inputs["ego_shape"], n_interp=args.road_border_n_interp + ) + denorm_inputs = args.observation_normalizer.inverse(inputs) + + if args.coeff_road_border_loss > 0: + rb_loss = compute_road_border_penalty( + ego_edge_points, + denorm_inputs["line_strings"], + margin=args.road_border_margin, + ) # [B, T] + loss["road_border_loss"] = rb_loss.mean() + else: + loss["road_border_loss"] = torch.tensor(0.0, device=prediction.device) + + if args.coeff_neighbor_collision_loss > 0: + nc_loss = compute_neighbor_collision_penalty( + ego_edge_points, + neighbors_future, + neighbors_future_valid, + denorm_inputs["neighbor_agents_past"], + margin_vehicle=args.neighbor_collision_margin_vehicle, + margin_pedestrian=args.neighbor_collision_margin_pedestrian, + margin_bicycle=args.neighbor_collision_margin_bicycle, + ) # [B, T] + loss["neighbor_collision_loss"] = nc_loss.mean() + else: + loss["neighbor_collision_loss"] = torch.tensor(0.0, device=prediction.device) + + assert not torch.isnan(loss["ego_planning_loss"]), "loss cannot be nan" + + turn_indicator_logit = decoder_output["turn_indicator_logit"] + turn_indicator_gt = make_turn_indicator_gt(inputs["turn_indicators"]) # [B,] + turn_indicator_loss = nn.functional.cross_entropy( + turn_indicator_logit, turn_indicator_gt, reduction="none" + ) + turn_indicator_change = inputs["turn_indicators"][:, -2] != inputs["turn_indicators"][:, -1] + turn_indicator_coeff = torch.where(turn_indicator_change, 1.0, 0.05) + turn_indicator_loss = (turn_indicator_loss * turn_indicator_coeff).mean() + loss["turn_indicator_loss"] = turn_indicator_loss + + independent_turn_indicator_logit = decoder_output["independent_turn_indicator_logit"] + independent_turn_indicator_loss = nn.functional.cross_entropy( + independent_turn_indicator_logit, turn_indicator_gt, reduction="none" + ) + independent_turn_indicator_loss = ( + independent_turn_indicator_loss * turn_indicator_coeff + ).mean() + loss["independent_turn_indicator_loss"] = independent_turn_indicator_loss + + with torch.no_grad(): + turn_indicator_accuracy = ( + (turn_indicator_logit.argmax(dim=-1) == turn_indicator_gt).float().mean() + ) + loss["turn_indicator_accuracy"] = turn_indicator_accuracy + + return loss diff --git a/diffusion_planner/diffusion_planner/train.py b/diffusion_planner/diffusion_planner/train.py index 453bcacb4..6d12832e9 100644 --- a/diffusion_planner/diffusion_planner/train.py +++ b/diffusion_planner/diffusion_planner/train.py @@ -6,7 +6,7 @@ import torch import wandb from timm.utils import ModelEma -from torch import optim +from torch import nn, optim from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data import DataLoader, DistributedSampler @@ -15,6 +15,11 @@ from diffusion_planner.train_config import TrainConfig from diffusion_planner.train_epoch import train_epoch from diffusion_planner.utils import ddp +from diffusion_planner.utils.checkpoint_viz import ( + render_checkpoint_trajectory_figure, + select_and_load_validation_sample, + select_and_load_validation_samples, +) from diffusion_planner.utils.data_augmentation import StatePerturbation from diffusion_planner.utils.data_augmentation_bridge import ( StatePerturbation as BridgeStatePerturbation, @@ -172,6 +177,159 @@ def closed_loop_validate(model, args, epoch: int, out_dir: str) -> None: ) +def render_checkpoint_visualization( + model, + args, + epoch: int, + checkpoint_dir: str, + checkpoint_name: str, + sample_label: str | None, + sample_path: Path, + sample_data, +) -> None: + """Render one validation sample and log the PNG to wandb.""" + if sample_label is None: + output_path = Path(checkpoint_dir) / "checkpoint_trajectory.png" + wandb_key = f"checkpoint_viz/{checkpoint_name}" + caption_label = checkpoint_name + else: + output_path = Path(checkpoint_dir) / f"checkpoint_trajectory_{sample_label}.png" + wandb_key = f"checkpoint_viz/{checkpoint_name}/{sample_label}" + caption_label = f"{checkpoint_name} [{sample_label}]" + try: + metrics = render_checkpoint_trajectory_figure( + model, + args, + sample_data, + sample_path, + output_path, + title=f"{caption_label} @epoch {epoch + 1}", + seed=args.seed, + ) + except Exception as exc: + print(f"Checkpoint visualization failed for {checkpoint_name} @epoch {epoch + 1}: {exc}") + return + + if args.use_wandb and wandb.run is not None: + wandb.log( + { + wandb_key: wandb.Image( + str(output_path), + caption=( + f"{caption_label} @epoch {epoch + 1} " + f"ADE={metrics['ADE']:.2f} FDE={metrics['FDE']:.2f} " + f"max_step={metrics['max_step']:.2f} roughness={metrics['roughness']:.2f}" + ), + ) + }, + step=epoch + 1, + ) + + +def render_checkpoint_visualizations( + model, + args, + epoch: int, + checkpoint_dir: str, + checkpoint_name: str, + viz_samples, +) -> None: + for sample_label, (sample_idx, sample_path, sample_data) in viz_samples.items(): + print( + f"Rendering checkpoint visualization [{checkpoint_name}/{sample_label}] " + f"sample_index={sample_idx} path={sample_path}" + ) + render_checkpoint_visualization( + model, + args, + epoch, + checkpoint_dir, + checkpoint_name, + sample_label, + sample_path, + sample_data, + ) + + +def render_checkpoint_legacy_visualization( + model, + args, + epoch: int, + checkpoint_dir: str, + checkpoint_name: str, + sample_path: Path, + sample_data, +) -> None: + render_checkpoint_visualization( + model, + args, + epoch, + checkpoint_dir, + checkpoint_name, + None, + sample_path, + sample_data, + ) + + +def warmstart_encoder(model, path, device): + """Load ONLY the encoder.* weights from a checkpoint (e.g. a diffusion model + trained on production data) into ``model.encoder``, leaving the decoder/head + untouched. Speeds up planTF training since the shared encoder is the bulk of + the params. ``strict=False`` so a slightly different config surfaces as + missing/unexpected keys instead of crashing. ``model`` must be the unwrapped + Diffusion_Planner (call before DDP wrap).""" + ckpt = torch.load(path, map_location=device) + sd = ckpt["model"] if isinstance(ckpt, dict) and "model" in ckpt else ckpt + # Check the DDP prefix first. Stripping ``len("encoder.")`` from a + # ``module.encoder.*`` key produces ``ncoder.*`` and silently turns every + # checkpoint tensor into an unexpected key under ``strict=False``. + prefixes = ("module.encoder.", "encoder.") + enc = {} + for key, value in sd.items(): + for prefix in prefixes: + if key.startswith(prefix): + enc[key[len(prefix) :]] = value + break + if not enc: + raise ValueError( + "[warm-start] no encoder.* or module.encoder.* keys found in " + f"checkpoint: {path}" + ) + + # ``strict=False`` still raises for same-named tensors with different + # shapes. This is expected when a DP checkpoint and the plantf data + # adapter expose a different number of map-feature channels. Transfer all + # compatible encoder weights and retain plantf's fresh initialization for + # those input projections. + target = model.encoder.state_dict() + compatible = { + key: value + for key, value in enc.items() + if key in target and target[key].shape == value.shape + } + skipped_shape = [ + key + for key, value in enc.items() + if key in target and target[key].shape != value.shape + ] + missing, unexpected = model.encoder.load_state_dict(compatible, strict=False) + print( + f"[warm-start] encoder: loaded {len(compatible)}/{len(enc)} keys from {path} | " + f"missing={len(missing)} unexpected={len(unexpected)} " + f"shape_skipped={len(skipped_shape)}" + ) + if not compatible: + raise ValueError( + "[warm-start] no compatible encoder tensors found in checkpoint: " + f"{path}" + ) + if skipped_shape: + print(f"[warm-start] kept fresh incompatible tensors: {skipped_shape}") + if unexpected: + print(f"[warm-start] NOTE unexpected keys (not in encoder): {unexpected[:5]}...") + + def model_training(args: TrainConfig): assert len(args.coeff_timestep) == 4, "coeff_timestep must be a list of 4 elements" @@ -294,6 +452,27 @@ def model_training(args: TrainConfig): 0 if valid_pair_loader is None else len(valid_pair_loader.dataset) ) ) + if args.enable_checkpoint_viz: + viz_sample_idx, viz_sample_path, viz_sample_data = select_and_load_validation_sample( + valid_set.data_list + ) + print( + f"Checkpoint legacy visualization sample: " + f"index={viz_sample_idx}, path={viz_sample_path}" + ) + viz_samples = select_and_load_validation_samples(valid_set.data_list) + for sample_label, ( + viz_sample_idx, + viz_sample_path, + _viz_sample_data, + ) in viz_samples.items(): + print( + f"Checkpoint visualization sample [{sample_label}]: " + f"index={viz_sample_idx}, path={viz_sample_path}" + ) + else: + viz_sample_path, viz_sample_data = None, None + viz_samples = {} if args.ddp: torch.distributed.barrier() @@ -302,6 +481,10 @@ def model_training(args: TrainConfig): diffusion_planner = Diffusion_Planner(args) diffusion_planner = diffusion_planner.to(rank if args.device == "cuda" else args.device) + # Warm-start the shared encoder from a pretrained checkpoint (before DDP wrap). + if getattr(args, "pretrained_encoder_path", None): + warmstart_encoder(diffusion_planner, args.pretrained_encoder_path, args.device) + if args.ddp: diffusion_planner = DDP(diffusion_planner, device_ids=[rank], find_unused_parameters=True) @@ -319,15 +502,34 @@ def model_training(args: TrainConfig): ) ) - # optimizer + # optimizer: weight decay on matmul weights only; biases and + # LayerNorm/BatchNorm/GroupNorm/Embedding params are excluded (standard + # transformer practice, matches original planTF). The previous single group + # relied on torch AdamW's default wd=0.01 applied to ALL params. + encoder_lr = args.encoder_learning_rate or args.learning_rate + encoder_param_ids = {id(param) for param in ddp.get_model(diffusion_planner, args.ddp).encoder.parameters()} + decay_params, no_decay_params = [], [] + encoder_decay_params, encoder_no_decay_params = [], [] + _no_decay_modules = (nn.LayerNorm, nn.BatchNorm1d, nn.BatchNorm2d, nn.GroupNorm, nn.Embedding) + for module in ddp.get_model(diffusion_planner, args.ddp).modules(): + for param_name, param in module.named_parameters(recurse=False): + if not param.requires_grad: + continue + is_no_decay = param_name.endswith("bias") or isinstance(module, _no_decay_modules) + if id(param) in encoder_param_ids: + (encoder_no_decay_params if is_no_decay else encoder_decay_params).append(param) + else: + (no_decay_params if is_no_decay else decay_params).append(param) params = [ - { - "params": ddp.get_model(diffusion_planner, args.ddp).parameters(), - "lr": args.learning_rate, - } + {"params": decay_params, "lr": args.learning_rate, "weight_decay": args.weight_decay, "group_name": "head_decay"}, + {"params": no_decay_params, "lr": args.learning_rate, "weight_decay": 0.0, "group_name": "head_no_decay"}, + {"params": encoder_decay_params, "lr": encoder_lr, "weight_decay": args.weight_decay, "group_name": "encoder_decay"}, + {"params": encoder_no_decay_params, "lr": encoder_lr, "weight_decay": 0.0, "group_name": "encoder_no_decay"}, ] optimizer = optim.AdamW(params) + if global_rank == 0: + print(f"Optimizer base LR: head={args.learning_rate}, encoder={encoder_lr}") scheduler = CosineAnnealingWarmUpRestarts(optimizer, train_epochs, args.warm_up_epoch) if args.resume_model_path is not None: @@ -339,8 +541,8 @@ def model_training(args: TrainConfig): # Override learning rate with the new value for param_group in optimizer.param_groups: - param_group["lr"] = args.learning_rate - print(f"Learning rate reset to {args.learning_rate}") + param_group["lr"] = encoder_lr if param_group.get("group_name", "").startswith("encoder_") else args.learning_rate + print(f"Learning rate reset: head={args.learning_rate}, encoder={encoder_lr}") else: init_epoch = 0 @@ -381,6 +583,8 @@ def model_training(args: TrainConfig): valid_loss_ego = agg["avg_loss_ego"] valid_loss_neighbor = agg["avg_loss_neighbor"] mean_ego_loss_dict = {f"valid_loss/{k}": v for k, v in agg["ego_means"].items()} + mean_plantf_metric_dict = {f"valid_mode/{k}": v for k, v in agg["plantf_means"].items()} + mean_traj_metric_dict = {f"valid_traj/{k}": v for k, v in agg["traj_means"].items()} mean_epdms_dict = wandb_epdms_metrics(agg["epdms_means"]) valid_loss_ego_position_lat_loss = mean_ego_loss_dict.get( "valid_loss/ego_position_lat_loss", 0.0 @@ -400,6 +604,10 @@ def model_training(args: TrainConfig): f"{turn_indicator_change_accuracy=:.3f}\n" f"{turn_indicator_change_total=:.3f}" ) + for key, value in mean_plantf_metric_dict.items(): + print(f"{key}={value:.4f}") + for key, value in mean_traj_metric_dict.items(): + print(f"{key}={value:.4f}") if replan_agg.get("replan_consistency_count", 0) > 0: print( "replan_position_consistency={:.3f}\n" @@ -413,22 +621,28 @@ def model_training(args: TrainConfig): # begin training for epoch in range(init_epoch, train_epochs): + # Encoder freeze schedule: freeze for the first freeze_encoder_epochs (train + # the head only), then unfreeze for joint fine-tuning. Toggling requires_grad + # is safe because the encoder params are already in the optimizer (built + # while trainable) — frozen params simply receive no gradient and are skipped + # by optimizer.step, then resume once unfrozen. + if args.freeze_encoder_epochs > 0: + frozen = epoch < args.freeze_encoder_epochs + enc = ddp.get_model(diffusion_planner, args.ddp).encoder + for p in enc.parameters(): + p.requires_grad = not frozen + if epoch in (init_epoch, args.freeze_encoder_epochs): + print( + f"[freeze] epoch {epoch}: encoder {'FROZEN (head-only)' if frozen else 'trainable (joint)'}" + ) + # Synchronize all processes before training if args.ddp: torch.distributed.barrier() - # Adjust learning rate for final 10 epochs - final_epoch_count = 10 - if epoch >= train_epochs - final_epoch_count: - base_lr = args.learning_rate - if epoch >= train_epochs - final_epoch_count // 2: # Last 5 epochs: LR * 1/100 - adjusted_lr = base_lr * 0.01 - else: # First 5 of final 10 epochs: LR * 1/10 - adjusted_lr = base_lr * 0.1 - for param_group in optimizer.param_groups: - param_group["lr"] = adjusted_lr - if global_rank == 0: - print(f"Final phase: Epoch {epoch + 1}, LR adjusted to {adjusted_lr}") + # (The old "final 10 epochs" manual lr step-down was removed: the lr + # schedule now does a proper cosine decay to eta_min, so a second manual + # decay would double-count. See utils/lr_schedule.py.) # training step train_loss, train_total_loss = train_epoch( @@ -445,6 +659,8 @@ def model_training(args: TrainConfig): valid_loss_ego = agg["avg_loss_ego"] valid_loss_neighbor = agg["avg_loss_neighbor"] mean_ego_loss_dict = {f"valid_loss/{k}": v for k, v in agg["ego_means"].items()} + mean_plantf_metric_dict = {f"valid_mode/{k}": v for k, v in agg["plantf_means"].items()} + mean_traj_metric_dict = {f"valid_traj/{k}": v for k, v in agg["traj_means"].items()} replan_loss_dict = {f"valid_loss/{k}": v for k, v in replan_agg.items()} mean_epdms_dict = wandb_epdms_metrics(agg["epdms_means"]) valid_loss_ego_position_lat_loss = mean_ego_loss_dict.get( @@ -466,6 +682,10 @@ def model_training(args: TrainConfig): f"{turn_indicator_change_accuracy=:.3f}\n" f"{turn_indicator_change_total=:.3f}" ) + for key, value in mean_plantf_metric_dict.items(): + print(f"{key}={value:.4f}") + for key, value in mean_traj_metric_dict.items(): + print(f"{key}={value:.4f}") if replan_agg.get("replan_consistency_count", 0) > 0: print( "replan_position_consistency={:.3f}\n" @@ -487,6 +707,8 @@ def model_training(args: TrainConfig): "valid_loss/turn_indicator_accuracy": turn_indicator_accuracy, "valid_loss/turn_indicator_change_accuracy": turn_indicator_change_accuracy, **mean_ego_loss_dict, + **mean_plantf_metric_dict, + **mean_traj_metric_dict, **replan_loss_dict, **mean_epdms_dict, }, @@ -500,6 +722,8 @@ def model_training(args: TrainConfig): "valid_loss_neighbor": valid_loss_neighbor, "valid_loss_ego_position_lat_loss": valid_loss_ego_position_lat_loss, "valid_loss_ego_position_lon_loss": valid_loss_ego_position_lon_loss, + **{k.replace("/", "_"): v for k, v in mean_plantf_metric_dict.items()}, + **{k.replace("/", "_"): v for k, v in mean_traj_metric_dict.items()}, **replan_agg, **{k.replace("/", "_"): v for k, v in mean_epdms_dict.items()}, } @@ -528,16 +752,35 @@ def model_training(args: TrainConfig): with open(os.path.join(curr_dir, "args.json"), "w", encoding="utf-8") as f: json.dump(args_dict, f, indent=4) # Export ONNX next to the checkpoint (regular weights, ORT validation skipped). - export_checkpoint_onnx_guarded( - config_json_path=os.path.join(curr_dir, "args.json"), - ckpt_path=f"{curr_dir}/best_model.pth", - output_dir=Path(curr_dir), - output_prefix="diffusion_planner", - use_ema=False, - use_simplify=False, - opset_version=20, - external_data=False, - ) + if args.enable_onnx_export: + export_checkpoint_onnx_guarded( + config_json_path=os.path.join(curr_dir, "args.json"), + ckpt_path=f"{curr_dir}/best_model.pth", + output_dir=Path(curr_dir), + output_prefix="diffusion_planner", + use_ema=False, + use_simplify=False, + opset_version=20, + external_data=False, + ) + if args.enable_checkpoint_viz: + render_checkpoint_legacy_visualization( + model_ema.ema if args.use_ema else diffusion_planner, + args, + epoch, + curr_dir, + f"epoch{epoch + 1:04d}", + viz_sample_path, + viz_sample_data, + ) + render_checkpoint_visualizations( + model_ema.ema if args.use_ema else diffusion_planner, + args, + epoch, + curr_dir, + f"epoch{epoch + 1:04d}", + viz_samples, + ) # Closed-loop validation runs on the same cadence as the checkpoint save; outputs # (videos + metrics) land next to the saved weights they correspond to. closed_loop_validate( @@ -555,16 +798,35 @@ def model_training(args: TrainConfig): with open(os.path.join(curr_dir, "args.json"), "w", encoding="utf-8") as f: json.dump(args_dict, f, indent=4) # Export ONNX next to the checkpoint (regular weights, ORT validation skipped). - export_checkpoint_onnx_guarded( - config_json_path=os.path.join(curr_dir, "args.json"), - ckpt_path=f"{curr_dir}/best_model.pth", - output_dir=Path(curr_dir), - output_prefix="diffusion_planner", - use_ema=False, - use_simplify=False, - opset_version=20, - external_data=False, - ) + if args.enable_onnx_export: + export_checkpoint_onnx_guarded( + config_json_path=os.path.join(curr_dir, "args.json"), + ckpt_path=f"{curr_dir}/best_model.pth", + output_dir=Path(curr_dir), + output_prefix="diffusion_planner", + use_ema=False, + use_simplify=False, + opset_version=20, + external_data=False, + ) + if args.enable_checkpoint_viz: + render_checkpoint_legacy_visualization( + model_ema.ema if args.use_ema else diffusion_planner, + args, + epoch, + curr_dir, + "best_model", + viz_sample_path, + viz_sample_data, + ) + render_checkpoint_visualizations( + model_ema.ema if args.use_ema else diffusion_planner, + args, + epoch, + curr_dir, + "best_model", + viz_samples, + ) scheduler.step() train_sampler.set_epoch(epoch + 1) diff --git a/diffusion_planner/diffusion_planner/train_config.py b/diffusion_planner/diffusion_planner/train_config.py index af8cfb5d4..57c7b0fb6 100644 --- a/diffusion_planner/diffusion_planner/train_config.py +++ b/diffusion_planner/diffusion_planner/train_config.py @@ -73,6 +73,14 @@ class TrainConfig: batch_size: int = 512 save_utd: int = 10 learning_rate: float = 1e-4 + # Optional encoder-specific base LR. When None, the encoder uses + # ``learning_rate`` just like the decoder/head. + encoder_learning_rate: Optional[float] = None + # AdamW weight decay. Applied only to matmul weights; biases and + # LayerNorm/BatchNorm/Embedding params are excluded (standard transformer + # practice, matches original planTF). Previously torch AdamW's default 0.01 + # was applied to ALL params including norms/biases, over-regularizing. + weight_decay: float = 1e-4 warm_up_epoch: int = 5 encoder_drop_path_rate: float = 0.1 decoder_drop_path_rate: float = 0.1 @@ -108,6 +116,83 @@ class TrainConfig: alpha_planning_loss: float = 1.0 alpha_neighbor_loss: float = 0.1 + # Mode classification loss weight (decoder_type="plantf" only) + alpha_mode_cls_loss: float = 1.0 + # planTF loss shaping (docs/plantf_dead_mode_improvement.md). Both default + # to the diffusion head's original loss behavior. The "dango" collapse + # turned out to be driven by modes=1 (single-mode L2 = mean regression) and + # was already resolved by modes>=2 with WTA; the earlier planTF-specific + # overrides (endpoint loss on, lon down-weighting off) made results WORSE in + # A/B and are reverted. Both remain as opt-in experiment knobs. + coeff_endpoint_fde_loss: float = 0.0 + plantf_use_lon_velocity_weight: bool = True + # Smoothness penalty on the xy second difference of the best planTF mode, + # to suppress the "comb" jitter of the per-timestep absolute regression. + # Off by default; try 0.1-1.0. See docs/plantf_original_comparison_and_roadmap.md. + coeff_smoothness_loss: float = 0.0 + # Feed the current ego motion state (vx, vy, ax, ay, steering, yaw_rate) into + # the planTF trajectory head so its absolute-waypoint regression is anchored + # to the current motion (the diffusion decoder gets this via its pinned + # current state; the planTF head otherwise does not). See + # docs/plantf_dead_mode_improvement.md. + plantf_use_ego_state_in_head: bool = True + # A2 (docs/plantf_original_comparison_and_roadmap.md): replace the ego encoder + # token with an embedding of the current ego motion state (vx,vy,ax,ay,steer, + # yaw_rate) instead of the position-history token, matching original planTF's + # use_ego_history=false path. state_dropout randomly zeroes those channels + # during training (original uses 0.75). Anchors the whole prediction to the + # current motion the way the diffusion decoder's pinned state does. + plantf_ego_state_token: bool = False + plantf_ego_state_dropout: float = 0.75 + # C1: feed agent (ego + neighbor) history as consecutive-frame xy deltas + # instead of absolute positions, giving the encoder temporal-difference + # inputs (original planTF vectorizes history this way). + plantf_input_delta: bool = False + # Zero the goal_pose input. The mini goal_pose is global-frame (data bug) while + # Autoware feeds ego-frame -> train/deploy mismatch. Masking makes the model + # plan from the route (route_lanes carries the goal). Bakes into ONNX. + plantf_mask_goal_pose: bool = False + # Original planTF predicts every agent's future xy relative to that agent's + # current xy. At inference the decoder adds the observed current xy back, + # so the returned DP-compatible trajectory remains in absolute ego-frame + # coordinates. Keep this opt-in for checkpoint compatibility: existing + # absolute-xy checkpoints must not be interpreted as relative-xy models. + plantf_relative_xy: bool = False + # --- planTF combinable ablation toggles (docs/plantf_head_development_notes.md §9) --- + # Trajectory head architecture. "mlp" (default): reshape the single ego token + # into K modes. "cross_attn": K mode queries cross-attend to ALL encoder + # tokens (map/agents/route), giving the head scene context the single-token + # bottleneck lacks (candidate fix for tail divergence). "basis": the mlp head + # but the trajectory is a Bezier curve of plantf_basis_control_points control + # points expanded over time — a temporal inductive bias (structural C∞ + # smoothness) the flat per-step head lacks, while keeping mlp's deploy + # robustness and ONNX-triviality (fixed basis matmul, no recurrence). + # "gru": recurrent head that unrolls the waypoints with a GRU (temporal + # recurrence in the architecture). Experimental — RNN heads are more + # deploy-fragile than mlp/basis; not the deploy default. + plantf_head_type: Literal["mlp", "cross_attn", "basis", "gru"] = "mlp" + # Number of Bezier control points for plantf_head_type="basis" (ignored + # otherwise). Fewer = smoother/stiffer, more = more expressive. 8 is a good + # default for an 8s / 80-step horizon. + plantf_basis_control_points: int = 8 + # Inference mode selection. When True, 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. Applies to the + # validation forward path only, not the ONNX deploy graph. + plantf_route_rerank: bool = False + plantf_route_rerank_topk: int = 3 + # Loss improvements (docs/plantf_head_development_notes.md §9), all opt-in and + # independently combinable: + # - plantf_tail_weight (>0): weight the per-timestep ego regression toward the + # tail (w_t = 1 + tail_weight * t/(T-1)). + # - plantf_smoothness_tail_weight (>0, #6): same tail weighting on the curvature + # (second-difference) penalty. + # - plantf_use_laplace_nll (#5): replace the ego smooth-L1 with a Laplace NLL + # using a head-predicted per-point log-scale (planTF's probabilistic + # regression; calibrates tail uncertainty and mode confidence). + plantf_tail_weight: float = 0.0 + plantf_smoothness_tail_weight: float = 0.0 + plantf_use_laplace_nll: bool = False # Velocity Representation & Hybrid Loss use_velocity_representation: bool = False @@ -127,8 +212,25 @@ class TrainConfig: num_heads: int = 8 hidden_dim: int = 256 diffusion_model_type: Literal["x_start", "flow_matching"] = "x_start" + # "plantf" replaces the diffusion decoder with a one-shot multi-modal + # regression head (planTF, Cheng et al. ICRA 2024) on the same encoder. + decoder_type: Literal["diffusion", "plantf"] = "diffusion" + # Number of ego trajectory modes (decoder_type="plantf" only) + num_modes: int = 6 predicted_neighbor_num: int = MAX_NUM_NEIGHBORS resume_model_path: Optional[str] = None + # Warm-start: load ONLY the encoder weights (encoder.* keys) from a checkpoint + # (e.g. a diffusion model trained on production data) into this model, leaving + # the decoder/head randomly initialized. The shared encoder is the bulk of the + # params, so this speeds up planTF convergence. IMPORTANT: use the SAME + # normalization the pretrained encoder was trained with, else the encoder sees + # an out-of-distribution input and its features are miscalibrated (unless you + # also fine-tune it, i.e. freeze_encoder_epochs=0). + pretrained_encoder_path: Optional[str] = None + # Freeze the encoder for the first N epochs (train the head only), then unfreeze + # for joint fine-tuning. Only sensible when the normalization matches the + # pretrained encoder (frozen features must match the input distribution). + freeze_encoder_epochs: int = 0 # --------------------------------------------------------- # Logging & Distributed Setup @@ -139,6 +241,12 @@ class TrainConfig: notes: str = "" ddp: bool = True port: str = "22323" + # Portability toggles: per-checkpoint trajectory rendering (matplotlib + + # per-checkpoint forward pass) and ONNX export. Both default on to preserve + # behavior, but can be disabled for lightweight / dependency-free training + # runs on other environments. + enable_checkpoint_viz: bool = True + enable_onnx_export: bool = True # Validation-only temporal stability metrics. Replan consistency requires full-sequence # Step-1 NPZ frames in valid_set_list; the default gap=1 avoids treating skip-N lists diff --git a/diffusion_planner/diffusion_planner/train_epoch.py b/diffusion_planner/diffusion_planner/train_epoch.py index be14b4782..73cf409d4 100644 --- a/diffusion_planner/diffusion_planner/train_epoch.py +++ b/diffusion_planner/diffusion_planner/train_epoch.py @@ -3,6 +3,7 @@ from tqdm import tqdm from diffusion_planner.model.module.decoder import compute_training_loss +from diffusion_planner.model.module.plantf_decoder import compute_plantf_training_loss from diffusion_planner.utils import ddp from diffusion_planner.utils.data_augmentation import StatePerturbation from diffusion_planner.utils.train_utils import compute_grad_stats, get_epoch_mean_loss @@ -65,7 +66,12 @@ def train_epoch(data_loader, model, optimizer, args, ema, aug: StatePerturbation # call the model optimizer.zero_grad() - loss = compute_training_loss(model, inputs, (ego_future, neighbors_future, mask), args) + if getattr(args, "decoder_type", "diffusion") == "plantf": + loss = compute_plantf_training_loss( + model, inputs, (ego_future, neighbors_future, mask), args + ) + else: + loss = compute_training_loss(model, inputs, (ego_future, neighbors_future, mask), args) loss["loss"] = ( args.alpha_neighbor_loss * loss["neighbor_prediction_loss"] @@ -75,6 +81,12 @@ def train_epoch(data_loader, model, optimizer, args, ema, aug: StatePerturbation + args.coeff_road_border_loss * loss["road_border_loss"] + args.coeff_neighbor_collision_loss * loss["neighbor_collision_loss"] ) + if "mode_cls_loss" in loss: + loss["loss"] = loss["loss"] + args.alpha_mode_cls_loss * loss["mode_cls_loss"] + if "endpoint_fde_loss" in loss: + loss["loss"] = loss["loss"] + args.coeff_endpoint_fde_loss * loss["endpoint_fde_loss"] + if "smoothness_loss" in loss: + loss["loss"] = loss["loss"] + args.coeff_smoothness_loss * loss["smoothness_loss"] # loss backward loss["loss"].backward() diff --git a/diffusion_planner/diffusion_planner/utils/checkpoint_viz.py b/diffusion_planner/diffusion_planner/utils/checkpoint_viz.py new file mode 100644 index 000000000..abf9d0894 --- /dev/null +++ b/diffusion_planner/diffusion_planner/utils/checkpoint_viz.py @@ -0,0 +1,191 @@ +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import torch + +from diffusion_planner.train_epoch import heading_to_cos_sin +from diffusion_planner.utils import ddp +from diffusion_planner.validate_model import _prepare_validation_inputs + +# matplotlib and visualize_inputs are imported lazily inside +# render_checkpoint_trajectory_figure so that importing this module (which +# train.py always does) does not require matplotlib when checkpoint viz is +# disabled (enable_checkpoint_viz=False). select_and_load_validation_sample and +# the metric helpers below use only numpy/torch. + + +def _trajectory_metrics(prediction_xy: np.ndarray, gt_xy: np.ndarray) -> dict[str, float]: + point_error = np.linalg.norm(prediction_xy - gt_xy, axis=-1) + steps = np.diff(np.vstack(([0.0, 0.0], prediction_xy)), axis=0) + second_difference = np.diff(prediction_xy, n=2, axis=0) + return { + "ADE": float(point_error.mean()), + "FDE": float(point_error[-1]), + "max_step": float(np.linalg.norm(steps, axis=-1).max()), + "roughness": float(np.linalg.norm(second_difference, axis=-1).mean()), + } + + +def _select_validation_sample( + data_list: list[str], + min_gt_endpoint_m: float | None = None, + max_gt_endpoint_m: float | None = None, +) -> tuple[int, Path]: + if not data_list: + raise ValueError("Validation list is empty; cannot create checkpoint visualization.") + + fallback = Path(data_list[0]) + for idx, path_str in enumerate(data_list): + path = Path(path_str) + with np.load(path, allow_pickle=True) as data: + gt_xy = np.asarray(data["ego_agent_future"], dtype=np.float32)[:80, :2] + endpoint = float(np.linalg.norm(gt_xy[-1])) + if min_gt_endpoint_m is not None and endpoint < min_gt_endpoint_m: + continue + if max_gt_endpoint_m is not None and endpoint >= max_gt_endpoint_m: + continue + return idx, path + return 0, fallback + + +def _load_raw_sample(npz_path: Path) -> dict[str, np.ndarray]: + with np.load(npz_path, allow_pickle=True) as data: + sample = {key: np.asarray(value) for key, value in data.items() if key != "version"} + return sample + + +def _batchify_raw_sample(raw_sample: dict[str, np.ndarray]) -> dict[str, torch.Tensor]: + batch = {} + for key, value in raw_sample.items(): + tensor = torch.as_tensor(value) + if tensor.ndim == 0: + tensor = tensor.unsqueeze(0) + else: + tensor = tensor.unsqueeze(0) + batch[key] = tensor + batch["ego_agent_past"] = heading_to_cos_sin(batch["ego_agent_past"]) + batch["goal_pose"] = heading_to_cos_sin(batch["goal_pose"]) + return batch + + +def render_checkpoint_trajectory_figure( + model, + args, + raw_sample: dict[str, np.ndarray], + sample_path: Path, + output_path: Path, + *, + title: str, + seed: int, +) -> dict[str, float]: + """Render a single validation scene with GT and the checkpoint prediction.""" + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + from diffusion_planner.utils.visualize_input import visualize_inputs + + device = torch.device(args.device) + output_path.parent.mkdir(parents=True, exist_ok=True) + + plot_inputs = _batchify_raw_sample(raw_sample) + prepared = _prepare_validation_inputs(plot_inputs, args, device) + + net = ddp.get_model(model, args.ddp) + was_training = net.training + net.eval() + try: + device_index = [torch.cuda.current_device()] if device.type == "cuda" else [] + with torch.no_grad(), torch.random.fork_rng(devices=device_index): + torch.manual_seed(seed) + _, outputs = net(prepared.inputs) + finally: + net.train(was_training) + + if "prediction" not in outputs: + raise KeyError("Model output does not contain `prediction`.") + + prediction_xy = outputs["prediction"][0, 0, :, :2].detach().cpu().numpy() + gt_xy = prepared.ego_future[0, :, :2].detach().cpu().numpy() + metrics = _trajectory_metrics(prediction_xy, gt_xy) + + fig, ax = plt.subplots(figsize=(11, 7)) + visualize_inputs(plot_inputs, save_path=None, ax=ax) + + origin = np.zeros((1, 2), dtype=np.float32) + ax.plot( + *np.vstack((origin, gt_xy)).T, + color="black", + linestyle="--", + linewidth=2.2, + label="Ground truth", + ) + ax.plot( + *np.vstack((origin, prediction_xy)).T, + color="tab:red", + linewidth=1.8, + marker=".", + markersize=3, + label="Checkpoint prediction", + ) + ax.scatter([0.0], [0.0], marker="*", s=110, color="gold", edgecolor="black", zorder=5) + + plotted_xy = np.vstack((origin, gt_xy, prediction_xy)) + xy_min = plotted_xy.min(axis=0) + xy_max = plotted_xy.max(axis=0) + span = np.maximum(xy_max - xy_min, 1.0) + padding = np.maximum(0.12 * span, 2.0) + ax.set_xlim(xy_min[0] - padding[0], xy_max[0] + padding[0]) + ax.set_ylim(xy_min[1] - padding[1], xy_max[1] + padding[1]) + ax.set_aspect("equal", adjustable="box") + ax.grid(True, alpha=0.25) + ax.set_xlabel("x in ego frame [m]") + ax.set_ylabel("y in ego frame [m]") + ax.set_title(f"{title}\n{sample_path.name}") + ax.legend(loc="best") + ax.text( + 0.01, + 0.01, + ( + f"ADE={metrics['ADE']:.2f} m, FDE={metrics['FDE']:.2f} m\n" + f"max step={metrics['max_step']:.2f} m, roughness={metrics['roughness']:.2f} m" + ), + transform=ax.transAxes, + fontsize=9, + verticalalignment="bottom", + bbox={"boxstyle": "round", "facecolor": "white", "alpha": 0.85}, + ) + fig.tight_layout() + fig.savefig(output_path, dpi=160) + plt.close(fig) + return metrics + + +def select_and_load_validation_sample( + data_list: list[str], min_gt_endpoint_m: float = 10.0 +) -> tuple[int, Path, dict[str, np.ndarray]]: + idx, path = _select_validation_sample(data_list, min_gt_endpoint_m=min_gt_endpoint_m) + return idx, path, _load_raw_sample(path) + + +def select_and_load_validation_samples( + data_list: list[str], +) -> dict[str, tuple[int, Path, dict[str, np.ndarray]]]: + """Select representative validation scenes for stop / slow / move bins.""" + sample_specs = { + "stop": (None, 2.0), + "slow": (2.0, 10.0), + "move": (10.0, None), + } + samples: dict[str, tuple[int, Path, dict[str, np.ndarray]]] = {} + for label, (min_gt_endpoint_m, max_gt_endpoint_m) in sample_specs.items(): + idx, path = _select_validation_sample( + data_list, + min_gt_endpoint_m=min_gt_endpoint_m, + max_gt_endpoint_m=max_gt_endpoint_m, + ) + samples[label] = (idx, path, _load_raw_sample(path)) + return samples diff --git a/diffusion_planner/diffusion_planner/utils/dataset.py b/diffusion_planner/diffusion_planner/utils/dataset.py index 905c2be3e..dc1d9b33a 100644 --- a/diffusion_planner/diffusion_planner/utils/dataset.py +++ b/diffusion_planner/diffusion_planner/utils/dataset.py @@ -16,9 +16,35 @@ def __getitem__(self, idx): data = np.load(self.data_list[idx], allow_pickle=True) data = dict(data) # npz to dict data.pop("version", None) + data["neighbor_agents_future"] = _neighbor_future_to_cos_sin( + data.get("neighbor_agents_future") + ) return data +def _neighbor_future_to_cos_sin(nf): + """Normalize neighbor_agents_future to 4 channels (x, y, cos, sin) at load time. + + Some converter versions store it as 3 channels (x, y, heading), others as 4 + (x, y, cos, sin). ``train_epoch`` applies ``heading_to_cos_sin`` AFTER the + DataLoader collate, so a batch mixing 3- and 4-channel datasets fails to + collate ("Trying to resize storage that is not resizable"). Converting here — + before collation — lets differently-generated datasets (e.g. odaiba 4ch + + mini 3ch) be combined. Padded (all-zero) rows stay all-zero so downstream + masking is unchanged; 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. + """ + if nf is None: + return nf + nf = np.asarray(nf, dtype=np.float32) + if nf.shape[-1] != 3: + return nf + valid = np.any(nf != 0, axis=-1, keepdims=True) + cs = np.concatenate([nf[..., :2], np.cos(nf[..., 2:3]), np.sin(nf[..., 2:3])], axis=-1) + return np.where(valid, cs, 0.0).astype(np.float32) + + class DiffusionPlannerPairData(Dataset): def __init__(self, data_list, expected_gap: int | None = None): paths = openjson(data_list) diff --git a/diffusion_planner/diffusion_planner/utils/lr_schedule.py b/diffusion_planner/diffusion_planner/utils/lr_schedule.py index 3b1185f6d..a50bb14de 100644 --- a/diffusion_planner/diffusion_planner/utils/lr_schedule.py +++ b/diffusion_planner/diffusion_planner/utils/lr_schedule.py @@ -1,15 +1,22 @@ -from torch.optim.lr_scheduler import LinearLR, MultiplicativeLR, SequentialLR +from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR -def CosineAnnealingWarmUpRestarts(optimizer, epoch, warm_up_epoch, start_factor=0.1): - assert epoch >= warm_up_epoch - T_warmup = warm_up_epoch - - warmup_scheduler = LinearLR(optimizer, start_factor=start_factor, total_iters=warm_up_epoch - 1) - fixed_scheduler = MultiplicativeLR(optimizer, lr_lambda=lambda epoch: 1.0) +def CosineAnnealingWarmUpRestarts(optimizer, epoch, warm_up_epoch, start_factor=0.1, eta_min=1e-6): + """Linear warmup for ``warm_up_epoch`` epochs, then cosine decay to + ``eta_min`` over the remaining epochs. - scheduler = SequentialLR( - optimizer, schedulers=[warmup_scheduler, fixed_scheduler], milestones=[T_warmup] - ) + Previously the post-warmup phase used ``MultiplicativeLR(lambda=1.0)``, i.e. + it held the learning rate CONSTANT after warmup (despite the name). That let + training collapse right after warmup — every model, diffusion included, + peaked in the warmup epochs and then degraded. Decaying the lr with a proper + cosine schedule (matching original planTF) keeps the lr low enough after the + initial phase to actually converge. See + docs/plantf_original_comparison_and_roadmap.md. - return scheduler + Called with ``scheduler.step()`` once per epoch, so the sub-schedulers count + in epochs. + """ + assert epoch >= warm_up_epoch + warmup = LinearLR(optimizer, start_factor=start_factor, total_iters=max(warm_up_epoch - 1, 1)) + cosine = CosineAnnealingLR(optimizer, T_max=max(epoch - warm_up_epoch, 1), eta_min=eta_min) + return SequentialLR(optimizer, schedulers=[warmup, cosine], milestones=[warm_up_epoch]) diff --git a/diffusion_planner/diffusion_planner/utils/normalizer.py b/diffusion_planner/diffusion_planner/utils/normalizer.py index e4d1318af..e75f2f294 100644 --- a/diffusion_planner/diffusion_planner/utils/normalizer.py +++ b/diffusion_planner/diffusion_planner/utils/normalizer.py @@ -57,8 +57,12 @@ def __call__(self, data): if k not in data: # Check if key `k` exists in `data` continue mask = torch.sum(torch.ne(data[k], 0), dim=-1) == 0 - norm_data[k] = (data[k] - v["mean"].to(data[k].device)) / v["std"].to(data[k].device) - norm_data[k][mask] = 0 + normalized = (data[k] - v["mean"].to(data[k].device)) / v["std"].to(data[k].device) + # Avoid boolean-index assignment here. ONNX lowers ``x[mask] = 0`` + # to a Where/Scatter pattern with an incompatible condition shape + # for TensorRT when ``data`` has a trailing feature dimension. + # Make that broadcast explicit instead. + norm_data[k] = torch.where(mask.unsqueeze(-1), torch.zeros_like(normalized), normalized) return norm_data def inverse(self, data): @@ -67,8 +71,10 @@ def inverse(self, data): if k not in data: # Check if key `k` exists in `data` continue mask = torch.sum(torch.ne(data[k], 0), dim=-1) == 0 - norm_data[k] = data[k] * v["std"].to(data[k].device) + v["mean"].to(data[k].device) - norm_data[k][mask] = 0 + denormalized = data[k] * v["std"].to(data[k].device) + v["mean"].to(data[k].device) + # Keep the inverse path ONNX/TensorRT-safe for the same reason as + # the forward normalization path above. + norm_data[k] = torch.where(mask.unsqueeze(-1), torch.zeros_like(denormalized), denormalized) return norm_data def to_dict(self): diff --git a/diffusion_planner/diffusion_planner/utils/onnx_export.py b/diffusion_planner/diffusion_planner/utils/onnx_export.py index 291438862..a7602b724 100644 --- a/diffusion_planner/diffusion_planner/utils/onnx_export.py +++ b/diffusion_planner/diffusion_planner/utils/onnx_export.py @@ -22,6 +22,7 @@ from diffusion_planner.dimensions import * from diffusion_planner.model.diffusion_planner import Diffusion_Planner +from diffusion_planner.model.module.plantf_decoder import PlanTFDecoder from diffusion_planner.utils.config import Config FULL_INPUT_NAMES = [ @@ -70,10 +71,17 @@ TURN_INDICATOR_INPUT_NAMES = ["encoding", "final_x0"] +# planTF head: one-shot, no external denoising loop and no separate +# turn-indicator graph. It nevertheless retains the DP decoder's four-input +# interface so relative-xy PlantF can reconstruct agent-current anchors and +# existing deployment bindings do not need a head-specific decoder contract. +PLANTF_DECODER_INPUT_NAMES = DECODER_INPUT_NAMES + FULL_OUTPUT_NAMES = ["prediction", "turn_indicator_logit"] ENCODER_OUTPUT_NAMES = ["encoding"] DECODER_OUTPUT_NAMES = ["model_output"] TURN_INDICATOR_OUTPUT_NAMES = ["turn_indicator_logit"] +PLANTF_DECODER_OUTPUT_NAMES = ["prediction", "probability", "turn_indicator_logit"] TensorDict = dict[str, torch.Tensor] NumpyDict = dict[str, np.ndarray] @@ -84,7 +92,8 @@ class ModelWrappers: full: nn.Module encoder: nn.Module decoder: nn.Module - turn_indicator: nn.Module + # None for the planTF head, whose decoder graph already emits the logit + turn_indicator: nn.Module | None @dataclass(frozen=True) @@ -219,6 +228,33 @@ def forward(self, encoding: torch.Tensor, final_x0: torch.Tensor) -> torch.Tenso return self.decoder._compute_turn_indicator(ego_trajectory, encoding_pooled) +class PlanTFDecoderONNXWrapper(nn.Module): + """One-shot planTF head: encoder output -> final prediction in a single call. + + There is no external denoising loop; the graph also emits the mode logits + and the turn-indicator logit, so no separate turn-indicator graph exists. + """ + + def __init__(self, model: Diffusion_Planner): + super().__init__() + self.decoder = model.decoder + + def forward( + self, + encoding: torch.Tensor, + sampled_trajectories: torch.Tensor, + diffusion_time: torch.Tensor, + neighbor_agents_past: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + prediction, probability, turn_indicator_logit = self.decoder.forward_deploy( + encoding, sampled_trajectories, diffusion_time, neighbor_agents_past + ) + # Keep diffusion_time in the exported graph despite PlanTF's one-shot + # decoder; deployment binds the same four DP decoder inputs. + keep_alive = 0.0 * diffusion_time[:, 0, 0, 0] + return prediction + keep_alive[:, None, None, None], probability, turn_indicator_logit + + class FullONNXWrapper(nn.Module): """Original all-in-one planner export.""" @@ -269,6 +305,62 @@ def forward( return decoder_outputs["prediction"], decoder_outputs["turn_indicator_logit"] +class PlanTFFullONNXWrapper(FullONNXWrapper): + """planTF full graph with the same 17-input contract as the diffusion one. + + The planTF decoder ignores the diffusion-only inputs (sampled_trajectories, + ego_current_state, delay) and the legacy exporter prunes unused graph inputs, + which would break the deployed Autoware node — it feeds all 17 inputs and both + ONNX Runtime and TensorRT reject feeding/binding tensors absent from the graph. + A zero-valued residual keeps them anchored. Scalar slices (not full sums) so the + dead branch cannot overflow/NaN under reduced-precision runtimes. + """ + + def forward( + self, + sampled_trajectories: torch.Tensor, + ego_agent_past: torch.Tensor, + ego_current_state: torch.Tensor, + neighbor_agents_past: torch.Tensor, + static_objects: torch.Tensor, + lanes: torch.Tensor, + lanes_speed_limit: torch.Tensor, + lanes_has_speed_limit: torch.Tensor, + route_lanes: torch.Tensor, + route_lanes_speed_limit: torch.Tensor, + route_lanes_has_speed_limit: torch.Tensor, + polygons: torch.Tensor, + line_strings: torch.Tensor, + goal_pose: torch.Tensor, + ego_shape: torch.Tensor, + turn_indicators: torch.Tensor, + delay: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + prediction, turn_indicator_logit = super().forward( + sampled_trajectories, + ego_agent_past, + ego_current_state, + neighbor_agents_past, + static_objects, + lanes, + lanes_speed_limit, + lanes_has_speed_limit, + route_lanes, + route_lanes_speed_limit, + route_lanes_has_speed_limit, + polygons, + line_strings, + goal_pose, + ego_shape, + turn_indicators, + delay, + ) + keep_alive = 0.0 * ( + sampled_trajectories[:, 0, 0, 0] + ego_current_state[:, 0] + delay[:, 0] + ) + return prediction + keep_alive.view(-1, 1, 1, 1), turn_indicator_logit + + def build_dummy_inputs() -> TensorDict: inputs = {} inputs["sampled_trajectories"] = torch.ones( @@ -345,6 +437,13 @@ def load_model(config_json_path: str, ckpt_path: str, use_ema: bool) -> Diffusio def build_wrappers(model: Diffusion_Planner) -> ModelWrappers: + if isinstance(model.decoder, PlanTFDecoder): + return ModelWrappers( + full=PlanTFFullONNXWrapper(model).eval(), + encoder=EncoderONNXWrapper(model).eval(), + decoder=PlanTFDecoderONNXWrapper(model).eval(), + turn_indicator=None, + ) return ModelWrappers( full=FullONNXWrapper(model).eval(), encoder=EncoderONNXWrapper(model).eval(), @@ -395,6 +494,39 @@ def build_export_specs( ] +def build_plantf_export_specs( + wrappers: ModelWrappers, + inputs: TensorDict, + decoder_inputs: TensorDict, + full_onnx_path: Path, + encoder_onnx_path: Path, + decoder_onnx_path: Path, +) -> list[ExportSpec]: + return [ + ExportSpec( + wrapper=wrappers.full, + inputs=inputs, + input_names=FULL_INPUT_NAMES, + output_names=FULL_OUTPUT_NAMES, + output_path=full_onnx_path, + ), + ExportSpec( + wrapper=wrappers.encoder, + inputs=inputs, + input_names=ENCODER_INPUT_NAMES, + output_names=ENCODER_OUTPUT_NAMES, + output_path=encoder_onnx_path, + ), + ExportSpec( + wrapper=wrappers.decoder, + inputs=decoder_inputs, + input_names=PLANTF_DECODER_INPUT_NAMES, + output_names=PLANTF_DECODER_OUTPUT_NAMES, + output_path=decoder_onnx_path, + ), + ] + + def build_dynamic_axes( input_names: list[str], output_names: list[str] ) -> dict[str, dict[int, str]]: @@ -490,26 +622,37 @@ def export_model_to_onnx( with torch.no_grad(): encoding = wrappers.encoder(*(export_inputs[name] for name in ENCODER_INPUT_NAMES)) - decoder_inputs = build_decoder_inputs(export_inputs, encoding) - with torch.no_grad(): - final_x0 = wrappers.decoder( - decoder_inputs["encoding"], - decoder_inputs["sampled_trajectories"], - decoder_inputs["diffusion_time"], - decoder_inputs["neighbor_agents_past"], + if wrappers.turn_indicator is None: # planTF head: no denoising loop + decoder_inputs = build_decoder_inputs(export_inputs, encoding) + export_specs = build_plantf_export_specs( + wrappers, + export_inputs, + decoder_inputs, + full_onnx_path, + encoder_onnx_path, + decoder_onnx_path, + ) + else: + decoder_inputs = build_decoder_inputs(export_inputs, encoding) + with torch.no_grad(): + final_x0 = wrappers.decoder( + decoder_inputs["encoding"], + decoder_inputs["sampled_trajectories"], + decoder_inputs["diffusion_time"], + decoder_inputs["neighbor_agents_past"], + ) + turn_indicator_inputs = build_turn_indicator_inputs(encoding, final_x0) + + export_specs = build_export_specs( + wrappers, + export_inputs, + decoder_inputs, + turn_indicator_inputs, + full_onnx_path, + encoder_onnx_path, + decoder_onnx_path, + turn_indicator_onnx_path, ) - turn_indicator_inputs = build_turn_indicator_inputs(encoding, final_x0) - - export_specs = build_export_specs( - wrappers, - export_inputs, - decoder_inputs, - turn_indicator_inputs, - full_onnx_path, - encoder_onnx_path, - decoder_onnx_path, - turn_indicator_onnx_path, - ) for spec in export_specs: export_spec(spec, use_simplify, opset_version, external_data) diff --git a/diffusion_planner/diffusion_planner/validate_model.py b/diffusion_planner/diffusion_planner/validate_model.py index 389dacc92..d2f4c32a0 100644 --- a/diffusion_planner/diffusion_planner/validate_model.py +++ b/diffusion_planner/diffusion_planner/validate_model.py @@ -27,6 +27,7 @@ ) from diffusion_planner.train_epoch import heading_to_cos_sin from diffusion_planner.utils import ddp +from diffusion_planner.utils.normalizer import StateNormalizer from planner_metrics.temporal_stability import ( compute_curvature_rate_batch, compute_mean_abs_jerk_batch, @@ -46,6 +47,160 @@ class _PreparedValidationBatch: turn_indicator_seq: torch.Tensor +def compute_plantf_mode_metrics( + trajectory: torch.Tensor, + probability: torch.Tensor, + ego_future: torch.Tensor, + state_normalizer: StateNormalizer, + miss_threshold_m: float = 2.0, +) -> dict[str, torch.Tensor]: + """Compute per-sample metrics for a PlAnTF multi-modal trajectory head. + + ``trajectory`` is in the normalized model space while ``ego_future`` is in + ego-centric metres. The returned tensors all have shape ``[B]`` so they can + be concatenated and averaged across validation batches/ranks. + """ + trajectory_world = trajectory * state_normalizer.std[0].to( + trajectory.device + ) + state_normalizer.mean[0].to(trajectory.device) + position_error = torch.linalg.vector_norm( + trajectory_world[..., :2] - ego_future[:, None, :, :2], dim=-1 + ) # [B, K, T] + ade = position_error.mean(dim=-1) + fde = position_error[..., -1] + + oracle_mode = ade.argmin(dim=-1) + selected_mode = probability.argmax(dim=-1) + batch_index = torch.arange(trajectory.shape[0], device=trajectory.device) + top1_ade = ade[batch_index, selected_mode] + top1_fde = fde[batch_index, selected_mode] + + metrics = { + "plantf_top1_ade": top1_ade, + "plantf_top1_fde": top1_fde, + } + + # These diagnostics answer multimodal questions only. With one mode they + # are exact constants/redundancies (min=top1, oracle gap=0, accuracy=1, + # entropy=0, usage_0=1), so do not emit them into W&B/TSV. + if trajectory.shape[1] == 1: + metrics["plantf_miss_rate_2m"] = (top1_fde > miss_threshold_m).float() + return metrics + + min_ade = ade.min(dim=-1).values + min_fde = fde.min(dim=-1).values + mode_prob = probability.softmax(dim=-1) + mode_entropy = -(mode_prob * mode_prob.clamp_min(1e-12).log()).sum(dim=-1) + metrics.update( + { + "plantf_min_ade_k": min_ade, + "plantf_min_fde_k": min_fde, + # Preserve the original multimodal definition: a sample only + # misses when *none* of its candidate trajectories reaches the + # endpoint threshold. + "plantf_miss_rate_2m": (min_fde > miss_threshold_m).float(), + "plantf_top1_oracle_ade_gap": top1_ade - min_ade, + "plantf_mode_accuracy": (selected_mode == oracle_mode).float(), + "plantf_mode_entropy": mode_entropy, + } + ) + for mode in range(trajectory.shape[1]): + # selected: which mode argmax(pi) picks. oracle: which mode is actually + # best (ADE min). A concentrated oracle distribution means the data is + # effectively uni/bi-modal (a small num_modes suffices); a spread oracle + # distribution means it is genuinely multi-modal and collapsing to one + # mode would average the trajectories away. See + # docs/plantf_dead_mode_improvement.md. + metrics[f"plantf_mode_usage_{mode}"] = (selected_mode == mode).float() + metrics[f"plantf_oracle_usage_{mode}"] = (oracle_mode == mode).float() + return metrics + + +_PROGRESS_BINS = (("stop", 0.0, 2.0), ("slow", 2.0, 10.0), ("move", 10.0, float("inf"))) +_PROGRESS_DT = 0.1 + + +def compute_trajectory_progress_metrics( + prediction_ego: torch.Tensor, + ego_future: torch.Tensor, + *, + include_redundant_summary_metrics: bool = True, +) -> dict[str, torch.Tensor]: + """Per-sample forward-progress / smoothness metrics for the ego prediction. + + Motivated by the planTF "dango" collapse (docs/plantf_dead_mode_improvement.md): + a near-stationary oscillating prediction can hide inside an averaged + position loss, especially on stop-heavy data. These metrics expose it + directly, and each one is additionally reported binned by the GT endpoint + distance (``stop`` < 2m, ``slow`` 2-10m, ``move`` >= 10m) using the + ``_available`` masking convention, so stopped scenes cannot mask a moving + scene failure. + + Args: + prediction_ego: [B, T, 4] denormalized ego prediction (ego-centric metres). + ego_future: [B, T, 4] ground-truth ego future (ego-centric metres). + + Returns: + Dict of per-sample [B] tensors, keys prefixed with ``traj_``. + """ + pred_xy = prediction_ego[..., :2] + gt_xy = ego_future[..., :2] + + pred_end = pred_xy[:, -1].norm(dim=-1) # [B] endpoint distance from ego + gt_end = gt_xy[:, -1].norm(dim=-1) + pred_len = (pred_xy[:, 1:] - pred_xy[:, :-1]).norm(dim=-1).sum(dim=-1) # [B] + gt_len = (gt_xy[:, 1:] - gt_xy[:, :-1]).norm(dim=-1).sum(dim=-1) + fde = (pred_xy[:, -1] - gt_xy[:, -1]).norm(dim=-1) + + # Mean second difference of the predicted positions: a straight + # constant-speed trajectory scores 0; a jagged / oscillating one is large. + second_diff = pred_xy[:, 2:] - 2.0 * pred_xy[:, 1:-1] + pred_xy[:, :-2] + roughness = second_diff.norm(dim=-1).mean(dim=-1) + + pred_speed = (pred_xy[:, 1:] - pred_xy[:, :-1]).norm(dim=-1) / _PROGRESS_DT # [B, T-1] + gt_speed = (gt_xy[:, 1:] - gt_xy[:, :-1]).norm(dim=-1) / _PROGRESS_DT + speed_mae = (pred_speed - gt_speed).abs().mean(dim=-1) + + # Ratios are clamped on the GT side so stopped scenes (GT ~ 0m) do not blow up. + metrics = { + "traj_endpoint_progress_ratio": pred_end / gt_end.clamp_min(1.0), + "traj_length_ratio": pred_len / gt_len.clamp_min(1.0), + "traj_second_diff_m": roughness, + "traj_speed_mae_mps": speed_mae, + } + # Keep the legacy DP summaries by default. PlantF already reports the + # same global FDE as ``plantf_top1_fde``; its raw endpoint values only + # support the ratios above and add no independent validation signal. + if include_redundant_summary_metrics: + metrics.update( + { + "traj_endpoint_m": pred_end, + "traj_gt_endpoint_m": gt_end, + "traj_fde_m": fde, + } + ) + binned_metrics = { + "endpoint_progress_ratio": metrics["traj_endpoint_progress_ratio"], + "length_ratio": metrics["traj_length_ratio"], + "fde_m": fde, + "second_diff_m": metrics["traj_second_diff_m"], + "speed_mae_mps": metrics["traj_speed_mae_mps"], + } + binned = ( + "endpoint_progress_ratio", + "length_ratio", + "fde_m", + "second_diff_m", + "speed_mae_mps", + ) + for bin_name, lo, hi in _PROGRESS_BINS: + mask = ((gt_end >= lo) & (gt_end < hi)).float() + for base in binned: + metrics[f"traj_{base}_{bin_name}"] = binned_metrics[base] + metrics[f"traj_{base}_{bin_name}_available"] = mask + return metrics + + def _prepare_validation_inputs(inputs, args, device, delay=0) -> _PreparedValidationBatch: inputs = {key: value.to(device) for key, value in inputs.items()} batch_size = inputs["ego_current_state"].shape[0] @@ -334,6 +489,27 @@ def validate_model(model, val_loader, args, return_pred=False) -> tuple[float, f all_gt[:, 1:][neighbor_mask] = 0.0 prediction = outputs["prediction"] + if "trajectory" in outputs and "probability" in outputs: + mode_metrics = compute_plantf_mode_metrics( + outputs["trajectory"], + outputs["probability"], + ego_future, + args.state_normalizer, + ) + for key, val in mode_metrics.items(): + total_result_dict[key].append(val.cpu()) + progress_metrics = compute_trajectory_progress_metrics( + prediction[:, 0], + ego_future, + # Preserve historical DP validation fields. For PlantF, FDE is + # already emitted as plantf_top1_fde and the raw endpoint values + # have no independent interpretation in the regular log. + include_redundant_summary_metrics=not ( + "trajectory" in outputs and "probability" in outputs + ), + ) + for key, val in progress_metrics.items(): + total_result_dict[key].append(val.cpu()) turn_indicator_logit = outputs["turn_indicator_logit"] turn_indicator = turn_indicator_logit.argmax(dim=-1) turn_indicator_gt = make_turn_indicator_gt(turn_indicator_seq) @@ -585,6 +761,35 @@ def aggregate_valid_metrics(valid_dict, device): epdms_means[f"{metric}_coverage"] = local_cnt / max(local_total, 1) epdms_means[metric] = local_sum / local_cnt if local_cnt > 0 else float("nan") + plantf_means = {} + for key, val in valid_dict.items(): + if not key.startswith("plantf_"): + continue + local_sum = ddp.all_reduce_sum(val.sum().item(), device) + local_cnt = ddp.all_reduce_sum(val.numel(), device) + plantf_means[key.removeprefix("plantf_")] = local_sum / max(local_cnt, 1) + + # Forward-progress metrics; same `_available` masking convention as epdms + # so GT-progress-binned means only average the samples inside each bin. + traj_means = {} + for key, val in valid_dict.items(): + if not key.startswith("traj_"): + continue + metric = key.removeprefix("traj_") + tensor = val.float() + if metric.endswith("_available"): + continue + available = valid_dict.get(f"{key}_available") + if available is None: + local_sum = ddp.all_reduce_sum(tensor.sum().item(), device) + local_cnt = ddp.all_reduce_sum(tensor.numel(), device) + traj_means[metric] = local_sum / max(local_cnt, 1) + continue + mask = available.float() > 0.5 + local_sum = ddp.all_reduce_sum(tensor[mask].sum().item() if mask.any() else 0.0, device) + local_cnt = ddp.all_reduce_sum(mask.sum().item(), device) + traj_means[metric] = local_sum / local_cnt if local_cnt > 0 else float("nan") + return { "avg_loss_ego": loss_ego_sum / max(samples_ego, 1), "avg_loss_neighbor": loss_nei_sum / max(samples_nei, 1), @@ -595,4 +800,6 @@ def aggregate_valid_metrics(valid_dict, device): "turn_indicator_change_total": int(turn_change_total), "ego_means": ego_means, "epdms_means": epdms_means, + "plantf_means": plantf_means, + "traj_means": traj_means, } diff --git a/diffusion_planner/tests/test_plantf_decoder.py b/diffusion_planner/tests/test_plantf_decoder.py new file mode 100644 index 000000000..770505a5c --- /dev/null +++ b/diffusion_planner/tests/test_plantf_decoder.py @@ -0,0 +1,873 @@ +"""PlanTF decoder head: forward contract, winner-takes-all loss, and wiring.""" + +from argparse import Namespace + +import pytest +import torch +from diffusion_planner.dimensions import ( + MAX_NUM_NEIGHBORS, + OUTPUT_T, + POSE_DIM, + TURN_INDICATOR_OUTPUT_DIM, +) +from diffusion_planner.model.diffusion_planner import Diffusion_Planner +from diffusion_planner.model.module.plantf_decoder import ( + PlanTFCrossAttnHead, + PlanTFGRUHead, + PlanTFTrajectoryHead, + bezier_basis, + compute_plantf_training_loss, +) +from diffusion_planner.utils.normalizer import ObservationNormalizer, StateNormalizer +from diffusion_planner.utils.onnx_export import ( + FULL_INPUT_NAMES, + FULL_OUTPUT_NAMES, + FullONNXWrapper, + PlanTFFullONNXWrapper, + build_dummy_inputs, + build_wrappers, + export_onnx, + onnx_export_backends, +) +from diffusion_planner.validate_model import compute_plantf_mode_metrics + +NUM_MODES = 3 +HIDDEN_DIM = 32 + + +def _state_normalizer(num_agents): + mean = torch.zeros(num_agents, 1, POSE_DIM) + mean[:, :, 0] = 1.0 # non-trivial mean so denormalization is exercised + std = torch.full((num_agents, 1, POSE_DIM), 2.0) + return StateNormalizer(mean, std) + + +def _config(): + return Namespace( + # encoder + hidden_dim=HIDDEN_DIM, + use_ego_history=True, + ego_history_dropout_rate=0.0, + use_turn_indicators=True, + agent_num=MAX_NUM_NEIGHBORS, + static_objects_num=5, + static_objects_state_dim=10, + lane_num=140, + lane_len=20, + route_num=25, + route_len=20, + polygon_num=10, + polygon_len=40, + line_string_num=60, + line_string_len=20, + time_len=31, + encoder_drop_path_rate=0.0, + encoder_mixer_depth=2, + encoder_fusion_depth=2, + num_heads=2, + # decoder + decoder_type="plantf", + num_modes=NUM_MODES, + predicted_neighbor_num=MAX_NUM_NEIGHBORS, + future_len=OUTPUT_T, + use_velocity_representation=False, + state_normalizer=_state_normalizer(1 + MAX_NUM_NEIGHBORS), + ) + + +def _inputs(batch_size=2): + torch.manual_seed(0) + return { + k: v.repeat(batch_size, *([1] * (v.dim() - 1))) for k, v in build_dummy_inputs().items() + } + + +def _loss_args(config): + return Namespace( + decoder_type="plantf", + state_normalizer=config.state_normalizer, + observation_normalizer=None, + ego_prediction_horizon=OUTPUT_T, + coeff_velocity=1.0, + coeff_timestep=[1.0, 1.0, 1.0, 1.0], + coeff_position_lat_loss=1.0, + coeff_position_lon_loss=1.0, + coeff_heading_l2_loss=1.0, + coeff_road_border_loss=0.0, + coeff_neighbor_collision_loss=0.0, + road_border_margin=0.25, + road_border_n_interp=2, + neighbor_collision_margin_vehicle=0.25, + neighbor_collision_margin_pedestrian=1.0, + neighbor_collision_margin_bicycle=0.5, + ) + + +def test_trajectory_head_shapes(): + torch.manual_seed(0) + head = PlanTFTrajectoryHead( + embed_dim=HIDDEN_DIM, num_modes=NUM_MODES, future_steps=OUTPUT_T, out_channels=4 + ) + loc, pi = head(torch.randn(2, HIDDEN_DIM)) + assert loc.shape == (2, NUM_MODES, OUTPUT_T, 4) + assert pi.shape == (2, NUM_MODES) + + +def test_inference_prediction_contract(): + torch.manual_seed(0) + model = Diffusion_Planner(_config()).eval() + inputs = _inputs() + + with torch.no_grad(): + _, outputs = model(inputs) + + B = 2 + P = 1 + MAX_NUM_NEIGHBORS + assert outputs["prediction"].shape == (B, P, OUTPUT_T, POSE_DIM) + assert outputs["trajectory"].shape == (B, NUM_MODES, OUTPUT_T, POSE_DIM) + assert outputs["probability"].shape == (B, NUM_MODES) + assert outputs["turn_indicator_logit"].shape == (B, TURN_INDICATOR_OUTPUT_DIM) + assert outputs["independent_turn_indicator_logit"].shape == (B, TURN_INDICATOR_OUTPUT_DIM) + assert torch.isfinite(outputs["prediction"]).all() + + # ego row of "prediction" must be the argmax mode, denormalized + norm = model.decoder._state_normalizer + best = outputs["probability"].argmax(dim=-1) + expected = outputs["trajectory"][torch.arange(B), best] * norm.std[0] + norm.mean[0] + assert torch.allclose(outputs["prediction"][:, 0], expected) + + +def test_relative_xy_is_anchored_at_each_agent_current_position(): + """Original-planTF relative xy output must reconstruct absolute DP xy.""" + config = _config() + config.plantf_relative_xy = True + model = Diffusion_Planner(config).eval() + inputs = _inputs(batch_size=1) + Pn = MAX_NUM_NEIGHBORS + # Use distinct, non-origin coordinates so an ego-only anchor cannot pass. + inputs["ego_current_state"][0, :2] = torch.tensor([3.0, -2.0]) + inputs["neighbor_agents_past"][0, :Pn, -1, :2] = torch.stack( + [torch.tensor([10.0 + i, -5.0 - i]) for i in range(Pn)] + ) + trajectory = torch.zeros(1, NUM_MODES, OUTPUT_T, POSE_DIM) + neighbors = torch.zeros(1, Pn, OUTPUT_T, POSE_DIM) + + anchored_ego, anchored_neighbors = model.decoder._anchor_relative_xy( + trajectory, neighbors, inputs + ) + norm = config.state_normalizer + expected_ego = (inputs["ego_current_state"][0, :2] - norm.mean[0, 0, :2]) / norm.std[0, 0, :2] + expected_neighbors = ( + inputs["neighbor_agents_past"][0, :Pn, -1, :2] - norm.mean[1:, 0, :2] + ) / norm.std[1:, 0, :2] + assert torch.allclose(anchored_ego[0, :, :, :2], expected_ego.view(1, 1, 2)) + assert torch.allclose(anchored_neighbors[0, :, :, :2], expected_neighbors[:, None, :]) + + +def test_relative_xy_runs_forward_and_training_loss(): + """Relative-xy mode must work through the normal model/loss integration.""" + torch.manual_seed(7) + config = _config() + config.plantf_relative_xy = True + model = Diffusion_Planner(config) + inputs = _inputs(batch_size=2) + # Current positions must be non-zero for ego and neighbors, otherwise an + # accidental no-op anchor would not be detected by this integration test. + inputs["ego_current_state"][:, :2] = torch.tensor([[2.0, -1.0], [4.0, 3.0]]) + inputs["neighbor_agents_past"][:, :, -1, 0] += 8.0 + + model.eval() + with torch.no_grad(): + _, outputs = model(inputs) + assert torch.isfinite(outputs["prediction"]).all() + + model.train() + B, Pn = 2, MAX_NUM_NEIGHBORS + ego_future = torch.randn(B, OUTPUT_T, POSE_DIM) + neighbors_future = torch.randn(B, Pn, OUTPUT_T, POSE_DIM) + mask = torch.zeros(B, Pn, OUTPUT_T, dtype=torch.bool) + loss = compute_plantf_training_loss( + model, inputs, (ego_future, neighbors_future, mask), _loss_args(config) + ) + total = loss["ego_planning_loss"] + loss["neighbor_prediction_loss"] + loss["mode_cls_loss"] + total.backward() + assert torch.isfinite(total) + + +def test_relative_xy_rejects_split_decoder_without_neighbor_current_state(): + """Split graph requires the shared DP decoder anchor inputs.""" + config = _config() + config.plantf_relative_xy = True + model = Diffusion_Planner(config).eval() + with pytest.raises(ValueError, match="sampled_trajectories"): + model.decoder.forward_deploy(torch.zeros(1, 1 + MAX_NUM_NEIGHBORS, HIDDEN_DIM)) + + +def test_relative_xy_split_decoder_matches_full_decoder_contract(): + """Relative mode preserves DP's four-input split-decoder ABI and output.""" + torch.manual_seed(11) + config = _config() + config.plantf_relative_xy = True + model = Diffusion_Planner(config).eval() + inputs = _inputs(batch_size=1) + inputs["ego_current_state"][0, :2] = torch.tensor([2.5, -1.5]) + inputs["neighbor_agents_past"][0, :, -1, :2] += 3.0 + sampled = torch.zeros(1, 1 + MAX_NUM_NEIGHBORS, OUTPUT_T + 1, POSE_DIM) + sampled[:, 0, 0] = inputs["ego_current_state"][:, :4] + encoding = model.encoder(inputs) + with torch.no_grad(): + full = model.decoder(encoding, inputs)["prediction"] + split, _, _ = model.decoder.forward_deploy( + encoding, + sampled, + torch.ones(1, 1 + MAX_NUM_NEIGHBORS, OUTPUT_T + 1, 1), + inputs["neighbor_agents_past"], + ) + assert torch.allclose(split, full) + + +def test_training_loss_keys_and_backward(): + torch.manual_seed(0) + config = _config() + model = Diffusion_Planner(config).train() + inputs = _inputs() + + B, Pn = 2, MAX_NUM_NEIGHBORS + ego_future = torch.randn(B, OUTPUT_T, POSE_DIM) + neighbors_future = torch.randn(B, Pn, OUTPUT_T, POSE_DIM) + mask = torch.zeros(B, Pn, OUTPUT_T, dtype=torch.bool) + mask[:, Pn // 2 :] = True # half of the neighbors invalid + + loss = compute_plantf_training_loss( + model, inputs, (ego_future, neighbors_future, mask), _loss_args(config) + ) + + for key in ( + "ego_planning_loss", + "neighbor_prediction_loss", + "mode_cls_loss", + "turn_indicator_loss", + "independent_turn_indicator_loss", + "road_border_loss", + "neighbor_collision_loss", + "turn_indicator_accuracy", + ): + assert key in loss, key + assert torch.isfinite(loss[key]), key + + total = loss["ego_planning_loss"] + loss["neighbor_prediction_loss"] + loss["mode_cls_loss"] + total.backward() + grads = [p.grad for p in model.parameters() if p.grad is not None] + assert len(grads) > 0 + assert all(torch.isfinite(g).all() for g in grads) + + +def test_winner_takes_all_picks_closest_mode(): + """With a stubbed decoder output whose mode 1 equals GT, the WTA loss must + select mode 1: near-zero ego loss and CE targeting mode 1.""" + torch.manual_seed(0) + config = _config() + args = _loss_args(config) + + B, Pn, K = 2, MAX_NUM_NEIGHBORS, NUM_MODES + inputs = _inputs(B) + ego_future = torch.randn(B, OUTPUT_T, POSE_DIM) + neighbors_future = torch.zeros(B, Pn, OUTPUT_T, POSE_DIM) + mask = torch.ones(B, Pn, OUTPUT_T, dtype=torch.bool) # no valid neighbors + + ego_gt_norm = args.state_normalizer(torch.cat([ego_future[:, None], neighbors_future], dim=1))[ + :, 0 + ] + + trajectory = torch.randn(B, K, OUTPUT_T, POSE_DIM) + trajectory[:, 1] = ego_gt_norm # mode 1 is exact + + class StubModel: + def __call__(self, merged_inputs): + return None, { + "trajectory": trajectory, + "probability": torch.zeros(B, K, requires_grad=True), + "neighbor_prediction": torch.zeros(B, Pn, OUTPUT_T, POSE_DIM), + "turn_indicator_logit": torch.zeros(B, TURN_INDICATOR_OUTPUT_DIM), + "independent_turn_indicator_logit": torch.zeros(B, TURN_INDICATOR_OUTPUT_DIM), + } + + loss = compute_plantf_training_loss( + StubModel(), inputs, (ego_future, neighbors_future, mask), args + ) + + assert loss["ego_planning_loss"] < 1e-5 + # uniform logits over K modes -> CE == log(K) regardless of the target, + # so verify the target itself via a peaked-logit check instead + assert loss["neighbor_prediction_loss"] == 0.0 + + peaked = torch.full((B, K), -10.0) + peaked[:, 1] = 10.0 + + class PeakedStub(StubModel): + def __call__(self, merged_inputs): + _, out = super().__call__(merged_inputs) + out["probability"] = peaked + return None, out + + loss_peaked = compute_plantf_training_loss( + PeakedStub(), inputs, (ego_future, neighbors_future, mask), args + ) + assert loss_peaked["mode_cls_loss"] < 1e-5 + + +def test_forward_deploy_matches_eval_forward(): + """The ONNX deploy path must produce the same prediction / probability / + turn-indicator logit as the regular eval forward.""" + torch.manual_seed(0) + model = Diffusion_Planner(_config()).eval() + inputs = _inputs() + + with torch.no_grad(): + encoding = model.encoder(inputs) + outputs = model.decoder(encoding, inputs) + prediction, probability, turn_indicator_logit = model.decoder.forward_deploy(encoding) + + assert torch.allclose(prediction, outputs["prediction"]) + assert torch.allclose(probability, outputs["probability"]) + assert torch.allclose(turn_indicator_logit, outputs["turn_indicator_logit"]) + + +def test_full_onnx_wrapper_keeps_diffusion_only_inputs(tmp_path): + """The Autoware node feeds all 17 full-graph inputs by name (ORT and TensorRT + both reject names absent from the graph), but the planTF head ignores + sampled_trajectories / ego_current_state / delay and the legacy exporter + prunes unused graph inputs. The planTF full wrapper must anchor them without + changing the outputs.""" + onnx = pytest.importorskip("onnx") + torch.manual_seed(0) + model = Diffusion_Planner(_config()).eval() + wrappers = build_wrappers(model) + assert isinstance(wrappers.full, PlanTFFullONNXWrapper) + assert wrappers.turn_indicator is None + + inputs = build_dummy_inputs() + args = tuple(inputs[name] for name in FULL_INPUT_NAMES) + with torch.no_grad(): + anchored_prediction, anchored_logit = wrappers.full(*args) + plain_prediction, plain_logit = FullONNXWrapper(model).eval()(*args) + assert torch.equal(anchored_prediction, plain_prediction) + assert torch.equal(anchored_logit, plain_logit) + + full_onnx_path = tmp_path / "full.onnx" + with onnx_export_backends(), torch.no_grad(): + export_onnx( + wrappers.full, + inputs, + FULL_INPUT_NAMES, + FULL_OUTPUT_NAMES, + full_onnx_path, + use_simplify=False, + opset_version=17, + external_data=False, + ) + graph_inputs = {i.name for i in onnx.load(str(full_onnx_path)).graph.input} + assert graph_inputs == set(FULL_INPUT_NAMES) + + +def test_ego_loss_is_plain_smooth_l1_of_best_mode(): + """The ego loss must be exactly the smooth L1 between the WTA best mode and + the GT over all channels (original planTF loss) — no lat/lon/heading + decomposition, no velocity down-weighting, no timestep weighting, no + endpoint term. See docs/plantf_dead_mode_improvement.md.""" + torch.manual_seed(0) + config = _config() + model = Diffusion_Planner(config).eval() # eval: no dropout, deterministic + inputs = _inputs() + B, Pn = 2, MAX_NUM_NEIGHBORS + futures = ( + torch.randn(B, OUTPUT_T, POSE_DIM), + torch.randn(B, Pn, OUTPUT_T, POSE_DIM), + torch.zeros(B, Pn, OUTPUT_T, dtype=torch.bool), + ) + loss = compute_plantf_training_loss(model, inputs, futures, _loss_args(config)) + + # The DP-specific loss terms must be gone. + assert "endpoint_fde_loss" not in loss + # Reproduce the ego smooth-L1 independently from the model's own outputs. + norm = config.state_normalizer + ego_gt = norm(futures[0][:, None])[:, 0] # [B, T, 4] normalized + with torch.no_grad(): + _, outputs = model( + {**inputs, "gt_trajectories": torch.zeros(B, 1 + Pn, OUTPUT_T + 1, POSE_DIM)} + ) + traj = outputs["trajectory"] # [B, K, T, 4] + ego_std_xy = norm.std[0][..., :2] + ade = ((traj[..., :2] - ego_gt[:, None, :, :2]) * ego_std_xy).norm(dim=-1).mean(-1) + best = ade.argmin(dim=-1) + best_traj = traj[torch.arange(B), best] + expected = torch.nn.functional.smooth_l1_loss(best_traj, ego_gt) + assert torch.allclose(loss["ego_planning_loss"], expected, atol=1e-5) + + +def test_trajectory_progress_metrics_expose_stationary_collapse(): + """A near-stationary oscillating prediction must score ~0 progress and a + large second-difference, with the GT-progress bins masking correctly.""" + from diffusion_planner.validate_model import compute_trajectory_progress_metrics + + T = OUTPUT_T + t = torch.arange(T, dtype=torch.float32) + gt = torch.zeros(2, T, POSE_DIM) + gt[0, :, 0] = (t + 1) * 0.2 # straight 16m advance -> "move" bin + gt[1, :, 0] = (t + 1) * 0.01 # crawl 0.8m -> "stop" bin + pred = torch.zeros(2, T, POSE_DIM) + pred[0, :, 0] = 0.3 * (t % 2) # dango: oscillates within 0.3m + + m = compute_trajectory_progress_metrics(pred, gt) + assert float(m["traj_gt_endpoint_m"][0]) == pytest.approx(16.0) + assert float(m["traj_endpoint_progress_ratio"][0]) < 0.05 + assert float(m["traj_length_ratio"][0]) > 1.0 # long path, no progress + assert float(m["traj_second_diff_m"][0]) > 0.5 # jagged + # Bin masks: sample 0 is "move", sample 1 is "stop". + assert m["traj_fde_m_move_available"].tolist() == [1.0, 0.0] + assert m["traj_fde_m_stop_available"].tolist() == [0.0, 1.0] + assert m["traj_fde_m_slow_available"].tolist() == [0.0, 0.0] + + +def test_ego_state_injection_anchors_prediction(): + """The current ego motion state must actually change the head output when + enabled (anchoring the absolute-waypoint regression to the current motion, + the way the diffusion decoder's pinned current state does), and be a no-op + when disabled. See docs/plantf_dead_mode_improvement.md.""" + torch.manual_seed(0) + + def _with_motion(base, value): + out = {k: (v.clone() if torch.is_tensor(v) else v) for k, v in base.items()} + out["ego_current_state"] = out["ego_current_state"].clone() + out["ego_current_state"][:, 4:10] = value # vx, vy, ax, ay, steering, yaw_rate + return out + + def _detrivialize(model): + # The loc/pi output layers are zero-initialized, so before training the + # head maps any input to 0. Perturb them to mimic a trained head, so the + # effect of the ego-state input is observable. + with torch.no_grad(): + model.decoder.trajectory_head.loc[-1].weight.normal_(0, 0.1) + model.decoder.trajectory_head.pi[-1].weight.normal_(0, 0.1) + + # Enabled: different current motion -> different prediction. + config = _config() + config.plantf_use_ego_state_in_head = True + model = Diffusion_Planner(config).eval() + _detrivialize(model) + base = _inputs() + with torch.no_grad(): + _, fast = model(_with_motion(base, 3.0)) + _, slow = model(_with_motion(base, 0.0)) + assert not torch.allclose(fast["trajectory"], slow["trajectory"]) + + # Disabled: the motion state is ignored. + config_off = _config() + config_off.plantf_use_ego_state_in_head = False + model_off = Diffusion_Planner(config_off).eval() + _detrivialize(model_off) + with torch.no_grad(): + _, fast_off = model_off(_with_motion(base, 3.0)) + _, slow_off = model_off(_with_motion(base, 0.0)) + assert torch.allclose(fast_off["trajectory"], slow_off["trajectory"]) + + +def test_a2_ego_state_token_and_c1_input_delta_preserve_contract(): + """A2 (ego token = current-state embedding) and C1 (xy-delta history inputs), + independently and together, must keep the inference contract: + denormalized [B, 1+Pn, T, 4], finite. See + docs/plantf_original_comparison_and_roadmap.md.""" + for flags in ( + {"plantf_ego_state_token": True}, + {"plantf_input_delta": True}, + {"plantf_ego_state_token": True, "plantf_input_delta": True}, + ): + torch.manual_seed(0) + config = _config() + config.plantf_ego_state_dropout = 0.75 + for key, value in flags.items(): + setattr(config, key, value) + model = Diffusion_Planner(config).eval() + with torch.no_grad(): + _, outputs = model(_inputs()) + assert outputs["prediction"].shape == ( + 2, + 1 + MAX_NUM_NEIGHBORS, + OUTPUT_T, + POSE_DIM, + ), flags + assert torch.isfinite(outputs["prediction"]).all(), flags + + +def test_a2_c1_training_backward(): + """A2 + C1 must train (finite gradients) with the planTF loss.""" + torch.manual_seed(0) + config = _config() + config.plantf_ego_state_token = True + config.plantf_input_delta = True + config.plantf_ego_state_dropout = 0.75 + model = Diffusion_Planner(config).train() + B, Pn = 2, MAX_NUM_NEIGHBORS + futures = ( + torch.randn(B, OUTPUT_T, POSE_DIM), + torch.randn(B, Pn, OUTPUT_T, POSE_DIM), + torch.zeros(B, Pn, OUTPUT_T, dtype=torch.bool), + ) + loss = compute_plantf_training_loss(model, _inputs(), futures, _loss_args(config)) + total = loss["ego_planning_loss"] + loss["neighbor_prediction_loss"] + loss["mode_cls_loss"] + total.backward() + grads = [p.grad for p in model.parameters() if p.grad is not None] + assert len(grads) > 0 + assert all(torch.isfinite(g).all() for g in grads) + + +def test_output_heads_are_zero_initialized(): + """Every mode must start at the normalized-space mean (zero output), not + Xavier noise: under winner-takes-all training rarely-winning modes keep + their initialization, and argmax(pi) can select them at inference. See + docs/plantf_dead_mode_improvement.md.""" + torch.manual_seed(0) + decoder = Diffusion_Planner(_config()).decoder + encoding = torch.randn(2, 1 + MAX_NUM_NEIGHBORS + 5, HIDDEN_DIM) + + trajectory, probability, neighbor_prediction = decoder._decode(encoding) + assert torch.all(trajectory == 0.0) + assert torch.all(probability == 0.0) + assert torch.all(neighbor_prediction == 0.0) + + +def test_velocity_representation_integrates_to_continuous_waypoints(): + """With use_velocity_representation the heads regress displacement and + _decode integrates it: consecutive waypoints must differ by the (normalized) + per-step displacement, giving the temporal continuity a per-step absolute + regression lacks. See docs/plantf_dead_mode_improvement.md.""" + from diffusion_planner.loss import velocity_to_waypoints + + torch.manual_seed(0) + config = _config() + config.use_velocity_representation = True + decoder = Diffusion_Planner(config).decoder + encoding = torch.randn(2, 1 + MAX_NUM_NEIGHBORS + 5, HIDDEN_DIM) + + # Raw head displacement, before integration. + raw_vel, _ = decoder.trajectory_head(encoding[:, 0]) # [B, K, T, 4] + trajectory, _, neighbor = decoder._decode(encoding) + + norm = decoder._state_normalizer + # De-normalize the integrated trajectory back to metres and confirm it is the + # cumulative sum of the raw displacement (heading passed through). + ego_world = trajectory * norm.std[0] + norm.mean[0] + assert torch.allclose(ego_world, velocity_to_waypoints(raw_vel), atol=1e-4) + # xy strictly integrates (waypoint[t]-waypoint[t-1] == displacement[t]), so a + # bounded displacement can never produce the comb jitter of absolute regression. + step = ego_world[:, :, 1:, :2] - ego_world[:, :, :-1, :2] + assert torch.allclose(step, raw_vel[:, :, 1:, :2], atol=1e-4) + assert trajectory.shape == (2, NUM_MODES, OUTPUT_T, 4) + assert neighbor.shape == (2, MAX_NUM_NEIGHBORS, OUTPUT_T, 4) + + +def test_velocity_representation_inference_contract(): + """The velocity head must still honor the Decoder inference contract: + a denormalized [B, 1+Pn, T, 4] prediction.""" + torch.manual_seed(0) + config = _config() + config.use_velocity_representation = True + model = Diffusion_Planner(config).eval() + with torch.no_grad(): + _, outputs = model(_inputs()) + assert outputs["prediction"].shape == (2, 1 + MAX_NUM_NEIGHBORS, OUTPUT_T, POSE_DIM) + assert torch.isfinite(outputs["prediction"]).all() + + +def test_velocity_representation_training_loss_backward(): + torch.manual_seed(0) + config = _config() + config.use_velocity_representation = True + model = Diffusion_Planner(config).train() + inputs = _inputs() + B, Pn = 2, MAX_NUM_NEIGHBORS + futures = ( + torch.randn(B, OUTPUT_T, POSE_DIM), + torch.randn(B, Pn, OUTPUT_T, POSE_DIM), + torch.zeros(B, Pn, OUTPUT_T, dtype=torch.bool), + ) + loss = compute_plantf_training_loss(model, inputs, futures, _loss_args(config)) + total = loss["ego_planning_loss"] + loss["neighbor_prediction_loss"] + loss["mode_cls_loss"] + total.backward() + grads = [p.grad for p in model.parameters() if p.grad is not None] + assert len(grads) > 0 + assert all(torch.isfinite(g).all() for g in grads) + + +def test_velocity_representation_rejected_removed(): + """Sanity: constructing a velocity-representation planTF decoder no longer + raises (the NotImplementedError guard was removed).""" + config = _config() + config.use_velocity_representation = True + Diffusion_Planner(config) # must not raise + + +def test_mode_metrics_separate_candidate_quality_from_mode_selection(): + """A good oracle candidate with the wrong selected mode must have low + minADE but worse top-1 ADE and zero mode-selection accuracy.""" + normalizer = StateNormalizer( + mean=torch.zeros(1, 1, POSE_DIM), + std=torch.full((1, 1, POSE_DIM), 2.0), + ) + ego_future = torch.zeros(2, OUTPUT_T, POSE_DIM) + trajectory_world = torch.zeros(2, NUM_MODES, OUTPUT_T, POSE_DIM) + + # sample 0: mode 1 is exact but the logits select mode 0, one metre away + trajectory_world[0, 0, :, 0] = 1.0 + # sample 1: mode 2 is exact and selected + trajectory_world[1, 0, :, 0] = 3.0 + trajectory_world[1, 1, :, 0] = 2.0 + probability = torch.tensor([[10.0, 0.0, 0.0], [0.0, 0.0, 10.0]]) + trajectory_norm = normalizer(trajectory_world) + + metrics = compute_plantf_mode_metrics(trajectory_norm, probability, ego_future, normalizer) + + assert torch.allclose(metrics["plantf_min_ade_k"], torch.zeros(2)) + assert torch.allclose(metrics["plantf_min_fde_k"], torch.zeros(2)) + assert torch.allclose(metrics["plantf_top1_ade"], torch.tensor([1.0, 0.0])) + assert torch.allclose(metrics["plantf_top1_fde"], torch.tensor([1.0, 0.0])) + assert torch.equal(metrics["plantf_mode_accuracy"], torch.tensor([0.0, 1.0])) + assert torch.equal(metrics["plantf_mode_usage_0"], torch.tensor([1.0, 0.0])) + assert torch.equal(metrics["plantf_mode_usage_2"], torch.tensor([0.0, 1.0])) + assert torch.equal(metrics["plantf_miss_rate_2m"], torch.zeros(2)) + # oracle (best-ADE) mode is 1 for sample 0 and 2 for sample 1 — distinct from + # the selected mode, which is the whole point of tracking it separately. + assert torch.equal(metrics["plantf_oracle_usage_1"], torch.tensor([1.0, 0.0])) + assert torch.equal(metrics["plantf_oracle_usage_2"], torch.tensor([0.0, 1.0])) + + +# --------------------------------------------------------------------------- +# Combinable ablation toggles (docs/plantf_head_development_notes.md §9) +# --------------------------------------------------------------------------- + + +def _config_with(**kw): + c = _config() + for k, v in kw.items(): + setattr(c, k, v) + return c + + +def _run_loss(config, **loss_kw): + torch.manual_seed(0) + model = Diffusion_Planner(config).train() + inputs = _inputs() + B, Pn = 2, MAX_NUM_NEIGHBORS + ego_future = torch.randn(B, OUTPUT_T, POSE_DIM) + neighbors_future = torch.randn(B, Pn, OUTPUT_T, POSE_DIM) + mask = torch.zeros(B, Pn, OUTPUT_T, dtype=torch.bool) + args = _loss_args(config) + for k, v in loss_kw.items(): + setattr(args, k, v) + return compute_plantf_training_loss(model, inputs, (ego_future, neighbors_future, mask), args) + + +def test_cross_attn_head_shapes(): + torch.manual_seed(0) + head = PlanTFCrossAttnHead( + embed_dim=HIDDEN_DIM, num_modes=NUM_MODES, future_steps=OUTPUT_T, num_heads=2 + ) + memory = torch.randn(2, 7, HIDDEN_DIM) + valid = torch.ones(2, 7, dtype=torch.bool) + loc, pi = head(memory, valid) + assert loc.shape == (2, NUM_MODES, OUTPUT_T, 4) + assert pi.shape == (2, NUM_MODES) + + +def test_cross_attn_head_predict_scale(): + head = PlanTFCrossAttnHead( + embed_dim=HIDDEN_DIM, + num_modes=NUM_MODES, + future_steps=OUTPUT_T, + num_heads=2, + predict_scale=True, + ) + out = head(torch.randn(2, 7, HIDDEN_DIM), torch.ones(2, 7, dtype=torch.bool)) + assert len(out) == 3 + assert out[2].shape == (2, NUM_MODES, OUTPUT_T, 4) + + +def test_mlp_head_predict_scale(): + head = PlanTFTrajectoryHead( + embed_dim=HIDDEN_DIM, num_modes=NUM_MODES, future_steps=OUTPUT_T, predict_scale=True + ) + out = head(torch.randn(2, HIDDEN_DIM)) + assert len(out) == 3 + assert out[2].shape == (2, NUM_MODES, OUTPUT_T, 4) + + +def test_cross_attn_decoder_forward_and_loss(): + config = _config_with(plantf_head_type="cross_attn") + model = Diffusion_Planner(config).eval() + _, outputs = model(_inputs()) + assert outputs["trajectory"].shape[1] == NUM_MODES + assert "prediction" in outputs + loss = _run_loss(config) + assert torch.isfinite(loss["ego_planning_loss"]) + + +def test_laplace_nll_finite_and_changes_loss(): + base = _run_loss(_config()) + nll = _run_loss(_config_with(plantf_use_laplace_nll=True), plantf_use_laplace_nll=True) + assert torch.isfinite(nll["ego_planning_loss"]) + assert not torch.isclose(base["ego_planning_loss"], nll["ego_planning_loss"]) + + +def test_tail_weighted_regression_changes_ego_loss(): + base = _run_loss(_config()) + tw = _run_loss(_config(), plantf_tail_weight=3.0) + assert not torch.isclose(base["ego_planning_loss"], tw["ego_planning_loss"]) + + +def test_tail_weighted_smoothness_changes_smoothness_loss(): + torch.manual_seed(0) + config = _config() + model = Diffusion_Planner(config).train() + # Zero-init makes the prediction constant (second-difference == 0), so the + # penalty is 0 regardless of weighting. Perturb the loc head so the best-mode + # trajectory actually curves and the penalty is non-zero. + with torch.no_grad(): + model.decoder.trajectory_head.loc[-1].weight.normal_(0, 0.1) + inputs = _inputs() + B, Pn = 2, MAX_NUM_NEIGHBORS + ego_future = torch.randn(B, OUTPUT_T, POSE_DIM) + neighbors_future = torch.randn(B, Pn, OUTPUT_T, POSE_DIM) + mask = torch.zeros(B, Pn, OUTPUT_T, dtype=torch.bool) + + def run(weight): + args = _loss_args(config) + args.plantf_smoothness_tail_weight = weight + return compute_plantf_training_loss( + model, inputs, (ego_future, neighbors_future, mask), args + )["smoothness_loss"] + + base = run(0.0) + weighted = run(3.0) + assert base > 0 + assert not torch.isclose(base, weighted) + + +def test_bezier_basis_partition_of_unity(): + # Bernstein weights sum to 1 at every timestep and reproduce the endpoints + # exactly (first control point at t=0, last at t=1). + basis = bezier_basis(num_control_points=8, num_steps=OUTPUT_T) + assert basis.shape == (8, OUTPUT_T) + assert torch.allclose(basis.sum(dim=0), torch.ones(OUTPUT_T), atol=1e-5) + assert torch.allclose(basis[:, 0], torch.tensor([1.0] + [0.0] * 7), atol=1e-5) + assert torch.allclose(basis[:, -1], torch.tensor([0.0] * 7 + [1.0]), atol=1e-5) + + +def test_basis_head_shapes_and_control_points(): + # The basis head regresses n_ctrl*C values internally but exposes the same + # [B, K, T, C] contract as the flat head. + torch.manual_seed(0) + n_ctrl = 6 + head = PlanTFTrajectoryHead( + embed_dim=HIDDEN_DIM, + num_modes=NUM_MODES, + future_steps=OUTPUT_T, + out_channels=4, + num_control_points=n_ctrl, + ) + assert head.loc[-1].out_features == n_ctrl * 4 # coefficients, not T*4 + loc, pi = head(torch.randn(2, HIDDEN_DIM)) + assert loc.shape == (2, NUM_MODES, OUTPUT_T, 4) + assert pi.shape == (2, NUM_MODES) + + +def test_basis_head_is_structurally_smoother_than_flat(): + # With identical random weights the Bezier expansion produces a far smaller + # xy second-difference (curvature) than the flat per-step head — the whole + # point of the temporal basis. + torch.manual_seed(0) + x = torch.randn(4, HIDDEN_DIM) + + def curvature(head): + loc, _ = head(x) + xy = loc[..., :2] + d2 = xy[:, :, 2:] - 2 * xy[:, :, 1:-1] + xy[:, :, :-2] + return d2.abs().mean() + + torch.manual_seed(1) + flat = PlanTFTrajectoryHead(HIDDEN_DIM, NUM_MODES, OUTPUT_T, out_channels=4) + torch.manual_seed(1) + basis = PlanTFTrajectoryHead( + HIDDEN_DIM, NUM_MODES, OUTPUT_T, out_channels=4, num_control_points=8 + ) + assert curvature(basis) < 0.25 * curvature(flat) + + +def test_basis_decoder_contract_and_zero_init(): + # End-to-end: plantf_head_type="basis" preserves the prediction contract, and + # the shared zero-init still zeroes the trajectory (control points -> 0). + config = _config_with(plantf_head_type="basis", plantf_basis_control_points=8) + model = Diffusion_Planner(config).eval() + assert model.decoder.trajectory_head.num_control_points == 8 + assert torch.count_nonzero(model.decoder.trajectory_head.loc[-1].weight) == 0 + _, outputs = model(_inputs()) + assert outputs["prediction"].shape == (2, 1 + MAX_NUM_NEIGHBORS, OUTPUT_T, POSE_DIM) + assert outputs["trajectory"].shape == (2, NUM_MODES, OUTPUT_T, 4) + + +def test_basis_head_with_velocity_and_backward(): + # basis + velocity representation trains (loss finite, gradients flow). + config = _config_with( + plantf_head_type="basis", + plantf_basis_control_points=8, + use_velocity_representation=True, + ) + loss = _run_loss(config) + total = loss["ego_planning_loss"] + loss["neighbor_prediction_loss"] + loss["mode_cls_loss"] + assert torch.isfinite(total) + total.backward() + + +def test_gru_head_shapes(): + torch.manual_seed(0) + head = PlanTFGRUHead( + embed_dim=HIDDEN_DIM, num_modes=NUM_MODES, future_steps=OUTPUT_T, out_channels=4 + ) + loc, pi = head(torch.randn(2, HIDDEN_DIM)) + assert loc.shape == (2, NUM_MODES, OUTPUT_T, 4) + assert pi.shape == (2, NUM_MODES) + + +def test_gru_decoder_contract_and_zero_init(): + config = _config_with(plantf_head_type="gru") + model = Diffusion_Planner(config).eval() + assert isinstance(model.decoder.trajectory_head, PlanTFGRUHead) + # shared zero-init still zeroes the final loc Linear -> zero trajectory prior + assert torch.count_nonzero(model.decoder.trajectory_head.loc[-1].weight) == 0 + _, outputs = model(_inputs()) + assert outputs["prediction"].shape == (2, 1 + MAX_NUM_NEIGHBORS, OUTPUT_T, POSE_DIM) + assert outputs["trajectory"].shape == (2, NUM_MODES, OUTPUT_T, 4) + + +def test_gru_head_with_velocity_and_backward(): + config = _config_with(plantf_head_type="gru", use_velocity_representation=True) + loss = _run_loss(config) + total = loss["ego_planning_loss"] + loss["neighbor_prediction_loss"] + loss["mode_cls_loss"] + assert torch.isfinite(total) + total.backward() + + +def test_route_rerank_falls_back_without_normalizer(): + config = _config_with(plantf_route_rerank=True) # observation_normalizer stays None + model = Diffusion_Planner(config).eval() + _, outputs = model(_inputs()) + assert "prediction" in outputs + + +def test_route_rerank_runs_with_identity_normalizer(): + inputs = _inputs() + c = inputs["route_lanes"].shape[-1] + obs = ObservationNormalizer({"route_lanes": {"mean": torch.zeros(c), "std": torch.ones(c)}}) + config = _config_with(plantf_route_rerank=True, observation_normalizer=obs) + model = Diffusion_Planner(config).eval() + _, outputs = model(inputs) + assert outputs["prediction"].shape[0] == inputs["route_lanes"].shape[0] diff --git a/diffusion_planner/train_predictor.py b/diffusion_planner/train_predictor.py index c904151d4..498d58293 100644 --- a/diffusion_planner/train_predictor.py +++ b/diffusion_planner/train_predictor.py @@ -86,6 +86,19 @@ def get_args(args_list=None): parser.add_argument("--batch_size", type=int, default=512) parser.add_argument("--save_utd", type=int, default=10) parser.add_argument("--learning_rate", type=float, default=1e-4) + parser.add_argument( + "--encoder_learning_rate", + type=float, + default=None, + help="optional encoder-only base LR; defaults to --learning_rate", + ) + parser.add_argument( + "--weight_decay", + type=float, + default=1e-4, + help="AdamW weight decay on matmul weights only (biases and " + "LayerNorm/BatchNorm/Embedding params are excluded)", + ) parser.add_argument("--warm_up_epoch", type=int, default=5) parser.add_argument("--encoder_drop_path_rate", type=float, default=0.1) parser.add_argument("--decoder_drop_path_rate", type=float, default=0.1) @@ -151,6 +164,127 @@ def get_args(args_list=None): parser.add_argument("--alpha_planning_loss", type=float, default=1.0) parser.add_argument("--alpha_neighbor_loss", type=float, default=0.1) + parser.add_argument( + "--alpha_mode_cls_loss", + type=float, + default=1.0, + help="mode classification loss weight (decoder_type='plantf' only)", + ) + parser.add_argument( + "--coeff_endpoint_fde_loss", + type=float, + default=0.0, + help="normalized-space endpoint loss weight (decoder_type='plantf' only); " + "off by default — it only constrains the endpoint and biases toward " + "straight lines. Forward progress is taught by the per-timestep loss.", + ) + parser.add_argument( + "--plantf_use_lon_velocity_weight", + type=boolean, + default=True, + help="longitudinal velocity down-weighting for the planTF head " + "(on by default = same as the diffusion head; A/B showed disabling it " + "hurt, so it is kept on)", + ) + parser.add_argument( + "--plantf_use_ego_state_in_head", + type=boolean, + default=True, + help="feed the current ego motion state (vx,vy,ax,ay,steering,yaw_rate) " + "into the planTF trajectory head to anchor the prediction to the current " + "motion (decoder_type='plantf' only)", + ) + parser.add_argument( + "--plantf_ego_state_token", + type=boolean, + default=False, + help="A2: replace the ego encoder token with an embedding of the current " + "ego motion state (original planTF use_ego_history=false path)", + ) + parser.add_argument( + "--plantf_ego_state_dropout", + type=float, + default=0.75, + help="dropout rate on the ego motion-state channels for A2 (original planTF 0.75)", + ) + parser.add_argument( + "--plantf_input_delta", + type=boolean, + default=False, + help="C1: feed agent (ego+neighbor) history as consecutive-frame xy deltas", + ) + parser.add_argument( + "--plantf_mask_goal_pose", + type=boolean, + default=False, + help="zero the goal_pose input (global-frame data bug); model plans from route", + ) + parser.add_argument( + "--plantf_relative_xy", + type=boolean, + default=False, + help="original-planTF xy representation: regress every agent future relative " + "to its current position, then add that observed position back at inference. " + "Use only for a newly trained PlantF checkpoint.", + ) + parser.add_argument( + "--coeff_smoothness_loss", + type=float, + default=0.0, + help="weight on the xy second-difference smoothness penalty for the planTF " + "best mode (suppresses comb jitter; try 0.1-1.0)", + ) + + # planTF combinable ablation toggles (docs/plantf_head_development_notes.md §9) + parser.add_argument( + "--plantf_head_type", + type=str, + default="mlp", + choices=["mlp", "cross_attn", "basis", "gru"], + help="trajectory head: 'mlp' (reshape ego token), 'cross_attn' (K mode " + "queries cross-attend to all encoder tokens), 'basis' (mlp head that " + "regresses Bezier control points expanded over time = temporal smoothness), " + "or 'gru' (recurrent head that unrolls the waypoints; experimental)", + ) + parser.add_argument( + "--plantf_basis_control_points", + type=int, + default=8, + help="number of Bezier control points for --plantf_head_type basis " + "(fewer=smoother/stiffer, more=more expressive; ignored otherwise)", + ) + parser.add_argument( + "--plantf_route_rerank", + type=boolean, + default=False, + help="at inference, pick the ego mode by route adherence among the top-k pi " + "modes instead of argmax(pi) (validation path only, not ONNX)", + ) + parser.add_argument( + "--plantf_route_rerank_topk", + type=int, + default=3, + help="number of top-pi modes considered by the route re-ranker", + ) + parser.add_argument( + "--plantf_tail_weight", + type=float, + default=0.0, + help="tail weighting on the ego regression loss (w_t = 1 + w*t/(T-1))", + ) + parser.add_argument( + "--plantf_smoothness_tail_weight", + type=float, + default=0.0, + help="tail weighting on the curvature (smoothness) penalty", + ) + parser.add_argument( + "--plantf_use_laplace_nll", + type=boolean, + default=False, + help="use Laplace NLL (head-predicted per-point log-scale) for the ego " + "regression instead of smooth-L1", + ) # Velocity representation & hybrid loss (HDP paper, Section IV-B) parser.add_argument( @@ -189,9 +323,37 @@ def get_args(args_list=None): choices=["x_start", "flow_matching"], default="x_start", ) + parser.add_argument( + "--decoder_type", + type=str, + choices=["diffusion", "plantf"], + default="diffusion", + help="'plantf' swaps the diffusion decoder for a one-shot multi-modal regression head", + ) + parser.add_argument( + "--num_modes", + type=int, + default=6, + help="number of ego trajectory modes (decoder_type='plantf' only)", + ) parser.add_argument("--predicted_neighbor_num", type=int, default=MAX_NUM_NEIGHBORS) parser.add_argument("--resume_model_path", type=str, help="path to resume model", default=None) + parser.add_argument( + "--pretrained_encoder_path", + type=str, + default=None, + help="warm-start: load ONLY the encoder.* weights from this checkpoint " + "(e.g. a diffusion model trained on production data); the head stays fresh. " + "Use the same normalization the pretrained encoder was trained with.", + ) + parser.add_argument( + "--freeze_encoder_epochs", + type=int, + default=0, + help="freeze the encoder for the first N epochs (train the head only), " + "then unfreeze for joint fine-tuning (only with matching normalization)", + ) parser.add_argument("--use_wandb", default=False, type=boolean) parser.add_argument( @@ -205,6 +367,21 @@ def get_args(args_list=None): ) parser.add_argument("--notes", default="", type=str) + # Portability toggles (disable for lightweight/dependency-free runs). + parser.add_argument( + "--enable_checkpoint_viz", + default=True, + type=boolean, + help="render a checkpoint trajectory PNG each save (matplotlib + an extra " + "forward pass); disable to avoid the dependency and overhead", + ) + parser.add_argument( + "--enable_onnx_export", + default=True, + type=boolean, + help="export ONNX graphs at each checkpoint save; disable to skip ONNX/onnxruntime", + ) + # distributed training parameters parser.add_argument("--ddp", default=True, type=boolean, help="use ddp or not") parser.add_argument("--port", default="22323", type=str, help="port") diff --git a/diffusion_planner/train_run.py b/diffusion_planner/train_run.py index 7c7e1125c..fbcefcf68 100644 --- a/diffusion_planner/train_run.py +++ b/diffusion_planner/train_run.py @@ -25,6 +25,18 @@ def parse_args() -> argparse.Namespace: p.add_argument("--resume_model_path", default=None, help="optional: resume from this .pth") p.add_argument("--wandb_run_id", default=None, help="optional: existing wandb run id") p.add_argument("--wandb_project_name", default=None, help="optional: wandb project name") + p.add_argument( + "--decoder_type", + default="diffusion", + choices=["diffusion", "plantf"], + help="optional: 'plantf' trains the one-shot multi-modal regression head", + ) + p.add_argument( + "--num_modes", + type=int, + default=6, + help="optional: number of ego trajectory modes (decoder_type='plantf' only)", + ) return p.parse_args() @@ -50,6 +62,8 @@ def main() -> None: optional += ["--wandb_run_id", args.wandb_run_id] if args.wandb_project_name: optional += ["--wandb_project_name", args.wandb_project_name] + if args.decoder_type != "diffusion": + optional += ["--decoder_type", args.decoder_type, "--num_modes", str(args.num_modes)] Path("/tmp/tmp_dist_init").unlink(missing_ok=True) diff --git a/diffusion_planner/valid_predictor.py b/diffusion_planner/valid_predictor.py index 91642f701..d32816272 100644 --- a/diffusion_planner/valid_predictor.py +++ b/diffusion_planner/valid_predictor.py @@ -244,6 +244,10 @@ def run_validation(valid_cfg: ValidConfig): ) if "ego_road_border_loss" in agg["ego_means"]: print(f"ego_road_border_loss_mean={agg['ego_means']['ego_road_border_loss']:.4f}") + for key, value in agg["plantf_means"].items(): + print(f"plantf_{key}={value:.4f}") + for key, value in agg["traj_means"].items(): + print(f"traj_{key}={value:.4f}") if replan_agg.get("replan_consistency_count", 0) > 0: print( "replan_position_consistency={:.4f} replan_heading_consistency={:.4f} " @@ -270,6 +274,8 @@ def run_validation(valid_cfg: ValidConfig): "turn_indicator_change_accuracy": turn_indicator_change_accuracy, "turn_indicator_change_total": turn_indicator_change_total, **agg["ego_means"], + **{f"plantf_{key}": value for key, value in agg["plantf_means"].items()}, + **{f"traj_{key}": value for key, value in agg["traj_means"].items()}, **replan_agg, **{f"epdms_{key}": value for key, value in agg["epdms_means"].items()}, } diff --git a/docs/plantf_head_integration_plan.md b/docs/plantf_head_integration_plan.md new file mode 100644 index 000000000..c635303ce --- /dev/null +++ b/docs/plantf_head_integration_plan.md @@ -0,0 +1,208 @@ +# PlAnTF Head Integration in Diffusion Planner + +This document records the current design and implementation status of the +PlAnTF-style regression head in Diffusion Planner (DP). For the production +command and parameter reference, see [PlantF Usage and Production Training +Guide](plantf_usage.md). + +## Status + +The PlantF decoder is implemented on `feat/plantf-decoder-head` and is selected +with `decoder_type=plantf`. The implementation is covered by +`diffusion_planner/tests/test_plantf_decoder.py`. + +Implemented: + +- a PlAnTF-style winner-takes-all ego trajectory head and direct neighbor + predictor; +- `mlp`, `basis`, `gru`, and `cross_attn` ego-head variants; +- DP-compatible `prediction`, turn-indicator, normalization, and safety-loss + interfaces; +- agent-relative XY output support (`plantf_relative_xy`); +- tail-weighted regression, second-difference smoothness, and masked neighbor + losses; +- full and split ONNX export wrappers; and +- open-loop validation, replan-consistency evaluation, and focused unit tests. + +Validated so far: + +- unit tests for PlantF decoding, loss paths, and ONNX wrappers; +- local small-subset training and validation smoke tests; and +- open-loop experiments with the current mode-one production recipe. + +Still required before production deployment: + +- a controlled EMA-versus-regular-weight evaluation; +- closed-loop comparison against the DP head on matched scenarios; and +- deployment validation in the target Autoware runtime. + +## Goal + +PlantF retains DP's data contract and encoder, while replacing the iterative +diffusion decoder with a one-shot PlAnTF-style trajectory regression decoder. +The objective is to compare a direct regression head and the diffusion head +under the same DP input representation, encoder, data, and output interface. + +The adaptation deliberately keeps the DP ecosystem intact: + +- input features such as `ego_agent_past`, `neighbor_agents_past`, lanes, + routes, polygons, and line strings are unchanged; +- the DP encoder and its token order are unchanged; +- downstream consumers continue to receive + `prediction: [B, 1 + P_n, T, 4]`; and +- the decoder is selected by a configuration switch rather than by a separate + model family. + +The output channels are `(x, y, cos(heading), sin(heading))`. The PlantF +implementation uses the DP horizon (`T=80` at 0.1 s intervals), whereas the +original planTF implementation predicts a much shorter set of sparse future +poses. This temporal-grid difference is important when interpreting the first +predicted point: the PlantF head does not explicitly pin it to the current +state. + +## Architecture + +| Component | DP diffusion head | PlantF head | +| --- | --- | --- | +| Encoder | DP MLP-Mixer and fusion encoder | Shared DP encoder | +| Ego decoder | DiT-like denoising model | One-shot multimodal regression head | +| Neighbor decoder | Joint diffusion output | Direct per-neighbor regression head | +| Candidate selection | Sampling result | Highest-probability trajectory mode | +| Inference | Iterative solver / denoising schedule | One decoder forward pass | +| Guidance and delay prefix | Supported by the diffusion path | Not part of the one-shot decoder | + +The first encoder token is the ego token. PlantF uses it to generate `K` +candidate ego trajectories and `K` mode logits. Neighbor tokens are passed to +the direct neighbor predictor. At inference, the argmax-probability ego mode is +concatenated with the neighbor predictions so that the final tensor remains +compatible with DP consumers. + +During training, the decoder additionally returns: + +- `trajectory: [B, K, T, 4]` in state-normalized coordinates; +- `probability: [B, K]` mode logits; and +- `neighbor_prediction: [B, P_n, T, 4]` in state-normalized coordinates. + +The deployed `prediction` is converted back to the DP ego-centric metric +coordinate system. + +## Training objective + +For each sample, the ego candidate with the lowest XY ADE is selected without +gradient. The selected trajectory receives Smooth L1 supervision on all four +output channels. Valid neighbor futures receive masked Smooth L1 supervision. + +The shared training loop also supports: + +- turn-indicator objectives inherited from DP; +- road-border penalty; +- optional neighbor-collision penalty; +- tail weighting of ego regression and smoothness terms; and +- a second-difference XY smoothness penalty. + +For `num_modes > 1`, a cross-entropy objective trains the mode probabilities. +For `num_modes=1`, this loss is mathematically zero and is intentionally not +computed or logged. + +The current production baseline uses `num_modes=1`. This is a deliberate +choice: current data did not provide sufficient evidence that a learned +mode-selection distribution improves the selected trajectory. Multi-mode +training remains supported and retains its mode diagnostics. + +## Coordinates and normalization + +PlantF uses the same `StateNormalizer` contract as DP. The normalization file +is therefore part of the model contract and must match across: + +1. pretrained encoder loading; +2. PlantF training and validation; +3. checkpoint resume; and +4. inference and ONNX export. + +With `plantf_relative_xy=True`, ego and neighbor futures are regressed relative +to their own observed current positions. The decoder restores those positions +before exposing the normal DP-shaped prediction. This is the recommended +representation for a newly trained PlantF head. Do not resume an absolute-XY +PlantF head with this flag enabled, or vice versa. + +## DP interface and feature support + +`Diffusion_Planner` chooses the decoder through `build_decoder`: + +```python +if decoder_type == "diffusion": + decoder = Decoder(config) +elif decoder_type == "plantf": + decoder = PlanTFDecoder(config) +``` + +This keeps training, validation, checkpoint visualization, and most deployment +call sites shared. PlantF does not use diffusion samples, diffusion time, or +the denoising delay internally. Some export wrappers retain these inputs only +to preserve an existing full-graph interface. + +The following features are diffusion-specific and should not be assumed to +affect a PlantF prediction: + +- iterative denoising steps; +- diffusion guidance; +- delay/prefix constraints applied inside the denoising loop; and +- intermediate denoising trajectory visualizations. + +## ONNX and Autoware considerations + +The recommended integration path is the full PlantF ONNX graph with a +single-step runtime. The full wrapper retains the legacy input names required +by the existing DP interface, including inputs that the one-shot decoder does +not use directly. + +The split multi-step DP runtime is not compatible with PlantF: it expects a +diffusion decoder inside a DPM-Solver loop and, in some configurations, a +separate turn-indicator graph. A dedicated one-shot split runtime would be +required to use split PlantF graphs in that path. + +Before deployment, verify all of the following with the actual target runtime: + +1. `predicted_neighbor_num=320`, yielding the expected 321-agent output axis; +2. full-graph input and output names and dynamic batch dimensions; +3. numerical agreement between PyTorch and ONNX Runtime; and +4. the intended behavior without diffusion guidance and delay-prefix support. + +## Validation and model selection + +For `num_modes=1`, PlantF logs top-1 ADE, FDE, and miss rate, plus trajectory +progress/length, speed MAE, second-difference smoothness, stop/slow/move +stratification, and replan consistency when consecutive validation frames are +available. + +Mode-only diagnostics that become constants with one mode are omitted. In +multi-mode runs, min-over-mode metrics, oracle gap, mode accuracy, entropy, +mode usage, and the classification loss remain available. + +The current global best checkpoint is selected by validation lateral error. It +is not a full driving-quality ranking. FDE, miss rate, moving-scene metrics, +smoothness, road-border behavior, neighbor clearance, and replan consistency +must be reviewed before selecting a deployment candidate. + +## Tests + +Run the PlantF test suite from the repository root: + +```bash +cd diffusion_planner +PYTHONPATH=. ../.venv/bin/python -m pytest tests/test_plantf_decoder.py -q +``` + +For a training-path smoke test, use a small local train/validation subset and +run one epoch with the same representation flags as the intended production +run. Confirm that `train_log.tsv` contains the expected mode-one metrics and +that trajectory overlays have the correct coordinate anchoring. + +## References + +- J. Cheng et al., [planTF](https://github.com/jchengai/planTF), + *Rethinking Imitation-based Planner for Autonomous Driving*, ICRA 2024. +- `diffusion_planner/model/module/plantf_decoder.py` +- `diffusion_planner/diffusion_planner/model/diffusion_planner.py` +- `diffusion_planner/diffusion_planner/validate_model.py` +- [PlantF Usage and Production Training Guide](plantf_usage.md) diff --git a/docs/plantf_usage.md b/docs/plantf_usage.md new file mode 100644 index 000000000..2fa2f482f --- /dev/null +++ b/docs/plantf_usage.md @@ -0,0 +1,168 @@ +# PlAnTF Head: Usage and Production Training Guide + +This document describes the PlAnTF-style regression head integrated into +Diffusion Planner. It explains how it differs from the original diffusion +head and gives the current production training recipe. + +## Scope and compatibility + +Set `--decoder_type plantf` to replace only the decoder. The following parts +remain shared with Diffusion Planner (DP): + +- input feature dictionary and coordinate conventions; +- DP encoder and token layout; +- output contract: `prediction` has shape `[B, 1 + P, T, 4]`, with + `(x, y, cos(heading), sin(heading))`; and +- turn-indicator output, road-border loss, and optional collision loss. + +Consequently, PlantF checkpoints are compatible with the existing validation, +visualization, and deployment interfaces, but **are not interchangeable with +diffusion-head weights**. In particular, a PlantF checkpoint must be loaded +with `decoder_type=plantf` and the same output representation options used at +training time. + +## PlantF versus the DP head + +| Aspect | DP diffusion head | PlantF head | +| --- | --- | --- | +| Ego trajectory generation | Iterative diffusion / flow-matching denoising | One-shot regression from the encoded scene | +| Ego candidates | One generated trajectory | `num_modes` candidate trajectories with mode logits | +| Deployed ego trajectory | Diffusion sampler output | Highest-probability mode (`argmax` of mode logits) | +| Ego training objective | Diffusion reconstruction objective | Winner-takes-all Smooth L1 regression on the best xy-ADE mode; mode classification is used only when `num_modes > 1` | +| Neighbor prediction | Jointly generated by the diffusion decoder | Direct per-neighbor regression head | +| Inference cost | Depends on the denoising schedule | A single decoder forward pass | +| Diffusion-only controls | Guidance, denoising steps, prefix/delay behavior | Not used by the one-shot PlantF decoder | + +PlantF exposes the candidate trajectories and probabilities during validation +for diagnostics. Its normal inference output remains DP-shaped: the selected +ego mode followed by the predicted neighbor trajectories. + +## Recommended production configuration + +The current production baseline is a two-GPU, mode-one MLP head. It was chosen +for stable direct-regression behavior and for comparability with the existing +DP encoder checkpoint. + +| Setting | Value | Rationale | +| --- | --- | --- | +| `decoder_type` | `plantf` | Enable the one-shot PlantF decoder. | +| `plantf_head_type` | `mlp` | Current production baseline. `basis`, `gru`, and `cross_attn` are experiments, not the default production recipe. | +| `num_modes` | `1` | Avoids under-supported mode selection on the current dataset. Mode-classification loss and mode-only diagnostics are omitted automatically. | +| `plantf_relative_xy` | `True` | Predict each agent future relative to its observed current position, then restore the position in the decoder. Train from scratch for this representation; do not reuse an absolute-xy PlantF head. | +| `plantf_use_ego_state_in_head` | `True` | Provides current ego motion state to the trajectory head. | +| `use_velocity_representation` | `False` | Use absolute future waypoints after relative-xy anchoring; this is the validated baseline. | +| `plantf_tail_weight` | `1.0` | Gives later waypoints more regression weight. | +| `coeff_smoothness_loss` | `1.0` | Penalizes second differences and reduces comb-like trajectories. | +| `plantf_smoothness_tail_weight` | `1.0` | Places additional smoothness emphasis at the tail. | +| `freeze_encoder_epochs` | `2` | Warm-start the new head before joint fine-tuning. | +| `ego_history_dropout_rate` | `0.0` | Current baseline disables this DP input dropout. | +| augmentation | `quintic`, probability `0.5` | Current DP data-augmentation baseline. | +| global batch size | `96` on two GPUs | The training code splits this to 48 samples per rank. | +| head / encoder LR | `5e-4` / `5e-5` | Current production baseline. | + +### Normalization is part of the checkpoint contract + +Use the exact same normalization file for the pretrained encoder, PlantF +training, validation, and inference. Do not change `plantf_relative_xy` or the +normalization file while resuming a run. A mismatch changes the meaning and +scale of encoder inputs and target coordinates. + +## Two-GPU production command + +Run this from `diffusion_planner/`. Replace all placeholder paths before +starting. `--batch_size` is the **global** batch size, not the per-GPU batch +size. + +```bash +CUDA_VISIBLE_DEVICES=0,1 torchrun --standalone --nproc_per_node=2 \ + --master_port=29543 train_predictor.py \ + --exp_name plantf_prod_mode1_mlp \ + --save_dir /ABS/PATH/TO/DP_exp/plantf_prod_mode1_mlp \ + --train_set_list /ABS/PATH/TO/path_list_train.json \ + --valid_set_list /ABS/PATH/TO/path_list_valid.json \ + --normalization_file_path /ABS/PATH/TO/normalization.json \ + --pretrained_encoder_path /ABS/PATH/TO/dp_encoder_checkpoint.pth \ + --decoder_type plantf \ + --plantf_head_type mlp \ + --num_modes 1 \ + --plantf_relative_xy True \ + --plantf_use_ego_state_in_head True \ + --plantf_mask_goal_pose False \ + --use_velocity_representation False \ + --plantf_tail_weight 1.0 \ + --coeff_smoothness_loss 1.0 \ + --plantf_smoothness_tail_weight 1.0 \ + --coeff_road_border_loss 1.0 \ + --coeff_neighbor_collision_loss 0.0 \ + --freeze_encoder_epochs 2 \ + --ego_history_dropout_rate 0.0 \ + --use_data_augment True \ + --augment_type quintic \ + --augment_prob 0.5 \ + --batch_size 96 \ + --learning_rate 5e-4 \ + --encoder_learning_rate 5e-5 \ + --warm_up_epoch 5 \ + --weight_decay 1e-4 \ + --train_epochs 40 \ + --save_utd 10 \ + --num_workers 4 \ + --seed 3407 \ + --ddp True \ + --use_ema True \ + --use_wandb True \ + --wandb_project_name DP_planTF \ + --enable_temporal_stability_eval True \ + --enable_replan_consistency_eval True \ + --enable_onnx_export False \ + --enable_checkpoint_viz False +``` + +`coeff_neighbor_collision_loss=0.0` is intentional for this baseline and +should be treated as an explicit experiment choice, not a safety guarantee. +Enable and validate this term separately before changing the production +recipe. + +## Checkpoints, EMA, and model selection + +- `latest.pth` is overwritten after every epoch and contains the model, EMA + state, optimizer, and scheduler state for resuming. +- `epochXXXX/best_model.pth` is a periodic snapshot written every `save_utd` + epochs. Despite the filename, it is not the global best checkpoint. +- `best_model/best_model.pth` is updated when validation lateral error reaches + a new minimum. +- Checkpoints store both regular and EMA weights. The current validation and + best-checkpoint selection use the regular model weights; choose EMA for + inference only after evaluating it explicitly with the same protocol. + +The current best selection criterion is lateral error. Do not interpret it as +a complete driving-quality ranking: inspect FDE, miss rate, moving-scene FDE, +smoothness, road-border behavior, and replan consistency before promoting a +checkpoint. + +## Validation signals for `num_modes=1` + +The following PlantF validation signals are retained: + +- `valid_mode/top1_ade`, `valid_mode/top1_fde`, and + `valid_mode/miss_rate_2m`; +- trajectory length/progress, speed MAE, and second-difference smoothness; +- stop / slow / move stratified trajectory metrics; and +- replan position and heading consistency when consecutive validation frames + are available. + +The following are intentionally omitted for mode one because they are exact +constants or duplicates: min-over-mode metrics, oracle gap, mode accuracy, +mode entropy, mode usage, global duplicate FDE, raw endpoint distances, and +the zero-valued mode-classification loss. These metrics and loss are retained +for `num_modes > 1`. + +## Before starting a long run + +1. Verify every path in the train and validation JSON files exists. +2. Confirm that the normalization file matches the pretrained encoder. +3. Run one epoch on a small local subset and inspect `train_log.tsv`. +4. Confirm W&B is online if remote monitoring is required. +5. Confirm the global batch size is divisible by the DDP world size. +6. After the first checkpoint, inspect ego and neighbor trajectory overlays; + open-loop losses alone do not reveal all coordinate or anchoring errors. diff --git a/ros_scripts/torch2onnx.py b/ros_scripts/torch2onnx.py index 50b42c0fb..656feacef 100644 --- a/ros_scripts/torch2onnx.py +++ b/ros_scripts/torch2onnx.py @@ -178,6 +178,28 @@ def validate_full_model( ) +def validate_plantf_split_models( + wrappers: ModelWrappers, + inputs: TensorDict, + encoder_onnx_path: Path, + decoder_onnx_path: Path, +) -> None: + with torch.no_grad(): + torch_encoding = wrappers.encoder(*(inputs[name] for name in ENCODER_INPUT_NAMES)) + torch_prediction, torch_probability, torch_turn_indicator = wrappers.decoder(torch_encoding) + + encoder_onnx_inputs = {name: inputs[name].cpu().numpy() for name in ENCODER_INPUT_NAMES} + onnx_encoding = run_ort_in_subprocess(encoder_onnx_path, encoder_onnx_inputs)[0] + compare("encoding", torch_encoding.cpu().numpy(), onnx_encoding) + + onnx_prediction, onnx_probability, onnx_turn_indicator = run_ort_in_subprocess( + decoder_onnx_path, {"encoding": onnx_encoding} + ) + compare("prediction", torch_prediction.cpu().numpy(), onnx_prediction) + compare("probability", torch_probability.cpu().numpy(), onnx_probability) + compare("turn_indicator_logit", torch_turn_indicator.cpu().numpy(), onnx_turn_indicator) + + def validate_split_models( wrappers: ModelWrappers, inputs: TensorDict, @@ -265,29 +287,35 @@ def convert_model( print("\nORT validation") wrappers = build_wrappers(model) validation_inputs = load_validation_inputs(eval_npz_path) - with torch.no_grad(): - validation_encoding = wrappers.encoder( - *(validation_inputs[name] for name in ENCODER_INPUT_NAMES) - ) - validation_decoder_inputs = build_decoder_inputs(validation_inputs, validation_encoding) validate_full_model(wrappers, validation_inputs, full_onnx_path) - validate_split_models( - wrappers, - validation_inputs, - validation_decoder_inputs, - encoder_onnx_path, - decoder_onnx_path, - turn_indicator_onnx_path, - ) - print( - "\nSuccessfully converted to ONNX:" - f"\n {full_onnx_path}" - f"\n {encoder_onnx_path}" - f"\n {decoder_onnx_path}" - f"\n {turn_indicator_onnx_path}\n" - ) + if wrappers.turn_indicator is None: # planTF head: no denoising loop + validate_plantf_split_models( + wrappers, + validation_inputs, + encoder_onnx_path, + decoder_onnx_path, + ) + exported = [full_onnx_path, encoder_onnx_path, decoder_onnx_path] + else: + with torch.no_grad(): + validation_encoding = wrappers.encoder( + *(validation_inputs[name] for name in ENCODER_INPUT_NAMES) + ) + validation_decoder_inputs = build_decoder_inputs(validation_inputs, validation_encoding) + validate_split_models( + wrappers, + validation_inputs, + validation_decoder_inputs, + encoder_onnx_path, + decoder_onnx_path, + turn_indicator_onnx_path, + ) + exported = [full_onnx_path, encoder_onnx_path, decoder_onnx_path, turn_indicator_onnx_path] + + paths = "".join(f"\n {p}" for p in exported) + print(f"\nSuccessfully converted to ONNX:{paths}\n") if __name__ == "__main__":