Merge branch 'main' into feat/align_data_struct

Emerge-Lab · Dec 16, 2024 · 01b2e2b · 01b2e2b
2 parents 42c6565 + c353cb7
commit 01b2e2b
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 2 deletions.
diff --git a/pygpudrive/datatypes/observation.py b/pygpudrive/datatypes/observation.py
@@ -94,7 +94,7 @@ def __init__(self, abs_self_obs_tensor: torch.Tensor):
         self.pos_x = abs_self_obs_tensor[:, :, 0]
         self.pos_y = abs_self_obs_tensor[:, :, 1]
         self.pos_z = abs_self_obs_tensor[:, :, 2]
-        self.rotation_as_quaternion = abs_self_obs_tensor[:, :, 3:6]
+        self.rotation_as_quaternion = abs_self_obs_tensor[:, :, 3:7]
         self.rotation_angle = abs_self_obs_tensor[:, :, 7]
         self.goal_x = abs_self_obs_tensor[:, :, 8]
         self.goal_y = abs_self_obs_tensor[:, :, 9]

diff --git a/src/consts.hpp b/src/consts.hpp
@@ -8,7 +8,7 @@ namespace gpudrive {
 
 namespace consts {
 
-inline constexpr madrona::CountT kMaxAgentCount = 32;
+inline constexpr madrona::CountT kMaxAgentCount = 128;
 inline constexpr madrona::CountT kMaxRoadEntityCount = 10000;
 inline constexpr madrona::CountT kMaxAgentMapObservationsCount = 128;
 

diff --git a/src/sim.cpp b/src/sim.cpp
@@ -668,6 +668,7 @@ inline void collectAbsoluteObservationsSystem(Engine &ctx,
     out.rotation.rotationFromAxis = utils::quatToYaw(rotation);
     out.goal = goal;
     out.vehicle_size = vehicleSize;
+    out.id = ctx.get<AgentID>(agent_iface.e).id;
 }
 
 void setupRestOfTasks(TaskGraphBuilder &builder, const Sim::Config &cfg,