Add feature flag

Emerge-Lab · May 8, 2024 · 92c91f7 · 92c91f7
1 parent 8dceca2
commit 92c91f7
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 2 deletions.
diff --git a/src/init.hpp b/src/init.hpp
@@ -95,6 +95,11 @@ namespace gpudrive
         Ignore
     };
 
+    enum class FindRoadObservationsWith {
+      KNearestEntitiesWithRadiusFiltering,
+      AllEntitiesWithRadiusFiltering
+    };
+
     struct Parameters
     {
         float polylineReductionThreshold;
@@ -104,6 +109,8 @@ namespace gpudrive
         CollisionBehaviour collisionBehaviour = CollisionBehaviour::AgentStop; // Default: AgentStop
         uint32_t maxNumControlledVehicles = 10000; // Arbitrary high number to by default control all vehicles 
         bool IgnoreNonVehicles = false; // Default: false
+        FindRoadObservationsWith roadObservationAlgorithm{
+            FindRoadObservationsWith::KNearestEntitiesWithRadiusFiltering};
     };
 
     struct WorldInit

diff --git a/src/mgr.cpp b/src/mgr.cpp
@@ -437,6 +437,17 @@ static std::vector<std::string> getMapFiles(const Manager::Config &cfg)
     return mapFiles;
 }
 
+bool isRoadObservationAlgorithmValid(FindRoadObservationsWith algo) {
+    madrona::CountT roadObservationsCount =
+        sizeof(AgentMapObservations) / sizeof(MapObservation);
+
+    return algo ==
+               FindRoadObservationsWith::KNearestEntitiesWithRadiusFiltering ||
+           (algo ==
+                FindRoadObservationsWith::KNearestEntitiesWithRadiusFiltering &&
+            roadObservationsCount == consts::kMaxRoadEntityCount);
+}
+
 Manager::Impl * Manager::Impl::init(
     const Manager::Config &mgr_cfg)
 {
@@ -445,6 +456,9 @@ Manager::Impl * Manager::Impl::init(
 
     std::vector<std::string> mapFiles = getMapFiles(mgr_cfg);
 
+    assert(isRoadObservationAlgorithmValid(
+        mgr_cfg.params.roadObservationAlgorithm));
+
     switch (mgr_cfg.execMode) {
     case ExecMode::CUDA: {
 #ifdef MADRONA_CUDA_SUPPORT

diff --git a/src/sim.cpp b/src/sim.cpp
@@ -213,8 +213,34 @@ inline void collectObservationsSystem(Engine &ctx,
         arrIndex++;
     }
 
-    selectKNearestRoadEntities<consts::kMaxAgentMapObservationsCount>(
-        ctx, rot, model.position, map_obs.obs);
+    const auto alg = ctx.data().params.roadObservationAlgorithm;
+    if (alg == FindRoadObservationsWith::KNearestEntitiesWithRadiusFiltering) {
+        selectKNearestRoadEntities<consts::kMaxAgentMapObservationsCount>(
+            ctx, rot, model.position, map_obs.obs);
+        return;
+    }
+
+    assert(alg == FindRoadObservationsWith::AllEntitiesWithRadiusFiltering);
+    arrIndex = 0; CountT roadIdx = 0;
+    while(roadIdx < ctx.data().numRoads) {
+        Entity road = ctx.data().roads[roadIdx++];
+        Vector2 relative_pos = Vector2{ctx.get<Position>(road).x, ctx.get<Position>(road).y} - model.position;
+        relative_pos = rot.inv().rotateVec({relative_pos.x, relative_pos.y, 0}).xy();
+        if(relative_pos.length() > ctx.data().params.observationRadius)
+        {
+            continue;
+        }
+        map_obs.obs[arrIndex] = ctx.get<MapObservation>(road);
+        map_obs.obs[arrIndex].position = map_obs.obs[arrIndex].position - model.position;
+        arrIndex++;
+    }
+    while (arrIndex < consts::kMaxRoadEntityCount)
+    {
+        map_obs.obs[arrIndex].position = Vector2{0.f, 0.f};
+        map_obs.obs[arrIndex].heading = 0.f;
+        map_obs.obs[arrIndex].type = (float)EntityType::None;
+        arrIndex++;
+    }
 }
 
 inline void movementSystem(Engine &e,