Temporary fix: Access average z position (elevation) from log replays

examples/eval/run_wosac_eval.py

@@ -13,6 +13,7 @@
 from gpudrive.env.env_torch import GPUDriveTorchEnv
 from gpudrive.env.dataset import SceneDataLoader
 from gpudrive.datatypes.observation import GlobalEgoState
+from gpudrive.datatypes.metadata import Metadata
 from gpudrive.datatypes.info import Info
 from gpudrive.utils.checkpoint import load_agent
 from gpudrive.visualize.utils import img_from_fig
@@ -30,6 +31,11 @@
 
 def get_state(env):
     """Obtain raw agent states."""
+    avg_z_pos = Metadata.from_tensor(
+        metadata_tensor=env.sim.metadata_tensor(),
+        backend=env.backend,
+        device=env.device,
+    ).avg_z
     ego_state = GlobalEgoState.from_tensor(
         env.sim.absolute_self_observation_tensor(),
         backend=env.backend,
@@ -41,7 +47,7 @@ def get_state(env):
     return (
         ego_state.pos_x + mean_x,
         ego_state.pos_y + mean_y,
-        ego_state.pos_z,
+        avg_z_pos, #ego_state.pos_z
         ego_state.rotation_angle,
         ego_state.id,
     )
@@ -229,9 +235,9 @@ def rollout(
     # Settings
     MAX_AGENTS = 64
     NUM_ENVS = 1
-    DEVICE = "cuda"  # where to run the env rollouts
+    DEVICE = "cpu"  # where to run the env rollouts
     NUM_ROLLOUTS_PER_BATCH = 1
-    NUM_DATA_BATCHES = 1
+    NUM_DATA_BATCHES = 10
     INIT_STEPS = 10
     DATASET_SIZE = 100
     RENDER = False
@@ -251,7 +257,7 @@ def rollout(
 
     # Load agent
     agent = load_agent(
-        path_to_cpt="checkpoints/model_guidance_log_replay__S_100__04_29_19_49_30_053_015179.pt", 
+        path_to_cpt="checkpoints/model_guidance_log_replay__S_100__04_29_19_49_30_053_015179.pt",
     ).to(DEVICE)
 
     # Override default environment settings to match those the agent was trained with
@@ -265,17 +271,14 @@ def rollout(
     }
 
     # Add fixed overrides specific to WOSAC evaluation
-    fixed_overrides = {
-        "init_steps": INIT_STEPS,
-    }
-
+    config_dict["init_steps"] = INIT_STEPS
+    config_dict["init_mode"] = "wosac_eval"
+
     logging.info(
         f"initializing env with init_mode = {config_dict['init_mode']}"
     )
 
-    env_config = dataclasses.replace(
-        default_config, **config_dict, **fixed_overrides
-    )
+    env_config = dataclasses.replace(default_config, **config_dict)
 
     # Make environment
     env = GPUDriveTorchEnv(

gpudrive/datatypes/metadata.py

@@ -9,11 +9,11 @@ class Metadata:
     `MetaData` in src/types.hpp
 
     Attributes (all masks are 0/1 int of shape (NumWorlds, NumAgents)):
-        id: Unique ID of the agent (non-negative int).
-        isSdc: Whether agent is self-driving car (1/0, -1 padding). 
-        isOfInterest: IDs of agents interested (1/0, -1 padding).
-        isModeled: Whether agent's trajectory needs to be predicted for WOSAC (1/0, -1 padding).
+        is_sdc: Whether agent is self-driving car (1/0, -1 padding). 
+        objects_of_interest: IDs of agents interested (1/0, -1 padding).
+        tracks_to_predict: Whether agent's trajectory needs to be predicted for WOSAC (1/0, -1 padding).
         difficulty: Difficulty of the agent's trajectory to be predicted (0/1/2 if isModeled, 0 if !isModeled, -1 padding).
+        avg_z: Average z-coordinate (elevation) from the first 10 positions of the agent.
     """
 
     def __init__(self, metadata_tensor: torch.Tensor):
@@ -22,7 +22,8 @@ def __init__(self, metadata_tensor: torch.Tensor):
         self.objects_of_interest = metadata_tensor[:, :, 1]
         self.tracks_to_predict = metadata_tensor[:, :, 2]
         self.difficulty = metadata_tensor[:, :, 3]
-
+        self.avg_z = metadata_tensor[:, :, 4]
+
     @classmethod
     def from_tensor(
         cls, metadata_tensor: madrona_gpudrive.madrona.Tensor, 
@@ -38,4 +39,4 @@ def from_tensor(
     @property
     def shape(self) -> tuple[int, ...]:
         """Shape (num_worlds, num_agents) of each metadata mask."""
-        return self.is_sdc.shape
+        return self.is_sdc.shape

gpudrive/env/env_torch.py

@@ -1518,6 +1518,13 @@ def get_scenario_ids(self):
 
     print(f"Number of controlled agents: {control_mask.sum()}")
 
+    metadata = Metadata.from_tensor(
+        env.sim.metadata_tensor(),
+        backend=env.backend,
+        device=env.device,
+    )
+    print(f"avg Z: {metadata.avg_z}")
+
     # Rollout
     obs = env.reset(mask=control_mask)
 

src/json_serialization.hpp

@@ -15,17 +15,37 @@ namespace madrona_gpudrive
         p.y = j.at("y").get<float>();
     }
 
+    // Extract Z coordinate from position JSON
+    float extract_z_from_position(const nlohmann::json &j)
+    {
+        if (j.contains("z")) {
+            return j.at("z").get<float>();
+        }
+        return 0.0f; // Default z value if not provided
+    }
+
     void from_json(const nlohmann::json &j, MapObject &obj)
     {
         obj.mean = {0,0};
         uint32_t i = 0;
+
+        float zPositions[10];
+
         for (const auto &pos : j.at("position"))
-        { 
+        {
             if (i < MAX_POSITIONS)
             {
                 from_json(pos, obj.position[i]);
                 obj.mean.x += (obj.position[i].x - obj.mean.x)/(i+1);
                 obj.mean.y += (obj.position[i].y - obj.mean.y)/(i+1);
+
+                // Store z position separately
+                if (i < 10)
+                {
+                    float z = extract_z_from_position(pos);
+                    zPositions[i] = z;
+                }
+
                 ++i;
             }
             else
@@ -101,11 +121,28 @@ namespace madrona_gpudrive
             from_json(j.at("mark_as_expert"), obj.markAsExpert);
         }
 
-         // Initialize metadata fields to 0
-        obj.metadata.isSdc = 0;
-        obj.metadata.isObjectOfInterest = 0;
-        obj.metadata.isTrackToPredict = 0;
-        obj.metadata.difficulty = 0;
+        // Initialize metadata fields to 0
+        MetaData::zero(obj.metadata);
+
+        // Calculate average z position from the first 10 positions
+        // Store in the avgZ field of map object
+        float avgZ = 0.0f;
+        float avg_count = 0.0f;
+        for (int zi = 0; zi < 10; zi++) {
+            // Only consider valid z positions
+            if (obj.valid[zi]) {
+                avgZ += zPositions[zi];
+                avg_count += 1;
+            }
+        }
+        if (avg_count == 0) {
+            avgZ = 0.0f; // Avoid division by zero
+        }
+        else
+        {
+            avgZ /= avg_count; // Calculate average
+        }
+        obj.metadata.avgZ = avgZ;
 
         // Initialize VBD trajectories to zeros
         for (int i = 0; i < consts::episodeLen; i++) {
@@ -151,7 +188,7 @@ namespace madrona_gpudrive
         else
             road.type = EntityType::None;
 
-        
+
         std::vector<MapVector2> geometry_points_;
         for(const auto &point: j.at("geometry"))
         {
@@ -232,7 +269,7 @@ namespace madrona_gpudrive
             {
                 if(i==MAX_GEOMETRY)
                     break;
-                road.geometry[i] = geometry_points_[i * sample_every_n_]; 
+                road.geometry[i] = geometry_points_[i * sample_every_n_];
             }
             road.numPoints = num_sampled_points;
         }
@@ -309,14 +346,14 @@ namespace madrona_gpudrive
 
         std::string scenario_id = j.at("scenario_id").get<std::string>();
         std::strncpy(map.scenarioId, scenario_id.c_str(), sizeof(map.scenarioId));
-        
+
         auto mean = calc_mean(j);
         map.mean = {mean.first, mean.second};
         map.numObjects = std::min(j.at("objects").size(), static_cast<size_t>(MAX_OBJECTS));
 
         const auto& metadata = j.at("metadata");
         int sdc_index = metadata.at("sdc_track_index").get<int>();
-        
+
         // Create id to object index mapping
         std::unordered_map<int, size_t> idToObjIdx;
         size_t idx = 0;
@@ -343,27 +380,27 @@ namespace madrona_gpudrive
         // Initialize SDC first if valid
         if (sdc_index >= 0 && sdc_index < j.at("objects").size()) {
             j.at("objects")[sdc_index].get_to(map.objects[0]);
-            map.objects[0].metadata.isSdc = 1;
-            
+            map.objects[0].metadata.isSdc = 1.0f;
+
             // Set additional metadata if needed
             int sdc_id = map.objects[0].id;
             if (tracks_to_predict_indices.find(sdc_index) != tracks_to_predict_indices.end()) {
-                map.objects[0].metadata.isTrackToPredict = 1;
+                map.objects[0].metadata.isTrackToPredict = 1.0f;
                 // Find and set difficulty
                 for (const auto& track : metadata.at("tracks_to_predict")) {
                     if (track.at("track_index").get<int>() == sdc_index) {
-                        map.objects[0].metadata.difficulty = track.at("difficulty").get<int>();
+                        map.objects[0].metadata.difficulty = static_cast<float>(track.at("difficulty").get<int>());
                         break;
                     }
                 }
             }
             if (objects_of_interest_ids.find(sdc_id) != objects_of_interest_ids.end()) {
-                map.objects[0].metadata.isObjectOfInterest = 1;
+                map.objects[0].metadata.isObjectOfInterest = 1.0f;
             }
-            
+
             idToObjIdx[sdc_id] = 0;
             idx = 1;
-            
+
             // Remove SDC from sets to avoid double processing
             tracks_to_predict_indices.erase(sdc_index);
             objects_of_interest_ids.erase(sdc_id);
@@ -372,25 +409,25 @@ namespace madrona_gpudrive
         // Initialize tracks_to_predict objects (excluding SDC)
         for (size_t i = 0; i < j.at("objects").size() && idx < map.numObjects; i++) {
             if (i == sdc_index) continue; // Skip SDC as it's already initialized
-            
+
             if (tracks_to_predict_indices.find(i) != tracks_to_predict_indices.end()) {
                 j.at("objects")[i].get_to(map.objects[idx]);
-                map.objects[idx].metadata.isTrackToPredict = 1;
-                
+                map.objects[idx].metadata.isTrackToPredict = 1.0f;
+
                 // Find and set difficulty
                 for (const auto& track : metadata.at("tracks_to_predict")) {
                     if (track.at("track_index").get<int>() == static_cast<int>(i)) {
-                        map.objects[idx].metadata.difficulty = track.at("difficulty").get<int>();
+                        map.objects[idx].metadata.difficulty = static_cast<float>(track.at("difficulty").get<int>());
                         break;
                     }
                 }
-                
+
                 // Check if also object of interest
                 if (objects_of_interest_ids.find(map.objects[idx].id) != objects_of_interest_ids.end()) {
-                    map.objects[idx].metadata.isObjectOfInterest = 1;
+                    map.objects[idx].metadata.isObjectOfInterest = 1.0f;
                     objects_of_interest_ids.erase(map.objects[idx].id);
                 }
-                
+
                 idToObjIdx[map.objects[idx].id] = idx;
                 idx++;
             }
@@ -399,12 +436,12 @@ namespace madrona_gpudrive
         // Initialize objects_of_interest (excluding those already processed)
         for (size_t i = 0; i < j.at("objects").size() && idx < map.numObjects; i++) {
             if (i == sdc_index) continue;
-            
+
             int obj_id = j.at("objects")[i].at("id").get<int>();
             if (objects_of_interest_ids.find(obj_id) != objects_of_interest_ids.end()) {
                 j.at("objects")[i].get_to(map.objects[idx]);
-                map.objects[idx].metadata.isObjectOfInterest = 1;
-                
+                map.objects[idx].metadata.isObjectOfInterest = 1.0f;
+
                 idToObjIdx[map.objects[idx].id] = idx;
                 idx++;
             }
@@ -413,15 +450,15 @@ namespace madrona_gpudrive
         // Initialize all remaining objects
         for (size_t i = 0; i < j.at("objects").size() && idx < map.numObjects; i++) {
             if (i == sdc_index) continue;
-            
+
             int obj_id = j.at("objects")[i].at("id").get<int>();
             if (idToObjIdx.find(obj_id) == idToObjIdx.end()) { // Check if not already processed
                 j.at("objects")[i].get_to(map.objects[idx]);
                 idToObjIdx[map.objects[idx].id] = idx;
                 idx++;
             }
         }
-        
+
         // Process roads
         map.numRoads = std::min(j.at("roads").size(), static_cast<size_t>(MAX_ROADS));
         size_t countRoadPoints = 0;

src/mgr.cpp

@@ -896,19 +896,20 @@ Tensor Manager::scenarioIdTensor() const {
 
 Tensor Manager::metadataTensor() const {
     return impl_->exportTensor(
-        ExportID::MetaData, TensorElementType::Int32,
+        ExportID::MetaData, TensorElementType::Float32,
         {impl_->numWorlds, consts::kMaxAgentCount, MetaDataExportSize}
     );
 }
 
 Tensor Manager::vbdTrajectoryTensor() const {
-    return impl_->exportTensor(ExportID::VBDTrajectory, TensorElementType::Float32,
-                               {
-                                   impl_->numWorlds,
-                                   consts::kMaxAgentCount,
-                                   consts::episodeLen,
-                                   5,
-                               });
+    return impl_->exportTensor(
+        ExportID::VBDTrajectory, TensorElementType::Float32,
+    {
+        impl_->numWorlds,
+        consts::kMaxAgentCount,
+        consts::episodeLen,
+        5,
+    });
 }
 
 void Manager::triggerReset(int32_t world_idx)

src/types.hpp

@@ -435,24 +435,25 @@ namespace madrona_gpudrive
     const size_t ScenarioIdExportSize = 32;
     static_assert(sizeof(ScenarioId) == sizeof(char32_t) * ScenarioIdExportSize);
 
-    //Metadata struct : using agent IDs.
     struct MetaData
     {
-        int32_t isSdc;
-        int32_t isObjectOfInterest;
-        int32_t isTrackToPredict;
-        int32_t difficulty;
-
+        float isSdc;                
+        float isObjectOfInterest;   
+        float isTrackToPredict;     
+        float difficulty;           
+        float avgZ;
+
         static inline void zero(MetaData& metadata)
         {
-            metadata.isSdc = -1;
-            metadata.isObjectOfInterest = -1;
-            metadata.isTrackToPredict = -1;
-            metadata.difficulty = -1;
+            metadata.isSdc = -1.0f;           
+            metadata.isObjectOfInterest = -1.0f; 
+            metadata.isTrackToPredict = -1.0f;   
+            metadata.difficulty = -1.0f;         
+            metadata.avgZ = 0.0f;
         }
     };
-    const size_t MetaDataExportSize = 4;
-    static_assert(sizeof(MetaData) == sizeof(int32_t) * MetaDataExportSize);
+    const size_t MetaDataExportSize = 5;
+    static_assert(sizeof(MetaData) == sizeof(int32_t) * (MetaDataExportSize - 1) + sizeof(float));
 
     struct AgentInterface : public madrona::Archetype<
                                 Action,