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 96d45c91..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,15 +639,82 @@ 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)); } } // 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); +} + +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::DeletedAgents); + 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)); + } + 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 + 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::DeletedAgents); + 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)); + } + 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); } + Tensor Manager::actionTensor() const { return impl_->exportTensor(ExportID::Action, TensorElementType::Float32, diff --git a/src/mgr.hpp b/src/mgr.hpp index d12b3a23..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 @@ -78,6 +79,9 @@ class Manager { float acceleration, float steering, float headAngle); MGR_EXPORT void setMaps(const std::vector &maps); + + MGR_EXPORT void deleteAgents(const std::unordered_map> &agentsToDelete); + // TODO: remove parameters MGR_EXPORT std::vector getShapeTensorFromDeviceMemory(); diff --git a/src/sim.cpp b/src/sim.cpp index a4aa80bd..9e0d4a64 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,6 +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::DeletedAgents); registry.exportColumn( (uint32_t)ExportID::Action); @@ -878,6 +880,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); diff --git a/src/sim.hpp b/src/sim.hpp index 33c41170..4427af8f 100755 --- a/src/sim.hpp +++ b/src/sim.hpp @@ -37,6 +37,7 @@ enum class ExportID : uint32_t { ResetMap, WorldMeans, MetaData, + DeletedAgents, NumExports }; diff --git a/src/types.hpp b/src/types.hpp index 12cf84c6..7e5ffd78 100755 --- a/src/types.hpp +++ b/src/types.hpp @@ -94,6 +94,10 @@ namespace gpudrive int32_t reset; }; + struct DeletedAgents { + int32_t deletedAgents[consts::kMaxAgentCount]; + }; + struct WorldMeans { madrona::math::Vector3 mean; // TODO: Z is 0 for now, but can be used for 3D in future };