From a51a18974e755d8d426ce6e0dcf38c3351a6f095 Mon Sep 17 00:00:00 2001 From: Daphne Cornelisse Date: Tue, 29 Apr 2025 17:45:51 -0400 Subject: [PATCH 1/8] Intended usage for z-axis --- examples/eval/run_wosac_eval.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/examples/eval/run_wosac_eval.py b/examples/eval/run_wosac_eval.py index 6e5178737..520eefcf1 100644 --- a/examples/eval/run_wosac_eval.py +++ b/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.trajectory import LogTrajectory from gpudrive.datatypes.info import Info from gpudrive.utils.checkpoint import load_agent from gpudrive.visualize.utils import img_from_fig @@ -30,6 +31,15 @@ def get_state(env): """Obtain raw agent states.""" + # Take the z_pos as the first element of the logs + log_replay_z_init = LogTrajectory.from_tensor( + env.sim.expert_trajectory_tensor(), + env.num_worlds, + env.max_agent_count, + env.backend, + env.device, + )# TODO: Take .pos_z + ego_state = GlobalEgoState.from_tensor( env.sim.absolute_self_observation_tensor(), backend=env.backend, From 5b0927b93bd0da83ae7017b80ccf995520434579 Mon Sep 17 00:00:00 2001 From: Daphne Cornelisse Date: Tue, 29 Apr 2025 18:46:01 -0400 Subject: [PATCH 2/8] Access average z pos (elevation) from logs --- gpudrive/datatypes/metadata.py | 3 +++ src/init.hpp | 1 + src/json_serialization.hpp | 33 +++++++++++++++++++++++++++++++-- src/types.hpp | 4 +++- 4 files changed, 38 insertions(+), 3 deletions(-) diff --git a/gpudrive/datatypes/metadata.py b/gpudrive/datatypes/metadata.py index b0e6809e8..cf5454a42 100644 --- a/gpudrive/datatypes/metadata.py +++ b/gpudrive/datatypes/metadata.py @@ -14,6 +14,7 @@ class Metadata: isOfInterest: IDs of agents interested (1/0, -1 padding). isModeled: 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). + avgZ: Average z-coordinate (elevation) from the first 10 positions of the agent. """ def __init__(self, metadata_tensor: torch.Tensor): @@ -22,6 +23,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] # Added average z-coordinate + @classmethod def from_tensor( diff --git a/src/init.hpp b/src/init.hpp index 8d4ff60dc..9bee2074e 100755 --- a/src/init.hpp +++ b/src/init.hpp @@ -38,6 +38,7 @@ namespace madrona_gpudrive MapVector2 mean; bool markAsExpert{false}; float vbd_trajectories[consts::episodeLen][5]; // x, y, yaw, vx, vy + std::vector zPositions; }; struct MapRoad diff --git a/src/json_serialization.hpp b/src/json_serialization.hpp index da69da5f1..e19a2f2ab 100644 --- a/src/json_serialization.hpp +++ b/src/json_serialization.hpp @@ -15,10 +15,23 @@ namespace madrona_gpudrive p.y = j.at("y").get(); } + // Extract Z coordinate from position JSON + float extract_z_from_position(const nlohmann::json &j) + { + if (j.contains("z")) { + return j.at("z").get(); + } + 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; + + // Clear z positions array - make sure to add this field to MapObject + obj.zPositions.clear(); + for (const auto &pos : j.at("position")) { if (i < MAX_POSITIONS) @@ -26,6 +39,11 @@ namespace madrona_gpudrive 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 + float z = extract_z_from_position(pos); + obj.zPositions.push_back(z); + ++i; } else @@ -101,12 +119,23 @@ namespace madrona_gpudrive from_json(j.at("mark_as_expert"), obj.markAsExpert); } - // Initialize metadata fields to 0 + // Initialize metadata fields to 0 obj.metadata.isSdc = 0; obj.metadata.isObjectOfInterest = 0; obj.metadata.isTrackToPredict = 0; obj.metadata.difficulty = 0; - + + // Calculate average z position from the first 10 positions (or fewer if less available) + // Store in the avgZ field of the MetaData struct for later access + float avgZ = 0.0f; + size_t zCount = std::min(static_cast(10), obj.zPositions.size()); + if (zCount > 0) { + for (size_t zi = 0; zi < zCount; zi++) { + avgZ += obj.zPositions[zi]; + } + avgZ /= static_cast(zCount); + } + obj.metadata.avgZ = avgZ; // Make sure MetaData struct has this field // Initialize VBD trajectories to zeros for (int i = 0; i < consts::episodeLen; i++) { for (int j = 0; j < 5; j++) { diff --git a/src/types.hpp b/src/types.hpp index a75d801ef..644c31796 100755 --- a/src/types.hpp +++ b/src/types.hpp @@ -442,6 +442,7 @@ namespace madrona_gpudrive int32_t isObjectOfInterest; int32_t isTrackToPredict; int32_t difficulty; + float avgZ; // Average z-coordinate (elevation) from first 10 positions static inline void zero(MetaData& metadata) { @@ -449,9 +450,10 @@ namespace madrona_gpudrive metadata.isObjectOfInterest = -1; metadata.isTrackToPredict = -1; metadata.difficulty = -1; + metadata.avgZ = 0.0f; } }; - const size_t MetaDataExportSize = 4; + const size_t MetaDataExportSize = 5; static_assert(sizeof(MetaData) == sizeof(int32_t) * MetaDataExportSize); struct AgentInterface : public madrona::Archetype< From d8d78b93f09fbd09b6a40a560340fbe0fc3852fd Mon Sep 17 00:00:00 2001 From: Daphne Cornelisse Date: Tue, 29 Apr 2025 19:01:00 -0400 Subject: [PATCH 3/8] Fix intended usage of avg_z --- examples/eval/run_wosac_eval.py | 19 ++++++++----------- gpudrive/datatypes/metadata.py | 5 ++--- 2 files changed, 10 insertions(+), 14 deletions(-) diff --git a/examples/eval/run_wosac_eval.py b/examples/eval/run_wosac_eval.py index 520eefcf1..f17c9383f 100644 --- a/examples/eval/run_wosac_eval.py +++ b/examples/eval/run_wosac_eval.py @@ -13,7 +13,7 @@ from gpudrive.env.env_torch import GPUDriveTorchEnv from gpudrive.env.dataset import SceneDataLoader from gpudrive.datatypes.observation import GlobalEgoState -from gpudrive.datatypes.trajectory import LogTrajectory +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 @@ -31,15 +31,12 @@ def get_state(env): """Obtain raw agent states.""" - # Take the z_pos as the first element of the logs - log_replay_z_init = LogTrajectory.from_tensor( - env.sim.expert_trajectory_tensor(), - env.num_worlds, - env.max_agent_count, - env.backend, - env.device, - )# TODO: Take .pos_z - + # TODO: Take the average elevation + 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, @@ -51,7 +48,7 @@ def get_state(env): return ( ego_state.pos_x + mean_x, ego_state.pos_y + mean_y, - ego_state.pos_z, + ego_state.pos_z, #TODO: Return avg_z_pos ego_state.rotation_angle, ego_state.id, ) diff --git a/gpudrive/datatypes/metadata.py b/gpudrive/datatypes/metadata.py index cf5454a42..019f1c638 100644 --- a/gpudrive/datatypes/metadata.py +++ b/gpudrive/datatypes/metadata.py @@ -23,9 +23,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] # Added average z-coordinate + self.avg_z = metadata_tensor[:, :, 4] - @classmethod def from_tensor( cls, metadata_tensor: madrona_gpudrive.madrona.Tensor, @@ -41,4 +40,4 @@ def from_tensor( @property def shape(self) -> tuple[int, ...]: """Shape (num_worlds, num_agents) of each metadata mask.""" - return self.is_sdc.shape \ No newline at end of file + return self.is_sdc.shape From 8a909cb2fa04c4b532cfee096ed01b2732f6664f Mon Sep 17 00:00:00 2001 From: kevin Date: Wed, 30 Apr 2025 16:21:43 -0400 Subject: [PATCH 4/8] fixes --- examples/eval/run_wosac_eval.py | 4 ++-- gpudrive/datatypes/metadata.py | 11 +++++----- src/json_serialization.hpp | 39 ++++++++++++++++++++++----------- src/types.hpp | 13 +++++------ 4 files changed, 39 insertions(+), 28 deletions(-) diff --git a/examples/eval/run_wosac_eval.py b/examples/eval/run_wosac_eval.py index f17c9383f..8fced9d34 100644 --- a/examples/eval/run_wosac_eval.py +++ b/examples/eval/run_wosac_eval.py @@ -33,7 +33,7 @@ def get_state(env): """Obtain raw agent states.""" # TODO: Take the average elevation avg_z_pos = Metadata.from_tensor( - metadata_tensor=env.sim.metadata_tensor(), + metadata_tensor=env.sim.metadata_tensor(), backend=env.backend, device=env.device, ).avg_z @@ -48,7 +48,7 @@ def get_state(env): return ( ego_state.pos_x + mean_x, ego_state.pos_y + mean_y, - ego_state.pos_z, #TODO: Return avg_z_pos + avg_z_pos, ego_state.rotation_angle, ego_state.id, ) diff --git a/gpudrive/datatypes/metadata.py b/gpudrive/datatypes/metadata.py index 019f1c638..5b6b4a094 100644 --- a/gpudrive/datatypes/metadata.py +++ b/gpudrive/datatypes/metadata.py @@ -9,12 +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). - avgZ: Average z-coordinate (elevation) from the first 10 positions of the agent. + avg_z: Average z-coordinate (elevation) from the first 10 positions of the agent. """ def __init__(self, metadata_tensor: torch.Tensor): @@ -23,7 +22,7 @@ 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] + self.avg_z = metadata_tensor[:, :, 4] @classmethod def from_tensor( diff --git a/src/json_serialization.hpp b/src/json_serialization.hpp index e19a2f2ab..eff70b43f 100644 --- a/src/json_serialization.hpp +++ b/src/json_serialization.hpp @@ -29,8 +29,7 @@ namespace madrona_gpudrive obj.mean = {0,0}; uint32_t i = 0; - // Clear z positions array - make sure to add this field to MapObject - obj.zPositions.clear(); + float zPositions[10]; for (const auto &pos : j.at("position")) { @@ -41,9 +40,12 @@ namespace madrona_gpudrive obj.mean.y += (obj.position[i].y - obj.mean.y)/(i+1); // Store z position separately - float z = extract_z_from_position(pos); - obj.zPositions.push_back(z); - + if (i < 10) + { + float z = extract_z_from_position(pos); + zPositions[i] = z; + } + ++i; } else @@ -124,18 +126,29 @@ namespace madrona_gpudrive obj.metadata.isObjectOfInterest = 0; obj.metadata.isTrackToPredict = 0; obj.metadata.difficulty = 0; + obj.metadata.avgZ = 0.0f; - // Calculate average z position from the first 10 positions (or fewer if less available) - // Store in the avgZ field of the MetaData struct for later access + // Calculate average z position from the first 10 positions + // Store in the avgZ field of map object float avgZ = 0.0f; - size_t zCount = std::min(static_cast(10), obj.zPositions.size()); - if (zCount > 0) { - for (size_t zi = 0; zi < zCount; zi++) { - avgZ += obj.zPositions[zi]; + 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; } - avgZ /= static_cast(zCount); } - obj.metadata.avgZ = avgZ; // Make sure MetaData struct has this field + if (avg_count == 0) { + avgZ = 0.0f; // Avoid division by zero + } + else + { + avgZ /= avg_count; // Calculate average + } + std::cout << "avgZ: " << avgZ << std::endl; + obj.metadata.avgZ = avgZ; + // Initialize VBD trajectories to zeros for (int i = 0; i < consts::episodeLen; i++) { for (int j = 0; j < 5; j++) { diff --git a/src/types.hpp b/src/types.hpp index 644c31796..5cbd3a5a6 100755 --- a/src/types.hpp +++ b/src/types.hpp @@ -435,14 +435,13 @@ 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 avgZ; // Average z-coordinate (elevation) from first 10 positions + float isSdc; + float isObjectOfInterest; + float isTrackToPredict; + float difficulty; + float avgZ; static inline void zero(MetaData& metadata) { @@ -454,7 +453,7 @@ namespace madrona_gpudrive } }; const size_t MetaDataExportSize = 5; - static_assert(sizeof(MetaData) == sizeof(int32_t) * MetaDataExportSize); + static_assert(sizeof(MetaData) == sizeof(int32_t) * (MetaDataExportSize - 1) + sizeof(float)); struct AgentInterface : public madrona::Archetype< Action, From c3911b8199054810bd47160c6f67fc837ee4dda5 Mon Sep 17 00:00:00 2001 From: kevin Date: Thu, 1 May 2025 20:24:55 -0400 Subject: [PATCH 5/8] z pos fixes --- gpudrive/env/env_torch.py | 9 ++++++++- src/init.hpp | 1 - src/types.hpp | 8 ++++---- 3 files changed, 12 insertions(+), 6 deletions(-) diff --git a/gpudrive/env/env_torch.py b/gpudrive/env/env_torch.py index cdafc62ca..e486199f7 100755 --- a/gpudrive/env/env_torch.py +++ b/gpudrive/env/env_torch.py @@ -1710,7 +1710,7 @@ def get_scenario_ids(self): # Create data loader train_loader = SceneDataLoader( - root="data/processed/wosac/debug", + root="data/processed/temp", batch_size=1, dataset_size=1, sample_with_replacement=False, @@ -1730,6 +1730,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) diff --git a/src/init.hpp b/src/init.hpp index 9bee2074e..8d4ff60dc 100755 --- a/src/init.hpp +++ b/src/init.hpp @@ -38,7 +38,6 @@ namespace madrona_gpudrive MapVector2 mean; bool markAsExpert{false}; float vbd_trajectories[consts::episodeLen][5]; // x, y, yaw, vx, vy - std::vector zPositions; }; struct MapRoad diff --git a/src/types.hpp b/src/types.hpp index 5cbd3a5a6..b70a1ac49 100755 --- a/src/types.hpp +++ b/src/types.hpp @@ -437,10 +437,10 @@ namespace madrona_gpudrive struct MetaData { - float isSdc; - float isObjectOfInterest; - float isTrackToPredict; - float difficulty; + int32_t isSdc; + int32_t isObjectOfInterest; + int32_t isTrackToPredict; + int32_t difficulty; float avgZ; static inline void zero(MetaData& metadata) From 3dbbfd054e793f2cccab5c8da14b79ff67d34e3f Mon Sep 17 00:00:00 2001 From: Daphne Cornelisse Date: Fri, 2 May 2025 15:03:58 -0400 Subject: [PATCH 6/8] Fix bug by converting all vals in tensor to floats --- gpudrive/env/env_torch.py | 2 +- src/json_serialization.hpp | 24 ++++++++++-------------- src/mgr.cpp | 17 +++++++++-------- src/types.hpp | 18 +++++++++--------- 4 files changed, 29 insertions(+), 32 deletions(-) diff --git a/gpudrive/env/env_torch.py b/gpudrive/env/env_torch.py index e486199f7..d3c772e34 100755 --- a/gpudrive/env/env_torch.py +++ b/gpudrive/env/env_torch.py @@ -1710,7 +1710,7 @@ def get_scenario_ids(self): # Create data loader train_loader = SceneDataLoader( - root="data/processed/temp", + root="data/processed/wosac/validation_json_1", batch_size=1, dataset_size=1, sample_with_replacement=False, diff --git a/src/json_serialization.hpp b/src/json_serialization.hpp index eff70b43f..f9e5b032d 100644 --- a/src/json_serialization.hpp +++ b/src/json_serialization.hpp @@ -122,11 +122,7 @@ namespace madrona_gpudrive } // Initialize metadata fields to 0 - obj.metadata.isSdc = 0; - obj.metadata.isObjectOfInterest = 0; - obj.metadata.isTrackToPredict = 0; - obj.metadata.difficulty = 0; - obj.metadata.avgZ = 0.0f; + MetaData::zero(obj.metadata); // Calculate average z position from the first 10 positions // Store in the avgZ field of map object @@ -381,26 +377,26 @@ namespace madrona_gpudrive for (const auto& obj_id : metadata.at("objects_of_interest")) { objects_of_interest_ids.insert(obj_id.get()); } - + // 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() == sdc_index) { - map.objects[0].metadata.difficulty = track.at("difficulty").get(); + map.objects[0].metadata.difficulty = static_cast(track.at("difficulty").get()); 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; @@ -417,19 +413,19 @@ namespace madrona_gpudrive 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() == static_cast(i)) { - map.objects[idx].metadata.difficulty = track.at("difficulty").get(); + map.objects[idx].metadata.difficulty = static_cast(track.at("difficulty").get()); 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); } @@ -445,7 +441,7 @@ namespace madrona_gpudrive int obj_id = j.at("objects")[i].at("id").get(); 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++; diff --git a/src/mgr.cpp b/src/mgr.cpp index 2fa511897..092cb184a 100755 --- a/src/mgr.cpp +++ b/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) diff --git a/src/types.hpp b/src/types.hpp index b70a1ac49..8d9fd25d8 100755 --- a/src/types.hpp +++ b/src/types.hpp @@ -437,18 +437,18 @@ namespace madrona_gpudrive 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; } }; From 8f55a783c5046ee10c5a745b87615abe968f201f Mon Sep 17 00:00:00 2001 From: Daphne Cornelisse Date: Fri, 2 May 2025 15:21:05 -0400 Subject: [PATCH 7/8] remove print statement --- src/json_serialization.hpp | 1 - 1 file changed, 1 deletion(-) diff --git a/src/json_serialization.hpp b/src/json_serialization.hpp index f9e5b032d..b16ee55dd 100644 --- a/src/json_serialization.hpp +++ b/src/json_serialization.hpp @@ -142,7 +142,6 @@ namespace madrona_gpudrive { avgZ /= avg_count; // Calculate average } - std::cout << "avgZ: " << avgZ << std::endl; obj.metadata.avgZ = avgZ; // Initialize VBD trajectories to zeros From 7a2c2b45a522f4505330c2d52e34a070441ec3b9 Mon Sep 17 00:00:00 2001 From: Daphne Cornelisse Date: Fri, 2 May 2025 15:26:50 -0400 Subject: [PATCH 8/8] Cleanup --- examples/eval/run_wosac_eval.py | 18 +++++------ src/json_serialization.hpp | 56 ++++++++++++++++----------------- 2 files changed, 35 insertions(+), 39 deletions(-) diff --git a/examples/eval/run_wosac_eval.py b/examples/eval/run_wosac_eval.py index 718ea0b88..d243dad66 100644 --- a/examples/eval/run_wosac_eval.py +++ b/examples/eval/run_wosac_eval.py @@ -31,7 +31,6 @@ def get_state(env): """Obtain raw agent states.""" - # TODO: Take the average elevation avg_z_pos = Metadata.from_tensor( metadata_tensor=env.sim.metadata_tensor(), backend=env.backend, @@ -48,7 +47,7 @@ def get_state(env): return ( ego_state.pos_x + mean_x, ego_state.pos_y + mean_y, - avg_z_pos, + avg_z_pos, #ego_state.pos_z ego_state.rotation_angle, ego_state.id, ) @@ -236,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 @@ -272,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( diff --git a/src/json_serialization.hpp b/src/json_serialization.hpp index b16ee55dd..20c61828a 100644 --- a/src/json_serialization.hpp +++ b/src/json_serialization.hpp @@ -28,17 +28,17 @@ namespace madrona_gpudrive { 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) { @@ -123,7 +123,7 @@ namespace madrona_gpudrive // 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; @@ -188,7 +188,7 @@ namespace madrona_gpudrive else road.type = EntityType::None; - + std::vector geometry_points_; for(const auto &point: j.at("geometry")) { @@ -269,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; } @@ -346,14 +346,14 @@ namespace madrona_gpudrive std::string scenario_id = j.at("scenario_id").get(); 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(MAX_OBJECTS)); const auto& metadata = j.at("metadata"); int sdc_index = metadata.at("sdc_track_index").get(); - + // Create id to object index mapping std::unordered_map idToObjIdx; size_t idx = 0; @@ -376,12 +376,12 @@ namespace madrona_gpudrive for (const auto& obj_id : metadata.at("objects_of_interest")) { objects_of_interest_ids.insert(obj_id.get()); } - + // 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.0f; - + 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()) { @@ -389,18 +389,18 @@ namespace madrona_gpudrive // Find and set difficulty for (const auto& track : metadata.at("tracks_to_predict")) { if (track.at("track_index").get() == sdc_index) { - map.objects[0].metadata.difficulty = static_cast(track.at("difficulty").get()); + map.objects[0].metadata.difficulty = static_cast(track.at("difficulty").get()); break; } } } if (objects_of_interest_ids.find(sdc_id) != objects_of_interest_ids.end()) { - map.objects[0].metadata.isObjectOfInterest = 1.0f; + 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); @@ -409,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.0f; - + 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() == static_cast(i)) { - map.objects[idx].metadata.difficulty = static_cast(track.at("difficulty").get()); + map.objects[idx].metadata.difficulty = static_cast(track.at("difficulty").get()); 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.0f; + map.objects[idx].metadata.isObjectOfInterest = 1.0f; objects_of_interest_ids.erase(map.objects[idx].id); } - + idToObjIdx[map.objects[idx].id] = idx; idx++; } @@ -436,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(); 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.0f; - + map.objects[idx].metadata.isObjectOfInterest = 1.0f; + idToObjIdx[map.objects[idx].id] = idx; idx++; } @@ -450,7 +450,7 @@ 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(); if (idToObjIdx.find(obj_id) == idToObjIdx.end()) { // Check if not already processed j.at("objects")[i].get_to(map.objects[idx]); @@ -458,7 +458,7 @@ namespace madrona_gpudrive idx++; } } - + // Process roads map.numRoads = std::min(j.at("roads").size(), static_cast(MAX_ROADS)); size_t countRoadPoints = 0;