Testing agent type in obs

Author: daphne-cornelisse

examples/eval/notebooks/obtain_guidance_data.py

@@ -7,12 +7,12 @@
 from gpudrive.env.dataset import SceneDataLoader
 import torch
 import numpy as np
-
+from gpudrive.datatypes.observation import LocalEgoState
 
 if __name__ == "__main__":
 
-    GUIDANCE_MODE = "vbd_online"
-    DATASET = "data/processed/wosac/validation_json_100"  # Ensure VBD trajectory structures are in here
+    GUIDANCE_MODE = "log_replay"
+    DATASET = "data/processed/wosac/validation_interactive/json"  # Ensure VBD trajectory structures are in here
     SAVE_PATH = "examples/eval/figures_data/guidance/"
 
     env_config = EnvConfig(
@@ -30,7 +30,7 @@
 
     train_loader = SceneDataLoader(
         root=DATASET,
-        batch_size=10,
+        batch_size=1,
         dataset_size=100,
         sample_with_replacement=False,
         shuffle=False,
@@ -41,11 +41,20 @@
         config=env_config,
         data_loader=train_loader,
         max_cont_agents=32,
-        device="cuda",
+        device="cpu",
     )
 
     obs = env.reset(env.cont_agent_mask)
 
+        
+    # Get agent types
+    agent_types = LocalEgoState.from_tensor(
+        self_obs_tensor=env.sim.self_observation_tensor(),
+        backend='torch',
+        device="cpu",
+        #mask=mask,
+    ).agent_type
+
     # Save for analysis
     reference_traj = torch.cat(
         [

gpudrive/datatypes/observation.py

@@ -41,6 +41,7 @@ def __init__(self, self_obs_tensor: torch.Tensor, mask=None):
             self.is_goal_reached = self_obs_tensor[:, 7]
             self.id = self_obs_tensor[:, 8]
             self.steer_angle = self_obs_tensor[:, 9]
+            self.agent_type = self_obs_tensor[:, 10].long()
         else:
             self.speed = self_obs_tensor[:, :, 0]
             self.vehicle_length = self_obs_tensor[:, :, 1] * AGENT_SCALE
@@ -52,6 +53,7 @@ def __init__(self, self_obs_tensor: torch.Tensor, mask=None):
             self.is_goal_reached = self_obs_tensor[:, :, 7]
             self.id = self_obs_tensor[:, :, 8]
             self.steer_angle = self_obs_tensor[:, :, 9]
+            self.agent_type = self_obs_tensor[:, :, 10].long()
 
     @classmethod
     def from_tensor(
@@ -89,7 +91,7 @@ def normalize(self):
             min_val=constants.MIN_REL_GOAL_COORD,
             max_val=constants.MAX_REL_GOAL_COORD,
         )
-        self.steer_angle /= (torch.pi / 3)
+        self.steer_angle /= torch.pi / 3
 
     @property
     def shape(self) -> tuple[int, ...]:

gpudrive/env/env_torch.py

@@ -1753,7 +1753,7 @@ def swap_data_batch(self, data_batch=None):
             controlled_agent_mask=self.cont_agent_mask,
             reference_trajectory=self.reference_trajectory,
         )
-        
+
         self.guidance_dropout_mask = self.create_guidance_dropout_mask()
 
     def get_expert_actions(self):
@@ -1898,7 +1898,7 @@ def render(self, focus_env_idx=0, focus_agent_idx=[0, 1]):
 
     # Create data loader
     train_loader = SceneDataLoader(
-        root="data/processed/validation",
+        root="data/processed/wosac/validation_interactive",
         batch_size=1,
         dataset_size=100,
         sample_with_replacement=False,

src/sim.cpp

@@ -194,6 +194,8 @@ inline void collectSelfObsSystem(Engine &ctx,
     const Info& info = ctx.get<Info>(agent_iface.e);
     self_obs.goalState = info.reachedGoal ? 1.f : 0.f;
     self_obs.steerAngle = vel.angular.z;
+    self_obs.type = (float)ctx.get<EntityType>(agent_iface.e);
+    std::cout << "Self Obs Type: " << self_obs.type << std::endl;
 }
 
 inline void collectPartnerObsSystem(Engine &ctx,

src/types.hpp

@@ -194,6 +194,7 @@ namespace madrona_gpudrive
         float goalState;
         float id;
         float steerAngle;
+        float type;
         static inline SelfObservation zero()
         {
             return SelfObservation{
@@ -203,11 +204,12 @@ namespace madrona_gpudrive
                 .collisionState = 0,
                 .goalState = 0,
                 .id = -1,
-                .steerAngle = 0};
+                .steerAngle = 0,
+                .type = static_cast<float>(EntityType::Padding)};
         }
     };
 
-    const size_t SelfObservationExportSize = 10; // 1 + 3 + 2 + 1 + 1 + 1
+    const size_t SelfObservationExportSize = 11; // 1 + 3 + 2 + 1 + 1 + 1
 
     static_assert(sizeof(SelfObservation) == sizeof(float) * SelfObservationExportSize);