add input_target coordinate input

Author: Gaaaavin

config/toy.yaml

@@ -10,7 +10,7 @@ training:
   amp: true
   normalize_step_length: true
   resume: false
-  direction_loss_weight: 5.0
+  direction_loss_weight: 3.0
   distance_loss_weight: 1.0
   angle_loss_weight: 1.0
 
@@ -21,7 +21,7 @@ scheduler:
 
 optimizer:
   name: 'adamw'
-  lr: 1e-5
+  lr: 1e-4
 
 model:
   do_rgb_normalize: true
@@ -35,14 +35,14 @@ model:
   output_coordinate_repr: 'euclidean'
 
   cord_embedding:
-    type: 'target'
+    type: 'input_target'
     num_freqs: 6
     include_input: true
 
   encoder_feat_dim: 768
 
   decoder:
-    type: 'diff_policy'
+    type: 'attention'
     len_traj_pred: 5
     num_heads: 4
     num_layers: 4

data/citywalk_dataset.py

@@ -159,6 +159,11 @@ def __getitem__(self, index):
             transformed_input_positions = self.input2target(input_poses, target_pose)
         elif self.cfg.model.cord_embedding.type == 'target':
             transformed_input_positions = self.transform_target_pose(target_pose, current_pose)[np.newaxis, [0, 2]]
+        elif self.cfg.model.cord_embedding.type == 'input_target':
+            transformed_input_positions = np.concatenate([
+                self.input2target(input_poses, target_pose), 
+                self.transform_target_pose(target_pose, current_pose)[np.newaxis, [0, 2]]
+            ], axis=0)
         else:
             raise NotImplementedError(f"Coordinate embedding type {self.cfg.model.cord_embedding} not implemented")
         waypoints_transformed = self.transform_waypoints(waypoint_poses, current_pose)

model/urban_nav.py

@@ -67,6 +67,9 @@ def __init__(self, cfg):
         elif self.cord_embedding_type == 'target':
             self.cord_embedding = PolarEmbedding(cfg)
             self.dim_cord_embedding = self.cord_embedding.out_dim
+        elif self.cord_embedding_type == 'input_target':
+            self.cord_embedding = PolarEmbedding(cfg)
+            self.dim_cord_embedding = self.cord_embedding.out_dim * (self.context_size + 1)
         else:
             raise NotImplementedError(f"Coordinate embedding type {self.cord_embedding_type} not implemented")
 

pl_modules/urban_nav_module.py

@@ -22,7 +22,7 @@ def __init__(self, cfg):
             raise ValueError(f"Unsupported coordinate representation: {self.output_coordinate_repr}")
 
         self.decoder = cfg.model.decoder.type
-        if self.decoder not in ["diff_policy", "transformer"]:
+        if self.decoder not in ["diff_policy", "attention"]:
             raise ValueError(f"Unsupported decoder: {self.decoder}")
 
         # Direction loss weight (you can adjust this value in your cfg)