Delete agents

Emerge-Lab · Jan 18, 2025 · 25fe4e8 · 25fe4e8
1 parent 1701596
commit 25fe4e8
Show file tree

Hide file tree

Showing 7 changed files with 51 additions and 10 deletions.
diff --git a/src/bindings.cpp b/src/bindings.cpp
@@ -127,7 +127,20 @@ namespace gpudrive
             .def("expert_trajectory_tensor", &Manager::expertTrajectoryTensor)
             .def("set_maps", &Manager::setMaps)
             .def("world_means_tensor", &Manager::worldMeansTensor)
-            .def("metadata_tensor", &Manager::metadataTensor);
+            .def("metadata_tensor", &Manager::metadataTensor)
+            .def("deleteAgents", [](Manager &self, nb::dict py_agents_to_delete) {
+                std::unordered_map<int32_t, std::vector<int32_t>> agents_to_delete;
+
+                // Convert Python dict to C++ unordered_map
+                for (auto item : py_agents_to_delete) {
+                    int32_t key = nb::cast<int32_t>(item.first);
+                    std::vector<int32_t> value = nb::cast<std::vector<int32_t>>(item.second);
+                    agents_to_delete[key] = value;
+                }
+
+                self.deleteAgents(agents_to_delete);
+            })
+            .def("deleted_agents_tensor", &Manager::deletedAgentsTensor);
     }
 
 }
diff --git a/src/level_gen.cpp b/src/level_gen.cpp
@@ -347,6 +347,15 @@ static inline bool shouldAgentBeCreated(Engine &ctx, const MapObject &agentInit)
         return false;
     }
 
+    auto& deletedAgents = ctx.singleton<DeletedAgents>().deletedAgents;
+    for (CountT i = 0; i < consts::kMaxAgentCount; i++)
+    {
+        if(deletedAgents[i] == agentInit.id)
+        {
+            return false;
+        }
+    }
+
     return true;
 }
 

diff --git a/src/mgr.cpp b/src/mgr.cpp
@@ -643,13 +643,27 @@ void Manager::setMaps(const std::vector<std::string> &maps)
     reset(worldIndices);
 }
 
-void Manager::setToDelete(const std::unordered_map<int32_t, std::vector<int32_t>> &agentsToDelete)
+Tensor Manager::deletedAgentsTensor() const
 {
+    return impl_->exportTensor(ExportID::DeletedAgents, TensorElementType::Int32,
+                               {
+                                   impl_->numWorlds,
+                                   consts::kMaxAgentCount,
+                               });
+}
+
+void Manager::deleteAgents(const std::unordered_map<int32_t, std::vector<int32_t>> &agentsToDelete)
+{
+
+    ResetMap resetmap{
+        1,
+    };
+
     if (impl_->cfg.execMode == madrona::ExecMode::CUDA)
     {
 #ifdef MADRONA_CUDA_SUPPORT
         auto &gpu_exec = static_cast<CUDAImpl *>(impl_.get())->gpuExec;
-        auto agentsToDeleteDevicePtr = (int32_t *)gpu_exec.getExported((uint32_t)ExportID::DeleteAgents);
+        auto agentsToDeleteDevicePtr = (int32_t *)gpu_exec.getExported((uint32_t)ExportID::DeletedAgents);
         for (const auto &[worldIdx, agents] : agentsToDelete)
         {
             assert(worldIdx < impl_->cfg.scenes.size());
@@ -659,6 +673,8 @@ void Manager::setToDelete(const std::unordered_map<int32_t, std::vector<int32_t>
             {
                 REQ_CUDA(cudaMemcpy(agentsToDeletePtr + i, &agents[i], sizeof(int32_t), cudaMemcpyHostToDevice));
             }
+            auto resetMapPtr = (ResetMap *)gpu_exec.getExported((uint32_t)ExportID::ResetMap) + worldIdx;
+            REQ_CUDA(cudaMemcpy(resetMapPtr, &resetmap, sizeof(ResetMap), cudaMemcpyHostToDevice));
         }
 #else
         // Handle the case where CUDA support is not available
@@ -668,7 +684,7 @@ void Manager::setToDelete(const std::unordered_map<int32_t, std::vector<int32_t>
     else
     {
         auto &cpu_exec = static_cast<CPUImpl *>(impl_.get())->cpuExec;
-        auto agentsToDeleteDevicePtr = (int32_t *)cpu_exec.getExported((uint32_t)ExportID::DeleteAgents);
+        auto agentsToDeleteDevicePtr = (int32_t *)cpu_exec.getExported((uint32_t)ExportID::DeletedAgents);
         for (const auto &[worldIdx, agents] : agentsToDelete)
         {
             assert(worldIdx < impl_->cfg.scenes.size());
@@ -678,13 +694,14 @@ void Manager::setToDelete(const std::unordered_map<int32_t, std::vector<int32_t>
             {
                 memcpy(agentsToDeletePtr + i, &agents[i], sizeof(int32_t));
             }
+            auto resetMapPtr = (ResetMap *)cpu_exec.getExported((uint32_t)ExportID::ResetMap) + worldIdx;
+            memcpy(resetMapPtr, &resetmap, sizeof(ResetMap));
         }
     }
 
     std::vector<int32_t> worldIndices(impl_->cfg.scenes.size());
     std::iota(worldIndices.begin(), worldIndices.end(), 0);
     reset(worldIndices);
-
 }
 
 

diff --git a/src/mgr.hpp b/src/mgr.hpp
@@ -69,6 +69,7 @@ class Manager {
     MGR_EXPORT madrona::py::Tensor expertTrajectoryTensor() const;
     MGR_EXPORT madrona::py::Tensor worldMeansTensor() const;
     MGR_EXPORT madrona::py::Tensor metadataTensor() const;
+    MGR_EXPORT madrona::py::Tensor deletedAgentsTensor() const;
     madrona::py::Tensor rgbTensor() const;
     madrona::py::Tensor depthTensor() const;
     // These functions are used by the viewer to control the simulation
@@ -79,7 +80,7 @@ class Manager {
                               float headAngle);
     MGR_EXPORT void setMaps(const std::vector<std::string> &maps);
 
-    MGR_EXPORT void setToDelete(const std::unordered_map<int32_t, std::vector<int32_t>> &agentsToDelete);
+    MGR_EXPORT void deleteAgents(const std::unordered_map<int32_t, std::vector<int32_t>> &agentsToDelete);
 
     // TODO: remove parameters
     MGR_EXPORT std::vector<Shape>

diff --git a/src/sim.cpp b/src/sim.cpp
@@ -62,6 +62,7 @@ void Sim::registerTypes(ECSRegistry &registry, const Config &cfg)
     registry.registerSingleton<Map>();
     registry.registerSingleton<ResetMap>();
     registry.registerSingleton<WorldMeans>();
+    registry.registerSingleton<DeletedAgents>();
 
     registry.registerArchetype<Agent>();
     registry.registerArchetype<PhysicsEntity>();
@@ -74,7 +75,7 @@ void Sim::registerTypes(ECSRegistry &registry, const Config &cfg)
     registry.exportSingleton<Map>((uint32_t)ExportID::Map);
     registry.exportSingleton<ResetMap>((uint32_t)ExportID::ResetMap);
     registry.exportSingleton<WorldMeans>((uint32_t)ExportID::WorldMeans);
-    registry.exportSingleton<DeleteAgents>((uint32_t)ExportID::DeleteAgents);
+    registry.exportSingleton<DeletedAgents>((uint32_t)ExportID::DeletedAgents);
 
     registry.exportColumn<AgentInterface, Action>(
         (uint32_t)ExportID::Action);

diff --git a/src/sim.hpp b/src/sim.hpp
@@ -37,7 +37,7 @@ enum class ExportID : uint32_t {
     ResetMap,
     WorldMeans,
     MetaData,
-    DeleteAgents,
+    DeletedAgents,
     NumExports
 };
 

diff --git a/src/types.hpp b/src/types.hpp
@@ -94,8 +94,8 @@ namespace gpudrive
         int32_t reset;
     };   
 
-    struct DeleteAgents {
-        int32_t deleteAgents[consts::kMaxAgentCount];
+    struct DeletedAgents {
+        int32_t deletedAgents[consts::kMaxAgentCount];
     };
 
     struct WorldMeans {