From 170159654b1bc98fd1197b6c8703b3c3a43a44c8 Mon Sep 17 00:00:00 2001 From: Aarav Pandya Date: Fri, 17 Jan 2025 22:03:02 -0500 Subject: [PATCH 1/4] Init --- src/mgr.cpp | 47 ++++++++++++++++++++++++++++++++++++++++++++++- src/mgr.hpp | 3 +++ src/sim.cpp | 1 + src/sim.hpp | 1 + src/types.hpp | 4 ++++ 5 files changed, 55 insertions(+), 1 deletion(-) diff --git a/src/mgr.cpp b/src/mgr.cpp index 96d45c91..54de9df2 100755 --- a/src/mgr.cpp +++ b/src/mgr.cpp @@ -638,11 +638,56 @@ void Manager::setMaps(const std::vector &maps) } // Vector of range on integers from 0 to the number of worlds - std::vector worldIndices(maps.size()); + std::vector worldIndices(impl_->cfg.scenes.size()); std::iota(worldIndices.begin(), worldIndices.end(), 0); reset(worldIndices); } +void Manager::setToDelete(const std::unordered_map> &agentsToDelete) +{ + if (impl_->cfg.execMode == madrona::ExecMode::CUDA) + { +#ifdef MADRONA_CUDA_SUPPORT + auto &gpu_exec = static_cast(impl_.get())->gpuExec; + auto agentsToDeleteDevicePtr = (int32_t *)gpu_exec.getExported((uint32_t)ExportID::DeleteAgents); + for (const auto &[worldIdx, agents] : agentsToDelete) + { + assert(worldIdx < impl_->cfg.scenes.size()); + assert(agents.size() <= consts::kMaxAgentCount); + int32_t *agentsToDeletePtr = agentsToDeleteDevicePtr + worldIdx * consts::kMaxAgentCount; + for (size_t i = 0; i < agents.size(); i++) + { + REQ_CUDA(cudaMemcpy(agentsToDeletePtr + i, &agents[i], sizeof(int32_t), cudaMemcpyHostToDevice)); + } + } +#else + // Handle the case where CUDA support is not available + FATAL("Madrona was not compiled with CUDA support"); +#endif + } + else + { + auto &cpu_exec = static_cast(impl_.get())->cpuExec; + auto agentsToDeleteDevicePtr = (int32_t *)cpu_exec.getExported((uint32_t)ExportID::DeleteAgents); + for (const auto &[worldIdx, agents] : agentsToDelete) + { + assert(worldIdx < impl_->cfg.scenes.size()); + assert(agents.size() <= consts::kMaxAgentCount); + int32_t *agentsToDeletePtr = agentsToDeleteDevicePtr + worldIdx * consts::kMaxAgentCount; + for (size_t i = 0; i < agents.size(); i++) + { + memcpy(agentsToDeletePtr + i, &agents[i], sizeof(int32_t)); + } + } + } + + std::vector worldIndices(impl_->cfg.scenes.size()); + std::iota(worldIndices.begin(), worldIndices.end(), 0); + reset(worldIndices); + +} + + Tensor Manager::actionTensor() const { return impl_->exportTensor(ExportID::Action, TensorElementType::Float32, diff --git a/src/mgr.hpp b/src/mgr.hpp index d12b3a23..10a04927 100755 --- a/src/mgr.hpp +++ b/src/mgr.hpp @@ -78,6 +78,9 @@ class Manager { float acceleration, float steering, float headAngle); MGR_EXPORT void setMaps(const std::vector &maps); + + MGR_EXPORT void setToDelete(const std::unordered_map> &agentsToDelete); + // TODO: remove parameters MGR_EXPORT std::vector getShapeTensorFromDeviceMemory(); diff --git a/src/sim.cpp b/src/sim.cpp index f19ffb36..02f7bdf3 100755 --- a/src/sim.cpp +++ b/src/sim.cpp @@ -74,6 +74,7 @@ void Sim::registerTypes(ECSRegistry ®istry, const Config &cfg) registry.exportSingleton((uint32_t)ExportID::Map); registry.exportSingleton((uint32_t)ExportID::ResetMap); registry.exportSingleton((uint32_t)ExportID::WorldMeans); + registry.exportSingleton((uint32_t)ExportID::DeleteAgents); registry.exportColumn( (uint32_t)ExportID::Action); diff --git a/src/sim.hpp b/src/sim.hpp index 33c41170..4473343c 100755 --- a/src/sim.hpp +++ b/src/sim.hpp @@ -37,6 +37,7 @@ enum class ExportID : uint32_t { ResetMap, WorldMeans, MetaData, + DeleteAgents, NumExports }; diff --git a/src/types.hpp b/src/types.hpp index 12cf84c6..4a45ddf9 100755 --- a/src/types.hpp +++ b/src/types.hpp @@ -94,6 +94,10 @@ namespace gpudrive int32_t reset; }; + struct DeleteAgents { + int32_t deleteAgents[consts::kMaxAgentCount]; + }; + struct WorldMeans { madrona::math::Vector3 mean; // TODO: Z is 0 for now, but can be used for 3D in future }; From 25fe4e882915130e3ff8a40829713f6e8625e011 Mon Sep 17 00:00:00 2001 From: Aarav Pandya Date: Fri, 17 Jan 2025 22:51:43 -0500 Subject: [PATCH 2/4] Delete agents --- src/bindings.cpp | 15 ++++++++++++++- src/level_gen.cpp | 9 +++++++++ src/mgr.cpp | 25 +++++++++++++++++++++---- src/mgr.hpp | 3 ++- src/sim.cpp | 3 ++- src/sim.hpp | 2 +- src/types.hpp | 4 ++-- 7 files changed, 51 insertions(+), 10 deletions(-) diff --git a/src/bindings.cpp b/src/bindings.cpp index 37224813..36b2c7d7 100755 --- 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> agents_to_delete; + + // Convert Python dict to C++ unordered_map + for (auto item : py_agents_to_delete) { + int32_t key = nb::cast(item.first); + std::vector value = nb::cast>(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 index ff6ed492..0e5e7c7f 100755 --- 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; + 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 index 54de9df2..599ac58b 100755 --- a/src/mgr.cpp +++ b/src/mgr.cpp @@ -643,13 +643,27 @@ void Manager::setMaps(const std::vector &maps) reset(worldIndices); } -void Manager::setToDelete(const std::unordered_map> &agentsToDelete) +Tensor Manager::deletedAgentsTensor() const { + return impl_->exportTensor(ExportID::DeletedAgents, TensorElementType::Int32, + { + impl_->numWorlds, + consts::kMaxAgentCount, + }); +} + +void Manager::deleteAgents(const std::unordered_map> &agentsToDelete) +{ + + ResetMap resetmap{ + 1, + }; + if (impl_->cfg.execMode == madrona::ExecMode::CUDA) { #ifdef MADRONA_CUDA_SUPPORT auto &gpu_exec = static_cast(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 { 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 else { auto &cpu_exec = static_cast(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 { 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 worldIndices(impl_->cfg.scenes.size()); std::iota(worldIndices.begin(), worldIndices.end(), 0); reset(worldIndices); - } diff --git a/src/mgr.hpp b/src/mgr.hpp index 10a04927..b9a80d71 100755 --- 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 &maps); - MGR_EXPORT void setToDelete(const std::unordered_map> &agentsToDelete); + MGR_EXPORT void deleteAgents(const std::unordered_map> &agentsToDelete); // TODO: remove parameters MGR_EXPORT std::vector diff --git a/src/sim.cpp b/src/sim.cpp index 02f7bdf3..cc81e64a 100755 --- a/src/sim.cpp +++ b/src/sim.cpp @@ -62,6 +62,7 @@ void Sim::registerTypes(ECSRegistry ®istry, const Config &cfg) registry.registerSingleton(); registry.registerSingleton(); registry.registerSingleton(); + registry.registerSingleton(); registry.registerArchetype(); registry.registerArchetype(); @@ -74,7 +75,7 @@ void Sim::registerTypes(ECSRegistry ®istry, const Config &cfg) registry.exportSingleton((uint32_t)ExportID::Map); registry.exportSingleton((uint32_t)ExportID::ResetMap); registry.exportSingleton((uint32_t)ExportID::WorldMeans); - registry.exportSingleton((uint32_t)ExportID::DeleteAgents); + registry.exportSingleton((uint32_t)ExportID::DeletedAgents); registry.exportColumn( (uint32_t)ExportID::Action); diff --git a/src/sim.hpp b/src/sim.hpp index 4473343c..4427af8f 100755 --- 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 index 4a45ddf9..7e5ffd78 100755 --- 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 { From 4e894d2ff44e63d7f6ba16b6de1251d1733997ed Mon Sep 17 00:00:00 2001 From: Aarav Pandya Date: Fri, 17 Jan 2025 22:54:04 -0500 Subject: [PATCH 3/4] Set default deleted agents idx to -1 --- src/sim.cpp | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/sim.cpp b/src/sim.cpp index cc81e64a..f2b777b8 100755 --- a/src/sim.cpp +++ b/src/sim.cpp @@ -882,6 +882,11 @@ Sim::Sim(Engine &ctx, auto& map = ctx.singleton(); map = *(init.map); + + auto& deletedAgents = ctx.singleton(); + for (auto i = 0; i < consts::kMaxAgentCount; i++) { + deletedAgents.deletedAgents[i] = -1; + } // Creates agents, walls, etc. createPersistentEntities(ctx); From 449a7cbc3dd94849aededc8f2b3f99a059ed3f9f Mon Sep 17 00:00:00 2001 From: Aarav Pandya Date: Fri, 17 Jan 2025 23:05:53 -0500 Subject: [PATCH 4/4] Reset deleted agents in setmaps --- src/mgr.cpp | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/src/mgr.cpp b/src/mgr.cpp index 599ac58b..f041d6c0 100755 --- a/src/mgr.cpp +++ b/src/mgr.cpp @@ -610,6 +610,11 @@ void Manager::setMaps(const std::vector &maps) auto resetMapPtr = (ResetMap *)gpu_exec.getExported((uint32_t)ExportID::ResetMap) + world_idx; REQ_CUDA(cudaMemcpy(resetMapPtr, &resetmap, sizeof(ResetMap), cudaMemcpyHostToDevice)); + + // reset agents to delete + auto agentsToDeleteDevicePtr = (int32_t *)gpu_exec.getExported((uint32_t)ExportID::DeletedAgents); + int32_t *agentsToDeletePtr = agentsToDeleteDevicePtr + world_idx * consts::kMaxAgentCount; + REQ_CUDA(cudaMemset(agentsToDeletePtr, -1, consts::kMaxAgentCount * sizeof(int32_t))); } #else @@ -634,6 +639,11 @@ void Manager::setMaps(const std::vector &maps) auto resetMapPtr = (ResetMap *)cpu_exec.getExported((uint32_t)ExportID::ResetMap) + world_idx; memcpy(resetMapPtr, &resetmap, sizeof(ResetMap)); + + // reset agents to delete + auto agentsToDeleteDevicePtr = (int32_t *)cpu_exec.getExported((uint32_t)ExportID::DeletedAgents); + int32_t *agentsToDeletePtr = agentsToDeleteDevicePtr + world_idx * consts::kMaxAgentCount; + memset(agentsToDeletePtr, -1, consts::kMaxAgentCount * sizeof(int32_t)); } }