Dynamic files world reset

src/bindings.cpp

@@ -124,7 +124,8 @@ namespace gpudrive
             .def("rgb_tensor", &Manager::rgbTensor)
             .def("depth_tensor", &Manager::depthTensor)
             .def("response_type_tensor", &Manager::responseTypeTensor)
-            .def("expert_trajectory_tensor", &Manager::expertTrajectoryTensor);
+            .def("expert_trajectory_tensor", &Manager::expertTrajectoryTensor)
+            .def("set_maps", &Manager::setMaps);
     }
 
 }

src/init.hpp

@@ -7,7 +7,7 @@ namespace gpudrive
 {
 
     // Constants computed from train files.
-    constexpr size_t MAX_OBJECTS = 515;
+    constexpr size_t MAX_OBJECTS = consts::kMaxAgentCount;
     constexpr size_t MAX_ROADS = 956;
     constexpr size_t MAX_POSITIONS = 91;
     constexpr size_t MAX_GEOMETRY = 1746;

src/level_gen.cpp

@@ -394,33 +394,38 @@ void createCameraEntity(Engine &ctx)
         camera,
         150.f, 0.001f,
         1.5f * math::up);
+
+    ctx.data().camera_agent = camera;
 }
 
-void createPersistentEntities(Engine &ctx, Map *map) {
+void createPersistentEntities(Engine &ctx) {
     // createFloorPlane(ctx);
 
+    const auto& map = ctx.singleton<Map>();
+
     if (ctx.data().enableRender)
     {
         createCameraEntity(ctx);
     }
 
     ctx.data().mean = {0, 0};
-    ctx.data().mean.x = map->mean.x;
-    ctx.data().mean.y = map->mean.y;
+    ctx.data().mean.x = map.mean.x;
+    ctx.data().mean.y = map.mean.y;
     ctx.data().numControlledAgents = 0;
+    ctx.singleton<ResetMap>().reset = 0;
 
     CountT agentIdx = 0;
-    for (CountT agentCtr = 0; agentCtr < map->numObjects; ++agentCtr) {
+    for (CountT agentCtr = 0; agentCtr < map.numObjects; ++agentCtr) {
         if(agentIdx >= consts::kMaxAgentCount)
             break;
         if (ctx.data().params.IgnoreNonVehicles)
         {
-            if (map->objects[agentCtr].type == EntityType::Pedestrian || map->objects[agentCtr].type == EntityType::Cyclist)
+            if (map.objects[agentCtr].type == EntityType::Pedestrian || map.objects[agentCtr].type == EntityType::Cyclist)
             {
                 continue;
             }
         }
-        const auto &agentInit = map->objects[agentCtr];
+        const auto &agentInit = map.objects[agentCtr];
         if (ctx.data().params.initOnlyValidAgentsAtFirstStep && agentInit.valid[0] == false)
         {
             continue;
@@ -434,11 +439,11 @@ void createPersistentEntities(Engine &ctx, Map *map) {
     ctx.data().numAgents = agentIdx;
 
     CountT roadIdx = 0;
-    for(CountT roadCtr = 0; roadCtr < map->numRoads; roadCtr++)
+    for(CountT roadCtr = 0; roadCtr < map.numRoads; roadCtr++)
     {
         if(roadIdx >= consts::kMaxRoadEntityCount)
             break;
-        const auto &roadInit = map->roads[roadCtr];
+        const auto &roadInit = map.roads[roadCtr];
         createRoadEntities(ctx, roadInit, roadIdx);
     }
     ctx.data().numRoads = roadIdx;
@@ -494,7 +499,40 @@ static void resetPersistentEntities(Engine &ctx)
     }
 }
 
-void generateWorld(Engine &ctx)
+void destroyWorld(Engine &ctx)
+{
+    for (CountT idx = 0; idx < ctx.data().numAgents; ++idx)
+    {
+        Entity agent = ctx.data().agents[idx];
+        ctx.destroyRenderableEntity(agent);
+    }
+    for (CountT idx = 0; idx < ctx.data().numRoads; idx++)
+    {
+        Entity road = ctx.data().roads[idx];
+        ctx.destroyRenderableEntity(road);
+    }
+    if (ctx.data().enableRender)
+    {
+        ctx.destroyRenderableEntity(ctx.data().camera_agent);
+    }
+    for (CountT idx = 0; idx < consts::kMaxAgentCount; ++idx)
+    {
+        Entity agent_iface = ctx.data().agent_ifaces[idx];
+        ctx.destroyEntity(agent_iface);
+    }
+    for (CountT idx = 0; idx < consts::kMaxRoadEntityCount; ++idx)
+    {
+        Entity road_iface = ctx.data().road_ifaces[idx];
+        ctx.destroyEntity(road_iface);
+    }
+    ctx.data().numAgents = 0;
+    ctx.data().numRoads = 0;
+    ctx.data().numControlledAgents = 0;
+    ctx.data().mean = {0, 0};
+}
+
+
+void resetWorld(Engine &ctx)
 {
     resetPersistentEntities(ctx);
 }

src/level_gen.hpp

@@ -4,10 +4,11 @@
 
 namespace gpudrive {
 
-void createPersistentEntities(Engine &ctx, Map *map);
+void createPersistentEntities(Engine &ctx);
 
-// First, destroys any non-persistent state for the current world and then
-// generates a new play area.
-void generateWorld(Engine &ctx);
+void resetWorld(Engine &ctx);
+
+// Destroys all entities in the world
+void destroyWorld(Engine &ctx);
 
 }

src/mgr.cpp

@@ -595,6 +595,62 @@ void Manager::reset(std::vector<int32_t> worldsToReset) {
     }
 }
 
+void Manager::setMaps(const std::vector<std::string> &maps)
+{
+    assert(impl_->cfg.scenes.size() == maps.size());
+    impl_->cfg.scenes = maps;
+
+    ResetMap resetmap{
+        1,
+    };
+
+    if (impl_->cfg.execMode == madrona::ExecMode::CUDA)
+    {
+#ifdef MADRONA_CUDA_SUPPORT
+        auto &gpu_exec = static_cast<CUDAImpl *>(impl_.get())->gpuExec;
+        for (size_t world_idx = 0; world_idx < maps.size(); world_idx++)
+        {
+            Map *map = static_cast<Map *>(MapReader::parseAndWriteOut(maps[world_idx],
+                                                                      ExecMode::CUDA, impl_->cfg.params.polylineReductionThreshold));
+            Map *mapDevicePtr = (Map *)gpu_exec.getExported((uint32_t)ExportID::Map) + world_idx;
+            REQ_CUDA(cudaMemcpy(mapDevicePtr, map, sizeof(Map), cudaMemcpyHostToDevice));
+            madrona::cu::deallocGPU(map);
+
+            auto resetMapPtr = (ResetMap *)gpu_exec.getExported((uint32_t)ExportID::ResetMap) + world_idx;
+            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<CPUImpl *>(impl_.get())->cpuExec;
+
+        for (size_t world_idx = 0; world_idx < maps.size(); world_idx++)
+        {
+            // Parse the map string into your MapData structure
+            Map *map = static_cast<Map *>(MapReader::parseAndWriteOut(maps[world_idx],
+                                                                      ExecMode::CPU, impl_->cfg.params.polylineReductionThreshold));
+
+            Map *mapDevicePtr = (Map *)cpu_exec.getExported((uint32_t)ExportID::Map) + world_idx;
+            memcpy(mapDevicePtr, map, sizeof(Map));
+            delete map;
+
+            auto resetMapPtr = (ResetMap *)cpu_exec.getExported((uint32_t)ExportID::ResetMap) + world_idx;
+            memcpy(resetMapPtr, &resetmap, sizeof(ResetMap));
+        }
+    }
+
+    // Vector of range on integers from 0 to the number of worlds
+    std::vector<int32_t> worldIndices(maps.size());
+    std::iota(worldIndices.begin(), worldIndices.end(), 0);
+    reset(worldIndices);
+}
+
 Tensor Manager::actionTensor() const
 {
     return impl_->exportTensor(ExportID::Action, TensorElementType::Float32,

src/mgr.hpp

@@ -75,6 +75,7 @@ class Manager {
     MGR_EXPORT void setAction(int32_t world_idx, int32_t agent_idx,
                               float acceleration, float steering,
                               float headAngle);
+    MGR_EXPORT void setMaps(const std::vector<std::string> &maps);
     // TODO: remove parameters
     MGR_EXPORT std::vector<Shape>
     getShapeTensorFromDeviceMemory(madrona::ExecMode mode);

src/sim.cpp

@@ -54,6 +54,8 @@ void Sim::registerTypes(ECSRegistry &registry, const Config &cfg)
     registry.registerComponent<RoadInterfaceEntity>();
     registry.registerSingleton<WorldReset>();
     registry.registerSingleton<Shape>();
+    registry.registerSingleton<Map>();
+    registry.registerSingleton<ResetMap>();
 
     registry.registerArchetype<Agent>();
     registry.registerArchetype<PhysicsEntity>();
@@ -63,6 +65,8 @@ void Sim::registerTypes(ECSRegistry &registry, const Config &cfg)
 
     registry.exportSingleton<WorldReset>((uint32_t)ExportID::Reset);
     registry.exportSingleton<Shape>((uint32_t)ExportID::Shape);
+    registry.exportSingleton<Map>((uint32_t)ExportID::Map);
+    registry.exportSingleton<ResetMap>((uint32_t)ExportID::ResetMap);
     registry.exportColumn<AgentInterface, Action>(
         (uint32_t)ExportID::Action);
     registry.exportColumn<AgentInterface, SelfObservation>(
@@ -94,7 +98,9 @@ void Sim::registerTypes(ECSRegistry &registry, const Config &cfg)
         (uint32_t)ExportID::Trajectory);
 }
 
-static inline void cleanupWorld(Engine &ctx) {}
+static inline void cleanupWorld(Engine &ctx) {
+    destroyWorld(ctx);
+}
 
 static inline void initWorld(Engine &ctx)
 {
@@ -106,22 +112,35 @@ static inline void initWorld(Engine &ctx)
     ctx.data().rng = RNG::make(episode_idx);
     ctx.data().curEpisodeIdx = episode_idx;
 
+    if(ctx.singleton<ResetMap>().reset == 1)
+    {
+        createPersistentEntities(ctx);
+        ctx.singleton<ResetMap>().reset = 0;
+        phys::PhysicsSystem::reset(ctx);
+    }
+
     // Defined in src/level_gen.hpp / src/level_gen.cpp
-    generateWorld(ctx);
+    resetWorld(ctx);
 }
 
 // This system runs in TaskGraphID::Reset and checks if the code external to the
 // application has forced a reset by writing to the WorldReset singleton. If a
 // reset is needed, cleanup the existing world and generate a new one.
 inline void resetSystem(Engine &ctx, WorldReset &reset)
 {
-    if (reset.reset == 0) {
-      return;
+    if (reset.reset == 0)
+    {
+        return;
     }
 
     reset.reset = 0;
 
-    cleanupWorld(ctx);
+    auto resetMap = ctx.singleton<ResetMap>();
+
+    if (resetMap.reset == 1)
+    {
+        cleanupWorld(ctx);
+    }
     initWorld(ctx);
 }
 
@@ -867,7 +886,8 @@ Sim::Sim(Engine &ctx,
     // Currently the physics system needs an upper bound on the number of
     // entities that will be stored in the BVH. We plan to fix this in
     // a future release.
-    auto max_total_entities = init.map->numObjects + init.map->numRoadSegments;
+    // auto max_total_entities = init.map->numObjects + init.map->numRoadSegments;
+    auto max_total_entities = consts::kMaxAgentCount + consts::kMaxRoadEntityCount;
 
     phys::PhysicsSystem::init(ctx, init.rigidBodyObjMgr,
         consts::deltaT, consts::numPhysicsSubsteps, -9.8f * math::up,
@@ -879,8 +899,9 @@ Sim::Sim(Engine &ctx,
         RenderingSystem::init(ctx, cfg.renderBridge);
     }
 
+    ctx.singleton<Map>() = *(init.map);
     // Creates agents, walls, etc.
-    createPersistentEntities(ctx, init.map);
+    createPersistentEntities(ctx);
 
     // Generate initial world state
     initWorld(ctx);

src/sim.hpp

@@ -33,6 +33,8 @@ enum class ExportID : uint32_t {
     Info,
     ResponseType,
     Trajectory,
+    Map,
+    ResetMap,
     NumExports
 };
 
@@ -124,6 +126,8 @@ struct Sim : public madrona::WorldBase {
     Entity agent_ifaces[consts::kMaxAgentCount];
     Entity road_ifaces[consts::kMaxRoadEntityCount];
 
+    Entity camera_agent;
+
     madrona::CountT numControlledAgents;
 
     madrona::math::Vector2 mean;

src/types.hpp

@@ -53,7 +53,9 @@ struct WorldReset {
     int32_t reset;
 };
 
-
+struct ResetMap {
+    int32_t reset;
+};
 
 struct ClassicAction {
     float acceleration;