Dynamic files world reset

pygpudrive/env/base_env.py

@@ -1,4 +1,5 @@
 import os
+from typing import List, Optional
 import gymnasium as gym
 from pygpudrive.env.config import RenderConfig, RenderMode
 from pygpudrive.env.viz import PyGameVisualizer
@@ -57,7 +58,10 @@ def _set_reward_params(self):
         """
         reward_params = gpudrive.RewardParams()
 
-        if self.config.reward_type == "sparse_on_goal_achieved" or self.config.reward_type == "weighted_combination":
+        if (
+            self.config.reward_type == "sparse_on_goal_achieved"
+            or self.config.reward_type == "weighted_combination"
+        ):
             reward_params.rewardType = gpudrive.RewardType.OnGoalAchieved
         else:
             raise ValueError(f"Invalid reward type: {self.config.reward_type}")
@@ -259,6 +263,28 @@ def render(self, world_render_idx=0, color_objects_by_actor=None):
         }:
             return self.visualizer.getRender()
 
+        def resample_scenarios(self, dataset: Optional[List[str]] = None):
+
+        """Resample the scenes."""
+
+        # Sample a set of scenarios
+        if dataset is None:
+            dataset = select_scenes(self.scene_config)
+
+        print(f"Resampling {len(dataset)} scenarios... This may take a while.")
+
+        # Resample the scenes
+        self.sim.set_maps(dataset)
+
+        print(f"Resampled {len(dataset)} scenarios.")
+
+        # Re-initialize the controlled agents mask
+        self.cont_agent_mask = self.get_controlled_agents_mask()
+        self.max_agent_count = self.cont_agent_mask.shape[1]
+        self.num_valid_controlled_agents_across_worlds = (
+            self.cont_agent_mask.sum().item()
+        )
+
     def close(self):
         """Destroy the simulator and visualizer."""
         del self.sim

pygpudrive/env/config.py

@@ -91,7 +91,7 @@ class EnvConfig:
 
     # Reward settings
     reward_type: str = (
-        "sparse_on_goal_achieved" # Alternatively, "weighted_combination" 
+        "sparse_on_goal_achieved"  # Alternatively, "weighted_combination"
     )
 
     dist_to_goal_threshold: float = (

pygpudrive/env/env_torch.py

@@ -27,6 +27,7 @@ def __init__(
     ):
         # Initialization of environment configurations
         self.config = config
+        self.scene_config = scene_config
         self.num_worlds = scene_config.num_scenes
         self.max_cont_agents = max_cont_agents
         self.device = device

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/headless.cpp

@@ -34,9 +34,9 @@ int main(int argc, char *argv[])
     }
 
     uint64_t num_steps = std::stoul(argv[2]);
-    std::vector<std::string> scenes = {"../data/examples/tfrecord-00001-of-01000_307.json",
-         "../data/examples/tfrecord-00003-of-01000_109.json",
-         "../data/examples/tfrecord-00012-of-01000_389.json"};
+    std::vector<std::string> scenes = {"../data/processed/examples/tfrecord-00001-of-01000_307.json",
+         "../data/processed/examples/tfrecord-00003-of-01000_109.json",
+         "../data/processed/examples/tfrecord-00012-of-01000_389.json"};
     uint64_t num_worlds = scenes.size();
 
     bool rand_actions = false;

src/init.hpp

@@ -6,7 +6,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

@@ -309,6 +309,8 @@ void createCameraEntity(Engine &ctx)
         camera,
         150.f, 0.001f,
         1.5f * math::up);
+
+    ctx.data().camera_agent = camera;
 }
 
 static inline bool shouldAgentBeCreated(Engine &ctx, const MapObject &agentInit)
@@ -326,22 +328,25 @@ static inline bool shouldAgentBeCreated(Engine &ctx, const MapObject &agentInit)
     return true;
 }
 
-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 && agentIdx < consts::kMaxAgentCount; ++agentCtr) {
-        const auto &agentInit = map->objects[agentCtr];
+    for (CountT agentCtr = 0; agentCtr < map.numObjects && agentIdx < consts::kMaxAgentCount; ++agentCtr) {
+        const auto &agentInit = map.objects[agentCtr];
 
         if (not shouldAgentBeCreated(ctx, agentInit))
         {
@@ -356,9 +361,9 @@ void createPersistentEntities(Engine &ctx, Map *map) {
     ctx.data().numAgents = agentIdx;
 
     CountT roadIdx = 0;
-    for(CountT roadCtr = 0; roadCtr < map->numRoads && roadIdx < consts::kMaxRoadEntityCount; roadCtr++)
+    for(CountT roadCtr = 0; roadCtr < map.numRoads && roadIdx < consts::kMaxRoadEntityCount; roadCtr++)
     {
-        const auto &roadInit = map->roads[roadCtr];
+        const auto &roadInit = map.roads[roadCtr];
         createRoadEntities(ctx, roadInit, roadIdx);
     }
     ctx.data().numRoads = roadIdx;
@@ -404,7 +409,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

@@ -6,11 +6,12 @@
 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);
 
     static inline Action getZeroAction(DynamicsModel model)
     {

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();

src/sim.cpp

@@ -55,6 +55,8 @@ void Sim::registerTypes(ECSRegistry &registry, const Config &cfg)
     registry.registerComponent<AgentID>();
     registry.registerSingleton<WorldReset>();
     registry.registerSingleton<Shape>();
+    registry.registerSingleton<Map>();
+    registry.registerSingleton<ResetMap>();
 
     registry.registerArchetype<Agent>();
     registry.registerArchetype<PhysicsEntity>();
@@ -64,6 +66,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>(
@@ -96,7 +100,7 @@ void Sim::registerTypes(ECSRegistry &registry, const Config &cfg)
 }
 
 static inline void cleanupWorld(Engine &ctx) {
-    // TODO: Implement cleanup for changing worlds during runtime
+    destroyWorld(ctx);
 }
 
 static inline void initWorld(Engine &ctx)
@@ -109,22 +113,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);
 }
 
@@ -838,7 +855,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,
@@ -850,8 +868,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

@@ -61,7 +61,10 @@ struct AgentID {
         int32_t reset;
     };
 
-    struct ClassicAction
+struct ResetMap {
+    int32_t reset;
+};   
+     struct ClassicAction
     {
         float acceleration;
         float steering;