[fix] removed redundant neighbor-related input

Author: 0ttwhy4

flow_planner/model/flow_planner_model/flow_planner.py

@@ -102,8 +102,7 @@ def extract_encoder_inputs(self, inputs):
         return encoder_inputs
 
     def extract_decoder_inputs(self, encoder_outputs, inputs):
-        model_extra = dict(neighbor_current_mask=inputs['neighbor_current_mask'],
-                           cfg_flags=inputs['cfg_flags'] if 'cfg_flags' in inputs.keys() else None,)
+        model_extra = dict(cfg_flags=inputs['cfg_flags'] if 'cfg_flags' in inputs.keys() else None,)
         model_extra.update(encoder_outputs)
         return model_extra
 
@@ -151,12 +150,10 @@ def forward_train(self, data: NuPlanDataSample):
         prediction = self.decoder(noised_traj_tokens, t, **decoder_model_extra)
 
         loss_dict = {}
-        
         batch_loss = self.basic_loss(prediction, target_tokens)
         loss_dict['batch_loss'] = batch_loss
 
         loss = torch.sum(batch_loss, dim=-1) # (B, action_num, action_length, dim)
-        loss_dict['neighbor_pred_loss'] = torch.tensor(0.0, device=loss.device)
         loss_dict['ego_planning_loss'] = loss.mean()
 
         if self.planner_params['action_overlap'] > 0: # TODO:

flow_planner/model/model_utils/input_preprocess.py

@@ -44,16 +44,6 @@ def sample_to_model_input(
         if ego_future.numel() != 0:
             ego_future = ego_future[..., 1:1+self.future_len, :3] # (x, y, heading)
 
-        # in the original input, the neighbor_future only include 10 neighbors, while the neighbor_agent_current include 32 neighbors
-        
-        neighbor_future = data.neighbor_future
-        if neighbor_future.numel() != 0:
-            neighbor_future_mask = torch.sum(torch.ne(neighbor_future[..., :self.future_len, :3],0), dim=-1) == 0
-            neighbor_future = neighbor_future[..., :self.future_len, :3] # (x, y, heading)
-
-            neighbor_future[neighbor_future_mask] = 0.
-            neighbor_future_valid = ~neighbor_future_mask
-
         model_inputs = {}
         model_inputs['ego_past'] = data.ego_past.to(device)
         model_inputs['neighbor_past'] = data.neighbor_past.to(device)
@@ -63,10 +53,6 @@ def sample_to_model_input(
         model_inputs['routes'] = data.routes.to(device)
         model_inputs['map_objects'] = data.map_objects.to(device)
 
-        neighbor_current = data.neighbor_past[..., :self.neighbor_pred_num, -1, :4]
-        neighbor_current_mask = torch.sum(torch.ne(neighbor_current[..., :4], 0), dim=-1) == 0
-        model_inputs['neighbor_current_mask'] = neighbor_current_mask.to(device)
-        
 
         ego_current_state = data.ego_current
         model_inputs['ego_current'] = ego_current_state
@@ -75,38 +61,20 @@ def sample_to_model_input(
             ego_current_xy_cos_sin[..., :2],
             torch.atan2(ego_current_xy_cos_sin[..., 3:4], ego_current_xy_cos_sin[..., 2:3])
         ], dim=-1)
-        neighbor_current = torch.cat([
-            neighbor_current[..., :2],
-            torch.atan2(neighbor_current[..., 3:4], neighbor_current[..., 2:3])
-        ], dim=-1)
-        current_states = torch.cat([
-            ego_current[:, None],
-            neighbor_current
-        ], dim=1)
+
+        current_states = ego_current[:, None]
 
         if is_training:
-            gt_future = torch.cat([
-                ego_future[:, None, :, :],
-                neighbor_future,
-            ], dim=1)
+            gt_future = ego_future[:, None, :, :]
 
             gt_with_current = torch.cat([
                     current_states[:, :, None, :],
                     gt_future
                 ], dim=2)
-            
-            neighbor_mask = torch.cat([
-                    neighbor_current_mask.unsqueeze(-1),
-                    neighbor_future_mask
-            ], dim=-1)
-            gt_with_current[:, 1:][neighbor_mask] = 0.
 
             gt_with_current.to(device)
-            neighbor_future_valid.to(device)
-
         else:
             gt_with_current = current_states[:, :, None, :].repeat(1, 1, self.future_len + 1, 1)
-            neighbor_future_valid = None
 
         if kinematic == 'waypoints':
             gt_with_current = torch.cat([
@@ -127,6 +95,4 @@ def sample_to_model_input(
 
         gt_with_current = self.state_preprocess(gt_with_current)
 
-        model_inputs.update({'neighbor_future_valid': neighbor_future_valid})
-
         return model_inputs, gt_with_current

flow_planner/script/core/flow_matching.yaml

@@ -11,5 +11,4 @@ input_aug:
   device: ${device}
 
 ego_planning_loss: 1.0
-neighbor_pred_loss: 0.0
 consistency_loss: 0.5