From aba4b6fd5ab7dfa2f2287e2c296313522e1b7599 Mon Sep 17 00:00:00 2001 From: Kyungbeom Kim <69628911+KKGB@users.noreply.github.com> Date: Fri, 24 Jan 2025 01:29:23 +0900 Subject: [PATCH 1/4] Add plot expert trajectories (#325) * Add: plot expert trajectories * Refactor: rename draw_only_controllable_veh variable --- pygpudrive/env/config.py | 4 +++ pygpudrive/env/env_torch.py | 37 +++++++++++++++--------- pygpudrive/visualize/core.py | 54 ++++++++++++++++++++++++++++++++++++ 3 files changed, 82 insertions(+), 13 deletions(-) diff --git a/pygpudrive/env/config.py b/pygpudrive/env/config.py index c094aac7..21111695 100755 --- a/pygpudrive/env/config.py +++ b/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" diff --git a/pygpudrive/env/env_torch.py b/pygpudrive/env/env_torch.py index ddd266c9..224d61f2 100755 --- a/pygpudrive/env/env_torch.py +++ b/pygpudrive/env/env_torch.py @@ -489,44 +489,55 @@ def get_expert_actions(self): 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() diff --git a/pygpudrive/visualize/core.py b/pygpudrive/visualize/core.py index 61aa4ed0..45b0b4e5 100644 --- a/pygpudrive/visualize/core.py +++ b/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, From e7565aea5cb00fe31fd738dcd0cdb82f09e39045 Mon Sep 17 00:00:00 2001 From: Aarav Pandya Date: Sat, 25 Jan 2025 15:52:03 -0500 Subject: [PATCH 2/4] Export map names --- src/bindings.cpp | 3 ++- src/init.hpp | 2 ++ src/json_serialization.hpp | 3 +++ src/mgr.cpp | 7 +++++++ src/mgr.hpp | 1 + src/sim.cpp | 7 +++++++ src/sim.hpp | 1 + src/types.hpp | 8 ++++++++ 8 files changed, 31 insertions(+), 1 deletion(-) diff --git a/src/bindings.cpp b/src/bindings.cpp index 37224813..9e788cf0 100755 --- a/src/bindings.cpp +++ b/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); } } diff --git a/src/init.hpp b/src/init.hpp index 26f83dfc..c528d937 100755 --- a/src/init.hpp +++ b/src/init.hpp @@ -60,6 +60,8 @@ namespace gpudrive uint32_t numRoadSegments; MapVector2 mean; + char mapName[32]; + // Constructor Map() = default; }; diff --git a/src/json_serialization.hpp b/src/json_serialization.hpp index cd46f737..f1fd8a12 100644 --- a/src/json_serialization.hpp +++ b/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::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(MAX_OBJECTS)); diff --git a/src/mgr.cpp b/src/mgr.cpp index 96d45c91..29693a18 100755 --- a/src/mgr.cpp +++ b/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, diff --git a/src/mgr.hpp b/src/mgr.hpp index d12b3a23..95ccd74f 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 mapNameTensor() const; madrona::py::Tensor rgbTensor() const; madrona::py::Tensor depthTensor() const; // These functions are used by the viewer to control the simulation diff --git a/src/sim.cpp b/src/sim.cpp index f19ffb36..d874e2a2 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::MapName); registry.exportColumn( (uint32_t)ExportID::Action); @@ -878,6 +880,11 @@ Sim::Sim(Engine &ctx, RenderingSystem::init(ctx, cfg.renderBridge); } + auto& mapName = ctx.singleton(); + for (int i = 0; i < sizeof(mapName.mapName); i++) { + mapName.mapName[i] = init.map->mapName[i]; + } + auto& map = ctx.singleton(); map = *(init.map); // Creates agents, walls, etc. diff --git a/src/sim.hpp b/src/sim.hpp index 33c41170..0168f1fd 100755 --- a/src/sim.hpp +++ b/src/sim.hpp @@ -37,6 +37,7 @@ enum class ExportID : uint32_t { ResetMap, WorldMeans, MetaData, + MapName, NumExports }; diff --git a/src/types.hpp b/src/types.hpp index 12cf84c6..c9ae827d 100755 --- a/src/types.hpp +++ b/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 { From dafece731ebacdafbcdc68ffd21075486bfe0b5b Mon Sep 17 00:00:00 2001 From: Aarav Pandya Date: Sat, 25 Jan 2025 15:57:15 -0500 Subject: [PATCH 3/4] Handle set maps case --- src/level_gen.cpp | 7 ++++++- src/sim.cpp | 5 ----- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/src/level_gen.cpp b/src/level_gen.cpp index ff6ed492..c9fae877 100755 --- a/src/level_gen.cpp +++ b/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(); + auto& mapName = ctx.singleton(); + for (int i = 0; i < sizeof(mapName.mapName); i++) { + mapName.mapName[i] = map.mapName[i]; + } + + if (ctx.data().enableRender) { createCameraEntity(ctx); diff --git a/src/sim.cpp b/src/sim.cpp index d874e2a2..2ba1f0c0 100755 --- a/src/sim.cpp +++ b/src/sim.cpp @@ -880,11 +880,6 @@ Sim::Sim(Engine &ctx, RenderingSystem::init(ctx, cfg.renderBridge); } - auto& mapName = ctx.singleton(); - for (int i = 0; i < sizeof(mapName.mapName); i++) { - mapName.mapName[i] = init.map->mapName[i]; - } - auto& map = ctx.singleton(); map = *(init.map); // Creates agents, walls, etc. From d6233717fc6415f7252a5f842a14822bb3e5c104 Mon Sep 17 00:00:00 2001 From: Daphne Cornelisse Date: Sun, 26 Jan 2025 12:54:28 -0500 Subject: [PATCH 4/4] Add env helper func to extract filenames from worlds --- pygpudrive/env/env_torch.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/pygpudrive/env/env_torch.py b/pygpudrive/env/env_torch.py index 224d61f2..a7821462 100755 --- a/pygpudrive/env/env_torch.py +++ b/pygpudrive/env/env_torch.py @@ -486,6 +486,25 @@ 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__": @@ -515,6 +534,7 @@ def get_expert_actions(self): # RUN obs = env.reset() + frames = {f"env_{i}": [] for i in range(NUM_WORLDS)} for t in range(TOTAL_STEPS):