Skip to content

Commit

Permalink
merged map names export
Browse files Browse the repository at this point in the history
  • Loading branch information
nadarenator committed Jan 26, 2025
2 parents 1f8d460 + d623371 commit e76201a
Show file tree
Hide file tree
Showing 12 changed files with 99 additions and 3 deletions.
2 changes: 2 additions & 0 deletions pygpudrive/env/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -187,4 +187,6 @@ class RenderConfig:
view_option: Enum = None
resolution: Tuple[int, int] = (1024, 1024)
draw_obj_idx: bool = False
draw_expert_trajectories: bool = False
draw_only_controllable_veh: bool = False
obj_idx_font_size: int = 9
23 changes: 21 additions & 2 deletions pygpudrive/env/env_torch.py
Original file line number Diff line number Diff line change
Expand Up @@ -647,12 +647,30 @@ 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
import mediapy as media

env_config = EnvConfig(dynamics_model="delta_local")
render_config = RenderConfig()
Expand Down Expand Up @@ -700,6 +718,7 @@ def get_expert_actions(self):
print(f"Step: {t}")

# Step the environment
expert_actions, _, _, _ = env.get_expert_actions()
env.step_dynamics(expert_actions[:, :, t, :])

# if (t + 1) % 2 == 0:
Expand Down
45 changes: 45 additions & 0 deletions pygpudrive/visualize/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -130,6 +130,13 @@ def plot_simulator_state(
agent_infos = (
self.sim_object.info_tensor().to_torch().clone().to(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)

figs = []

Expand Down Expand Up @@ -675,6 +682,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,
Expand Down
1 change: 1 addition & 0 deletions src/bindings.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -141,6 +141,7 @@ namespace gpudrive
self.deleteAgents(agents_to_delete);
})
.def("deleted_agents_tensor", &Manager::deletedAgentsTensor);
.def("map_name_tensor", &Manager::mapNameTensor);
}

}
2 changes: 2 additions & 0 deletions src/init.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,8 @@ namespace gpudrive
uint32_t numRoadSegments;
MapVector2 mean;

char mapName[32];

// Constructor
Map() = default;
};
Expand Down
3 changes: 3 additions & 0 deletions src/json_serialization.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -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));
Expand Down
7 changes: 6 additions & 1 deletion src/level_gen.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -361,9 +361,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);
Expand Down
7 changes: 7 additions & 0 deletions src/mgr.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -868,6 +868,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,
Expand Down
1 change: 1 addition & 0 deletions src/mgr.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,7 @@ class Manager {
MGR_EXPORT madrona::py::Tensor worldMeansTensor() const;
MGR_EXPORT madrona::py::Tensor metadataTensor() const;
MGR_EXPORT madrona::py::Tensor deletedAgentsTensor() 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
Expand Down
2 changes: 2 additions & 0 deletions src/sim.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,7 @@ void Sim::registerTypes(ECSRegistry &registry, const Config &cfg)
registry.registerSingleton<ResetMap>();
registry.registerSingleton<WorldMeans>();
registry.registerSingleton<DeletedAgents>();
registry.registerSingleton<MapName>();

registry.registerArchetype<Agent>();
registry.registerArchetype<PhysicsEntity>();
Expand All @@ -76,6 +77,7 @@ void Sim::registerTypes(ECSRegistry &registry, const Config &cfg)
registry.exportSingleton<ResetMap>((uint32_t)ExportID::ResetMap);
registry.exportSingleton<WorldMeans>((uint32_t)ExportID::WorldMeans);
registry.exportSingleton<DeletedAgents>((uint32_t)ExportID::DeletedAgents);
registry.exportSingleton<MapName>((uint32_t)ExportID::MapName);

registry.exportColumn<AgentInterface, Action>(
(uint32_t)ExportID::Action);
Expand Down
1 change: 1 addition & 0 deletions src/sim.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@ enum class ExportID : uint32_t {
WorldMeans,
MetaData,
DeletedAgents,
MapName,
NumExports
};

Expand Down
8 changes: 8 additions & 0 deletions src/types.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -381,6 +381,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
{
Expand Down

0 comments on commit e76201a

Please sign in to comment.