Export the entity type of the agent in Info Tensor

Emerge-Lab · May 1, 2024 · d09a632 · d09a632
1 parent e25da92
commit d09a632
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/src/level_gen.cpp b/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);
     }
 

diff --git a/src/types.hpp b/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);