feat: plantf decoder head

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):

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,15 +123,15 @@ 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,
             class_type=CLASS_TYPE_LINE_STRING,
             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,19 +667,21 @@ 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
         channels_mlp_dim = 128
 
         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):