Export map names

pygpudrive/env/config.py

@@ -181,6 +181,8 @@ class RenderConfig:
         resolution (Tuple[int, int]): Resolution of the rendered image.
         line_thickness (int): Thickness of the road lines in the rendering.
         draw_obj_idx (bool): Whether to draw object indices on objects.
+        draw_expert_trajectories (bool): Whether to draw expert trajectories.
+        draw_only_controllable_veh (bool): Whether to draw only the trajectories of controllable vehicles.
         obj_idx_font_size (int): Font size for object indices.
         color_scheme (str): Color mode for the rendering ("light" or "dark").
     """
@@ -190,6 +192,8 @@ class RenderConfig:
     resolution: Tuple[int, int] = (1024, 1024)
     line_thickness: int = 0.7
     draw_obj_idx: bool = False
+    draw_expert_trajectories: bool = False
+    draw_only_controllable_veh: bool = False
     obj_idx_font_size: int = 9
     color_scheme: str = "light"
 

pygpudrive/env/env_torch.py

@@ -486,47 +486,78 @@ def get_expert_actions(self):
             log_trajectory.vel_xy,
             log_trajectory.yaw,
         )
+
+    def get_env_filenames(self):
+        """Obtain the tfrecord filename for each world, mapping world indices to map names."""
+
+        map_name_integers = self.sim.map_name_tensor().to_torch()
+
+        filenames = {}
+
+        # Iterate through the number of worlds
+        for i in range(self.num_worlds):
+            tensor = map_name_integers[i]
+
+            # Convert ints to characters, ignoring zeros
+            map_name = ''.join([chr(i) for i in tensor.tolist() if i != 0])
+
+            # Map the world index to the corresponding map name
+            filenames[i] = map_name
+
+        return filenames
 
 
 if __name__ == "__main__":
+    import mediapy
+    from pygpudrive.visualize.utils import img_from_fig
 
     # CONFIGURE
     TOTAL_STEPS = 90
-    MAX_CONTROLLED_AGENTS = 32
-    NUM_WORLDS = 1
+    MAX_CONTROLLED_AGENTS = 2
+    NUM_WORLDS = 2
 
     env_config = EnvConfig(dynamics_model="delta_local")
-    render_config = RenderConfig()
-    scene_config = SceneConfig("data/processed/training", NUM_WORLDS)
+    scene_config = SceneConfig("data/processed/examples", NUM_WORLDS)
+    render_config = RenderConfig(
+        draw_expert_trajectories=True,
+        draw_only_controllable_veh=True,
+    )
 
     # MAKE ENV
     env = GPUDriveTorchEnv(
         config=env_config,
         scene_config=scene_config,
         max_cont_agents=MAX_CONTROLLED_AGENTS,  # Number of agents to control
-        device="cpu",
-        render_config=render_config,
+        action_type="continuous",
+        device="cuda",
     )
 
     # RUN
     obs = env.reset()
-    frames = []
-
-    expert_actions, _, _, _ = env.get_expert_actions()
+
+    frames = {f"env_{i}": [] for i in range(NUM_WORLDS)}
 
     for t in range(TOTAL_STEPS):
         print(f"Step: {t}")
 
         # Step the environment
+        expert_actions, _, _, _ = env.get_expert_actions()
         env.step_dynamics(expert_actions[:, :, t, :])
 
-        frames.append(env.render())
+        figs = env.vis.plot_simulator_state(
+            env_indices=list(range(NUM_WORLDS)),
+            time_steps=[t]*NUM_WORLDS,
+            figsize=(6, 6),
+            zoom_radius=100,
+        )
+
+        for i in range(NUM_WORLDS):
+            frames[f"env_{i}"].append(img_from_fig(figs[i]))
 
         obs = env.get_obs()
         reward = env.get_rewards()
         done = env.get_dones()
 
-    # import imageio
-    imageio.mimsave("world1.gif", np.array(frames))
-
+    for i in range(NUM_WORLDS):
+        mediapy.write_video(f"world{i}.mp4", frames[f"env_{i}"], fps=5)
     env.close()

pygpudrive/visualize/core.py

@@ -17,6 +17,7 @@
     GlobalEgoState,
     PartnerObs,
 )
+from pygpudrive.datatypes.trajectory import LogTrajectory
 from pygpudrive.datatypes.control import ResponseType
 from pygpudrive.visualize.color import (
     ROAD_GRAPH_COLORS,
@@ -94,6 +95,13 @@ def plot_simulator_state(
             backend=self.backend,
             device=self.device,
         )
+        log_trajectory = LogTrajectory.from_tensor(
+            expert_traj_tensor=self.sim_object.expert_trajectory_tensor(),
+            num_worlds=len(env_indices),
+            max_agents=self.controlled_agents.shape[1],
+            backend=self.backend,
+        )
+        expert_trajectories = log_trajectory.pos_xy.to(self.device)
 
         agent_infos = self.sim_object.info_tensor().to_torch().to(self.device)
 
@@ -158,6 +166,14 @@ def plot_simulator_state(
                 marker_size_scale=marker_scale,
             )
 
+            # Draw expert trajectories
+            self._plot_expert_trajectories(
+                ax=ax,
+                env_idx=env_idx,
+                expert_trajectories=expert_trajectories,
+                response_type=response_type,
+            )
+
             # Draw the agents
             self._plot_filtered_agent_bounding_boxes(
                 ax=ax,
@@ -492,6 +508,44 @@ def _plot_filtered_agent_bounding_boxes(
             label=label,
         )
 
+    def _plot_expert_trajectories(
+        self,
+        ax: matplotlib.axes.Axes,
+        env_idx: int,
+        expert_trajectories: torch.Tensor,
+        response_type: Any,
+    ) -> None:
+        """Plot expert trajectories.
+        Args:
+            ax: Matplotlib axis for plotting.
+            env_idx: Environment index to select specific environment agents.
+            expert_trajectories: The global state of expert from `LogTrajectory`.
+        """
+        if self.vis_config.draw_expert_trajectories:
+            controlled_mask = self.controlled_agents[env_idx, :]
+            non_controlled_mask = ~response_type.static[env_idx, :] & response_type.moving[env_idx, :] & ~controlled_mask
+            mask = (
+                controlled_mask
+                if self.vis_config.draw_only_controllable_veh
+                else controlled_mask | non_controlled_mask
+            )
+            agent_indices = torch.where(mask)[0]
+            trajectories = expert_trajectories[env_idx][mask]
+            for idx, trajectory in zip(agent_indices, trajectories):
+                color = AGENT_COLOR_BY_STATE["ok"] if controlled_mask[idx] else AGENT_COLOR_BY_STATE["log_replay"]
+                for step in trajectory:
+                    x, y = step[:2].numpy()
+                    if x < OUT_OF_BOUNDS and y < OUT_OF_BOUNDS:
+                        ax.add_patch(
+                            Circle(
+                                (x, y),
+                                radius=0.3,
+                                color=color,
+                                fill=True,
+                                alpha=0.5,
+                            )
+                        )
+
     def plot_agent_observation(
         self,
         agent_idx: int,

src/bindings.cpp

@@ -127,7 +127,8 @@ 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("map_name_tensor", &Manager::mapNameTensor);
     }
 
 }

src/init.hpp

@@ -60,6 +60,8 @@ namespace gpudrive
         uint32_t numRoadSegments;
         MapVector2 mean;
 
+        char mapName[32];
+
         // Constructor
         Map() = default;
     };

src/json_serialization.hpp

@@ -279,6 +279,9 @@ namespace gpudrive
 
     void from_json(const nlohmann::json &j, Map &map, float polylineReductionThreshold)
     {
+        std::string name = j.at("name").get<std::string>();
+        std::strncpy(map.mapName, name.c_str(), sizeof(map.mapName));
+
         auto mean = calc_mean(j);
         map.mean = {mean.first, mean.second};
         map.numObjects = std::min(j.at("objects").size(), static_cast<size_t>(MAX_OBJECTS));

src/level_gen.cpp

@@ -352,9 +352,14 @@ static inline bool shouldAgentBeCreated(Engine &ctx, const MapObject &agentInit)
 
 void createPersistentEntities(Engine &ctx) {
     // createFloorPlane(ctx);
-
     const auto& map = ctx.singleton<Map>();
 
+    auto& mapName = ctx.singleton<MapName>();
+    for (int i = 0; i < sizeof(mapName.mapName); i++) {
+        mapName.mapName[i] = map.mapName[i];
+    }
+
+
     if (ctx.data().enableRender)
     {
         createCameraEntity(ctx);

src/mgr.cpp

@@ -796,6 +796,13 @@ Tensor Manager::expertTrajectoryTensor() const {
         {impl_->numWorlds, consts::kMaxAgentCount, TrajectoryExportSize});
 }
 
+Tensor Manager::mapNameTensor() const {
+    return impl_->exportTensor(
+        ExportID::MapName, TensorElementType::Int32,
+        {impl_->numWorlds, MapNameExportSize}
+    );
+}
+
 Tensor Manager::metadataTensor() const {
     return impl_->exportTensor(
         ExportID::MetaData, TensorElementType::Int32,

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 mapNameTensor() const;
     madrona::py::Tensor rgbTensor() const;
     madrona::py::Tensor depthTensor() const;
     // These functions are used by the viewer to control the simulation

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<MapName>();
 
     registry.registerArchetype<Agent>();
     registry.registerArchetype<PhysicsEntity>();
@@ -74,6 +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<MapName>((uint32_t)ExportID::MapName);
 
     registry.exportColumn<AgentInterface, Action>(
         (uint32_t)ExportID::Action);

src/sim.hpp

@@ -37,6 +37,7 @@ enum class ExportID : uint32_t {
     ResetMap,
     WorldMeans,
     MetaData,
+    MapName,
     NumExports
 };
 

src/types.hpp

@@ -377,6 +377,14 @@ namespace gpudrive
 
     static_assert(sizeof(AbsoluteSelfObservation) == sizeof(float) * AbsoluteSelfObservationExportSize);
 
+    struct MapName
+    {
+        char32_t mapName[32];
+    };
+
+    const size_t MapNameExportSize = 32;
+    static_assert(sizeof(MapName) == sizeof(char32_t) * MapNameExportSize);
+
     //Metadata struct : using agent IDs.
     struct MetaData
     {