Export the entity type of the agent in Info Tensor (#98)

* Export the entity type of the agent in Info Tensor

* Export Entity Type bindings

src/bindings.cpp

@@ -54,6 +54,19 @@ namespace gpudrive
             .value("AgentRemoved", CollisionBehaviour::AgentRemoved)
             .value("Ignore", CollisionBehaviour::Ignore);
 
+        nb::enum_<EntityType>(m, "EntityType")
+            .value("_None", EntityType::None) 
+            .value("RoadEdge", EntityType::RoadEdge) 
+            .value("RoadLine", EntityType::RoadLine) 
+            .value("RoadLane", EntityType::RoadLane) 
+            .value("CrossWalk", EntityType::CrossWalk) 
+            .value("SpeedBump", EntityType::SpeedBump) 
+            .value("StopSign", EntityType::StopSign) 
+            .value("Vehicle", EntityType::Vehicle) 
+            .value("Pedestrian", EntityType::Pedestrian) 
+            .value("Cyclist", EntityType::Cyclist) 
+            .value("Padding", EntityType::Padding) 
+            .value("NumTypes", EntityType::NumTypes);
 
         // Bindings for Manager class
         nb::class_<Manager>(m, "SimManager")

src/level_gen.cpp

@@ -40,6 +40,7 @@ static inline void resetAgent(Engine &ctx, Entity agent) {
     ctx.get<Done>(agent).v = 0;
     ctx.get<Reward>(agent).v = 0;
     ctx.get<Info>(agent) = Info{};
+    ctx.get<Info>(agent).type = (int32_t)ctx.get<EntityType>(agent);
 
 #ifndef GPUDRIVE_DISABLE_NARROW_PHASE
     ctx.get<CollisionDetectionEvent>(agent).hasCollided.store_release(0);
@@ -354,6 +355,7 @@ static void resetPaddingEntities(Engine &ctx) {
         ctx.get<Done>(agent).v = 0;
         ctx.get<StepsRemaining>(agent).t = consts::episodeLen;
         ctx.get<Info>(agent) = Info{};
+        ctx.get<Info>(agent).type = (int32_t)ctx.get<EntityType>(agent);
         registerRigidBodyEntity(ctx, agent, SimObject::Agent);
     }
 

src/types.hpp

@@ -88,9 +88,10 @@ struct Info{
     int collidedWithVehicle;
     int collidedWithNonVehicle;
     int reachedGoal;
+    int type;
 };
 
-const size_t InfoExportSize = 4;
+const size_t InfoExportSize = 5;
 
 static_assert(sizeof(Info) == sizeof(int) * InfoExportSize);