Merge pull request #43 from inverted-ai/map-file-management

Map file management
inverted-ai · Apr 4, 2024 · ce19713 · ce19713
2 parents 6e324aa + 51dd8d5
commit ce19713
Show file tree

Hide file tree

Showing 32 changed files with 549,508 additions and 78 deletions.
diff --git a/README.md b/README.md
@@ -67,11 +67,21 @@ contacting us. For academics, we may be able to offer free API keys.
 
 ## Maps and Map Formats
 
-TorchDriveSim uses Lanelet2 as its map format, but it runs without Lanelet2 on any of the pre-defined
-maps available as meshes in this repo, although it won't be able to detect wrong way infractions and use certain
-heuristic behavioral models. To use those features, and to use TorchDriveSim with your own maps, you'll need to install
-Lanelet2 with its Python bindings. You can either use the official distribution or the fork hosted by Inverted AI,
-which allows for installation without ROS.
+Several CARLA maps (`carla_Town01`, `carla_Town02`, `carla_Town03`, `carla_Town04`, `carla_Town10HD`)
+are included in `torchdrivesim` itself and can be loaded
+by name. To include other maps, place the files in the format described below somewhere in a folder referenced by the
+`TDS_RESOURCE_PATH` environment variable. Generally, a map is defined by a folder with the following structure:
+```
+MAPNAME/
+  metadata.json  # custom metadata format
+  MAPNAME.osm  # road network in Lanelet2 format
+  MAPNAME_mesh.json  # custom road mesh format, can be derived from the .osm file
+  MAPNAME_stoplines.json  # custom format specifying traffic lights, stop signs, and yield signs, if needed
+```
+
+See the bundled maps in `torchdrivesim/resources/maps` for examples. There is currently no tooling available
+for creating TorchDriveSim-compatible maps, but you can try the experimental OpenDRIVE
+[converter](https://github.com/inverted-ai/map-converter).
 
 ## Scenario Definition
 

diff --git a/examples/extended_state.py b/examples/extended_state.py
@@ -17,6 +17,7 @@
 from torchdrivesim.behavior.heuristic import heuristic_initialize
 from torchdrivesim.kinematic import KinematicBicycle, TeleportingKinematicModel, KinematicModel
 from torchdrivesim.lanelet2 import load_lanelet_map, road_mesh_from_lanelet_map, lanelet_map_to_lane_mesh
+from torchdrivesim.map import find_map_config
 from torchdrivesim.mesh import BirdviewMesh
 from torchdrivesim.rendering import renderer_from_config, RendererConfig
 from torchdrivesim.simulator import TorchDriveConfig, Simulator, HomogeneousWrapper
@@ -25,14 +26,10 @@
 
 @dataclass
 class InitializationVisualizationConfig:
-    driving_surface_mesh_path: str = os.path.join(
-        os.path.dirname(os.path.realpath(__file__)), "../resources/maps/carla/meshes/Town03_driving_surface_mesh.pkl"
-    )
-    map_name: str = "Town03"
+    map_name: str = "carla_Town03"
     res: int = 1024
     fov: float = 200
     center: Optional[Tuple[float, float]] = None
-    map_origin: Tuple[float, float] = (0, 0)
     orientation: float = np.pi / 2
     save_path: str = './simulation.gif'
     method: str = 'iai'
@@ -140,13 +137,12 @@ def fit_param(self, df: pandas.DataFrame, precision: float = 0.01) -> Any:
 def visualize_map(cfg: InitializationVisualizationConfig):
     device = 'cuda'
     res = Resolution(cfg.res, cfg.res)
-    driving_surface_mesh = BirdviewMesh.unpickle(cfg.driving_surface_mesh_path).to(device)
-    simulator_cfg = TorchDriveConfig(left_handed_coordinates=cfg.left_handed)
+    map_cfg = find_map_config(cfg.map_name)
+    driving_surface_mesh = map_cfg.road_mesh.to(device)
+    simulator_cfg = TorchDriveConfig(left_handed_coordinates=map_cfg.left_handed_coordinates,
+                                     renderer=RendererConfig(left_handed_coordinates=map_cfg.left_handed_coordinates))
 
-    if cfg.map_name.startswith('Town'):
-        location = f'carla:{":".join(cfg.map_name.split("_"))}'
-    else:
-        location = f'canada:vancouver:{cfg.map_name}'
+    location = map_cfg.iai_location_name
     agent_attributes, agent_states, recurrent_states =\
         iai_initialize(location=location, agent_count=cfg.agent_count, center=tuple(cfg.center) if cfg.center is not None else None)
     agent_states = torch.cat([agent_states, 10*torch.ones_like(agent_states[..., :1])], dim=-1)

diff --git a/examples/gym_env.py b/examples/gym_env.py
@@ -1,6 +1,7 @@
 """
 An example showing how to define an OpenAI gym environment based on TorchDriveSim.
 It uses the IAI API to provide behaviors for other vehicles and requires an access key to run.
+See https://github.com/inverted-ai/torchdriveenv for a production-quality RL environment.
 """
 import contextlib
 import os
@@ -18,6 +19,7 @@
 
 from torchdrivesim.behavior.iai import iai_initialize, IAIWrapper
 from torchdrivesim.kinematic import KinematicBicycle
+from torchdrivesim.map import find_map_config
 from torchdrivesim.mesh import BirdviewMesh
 from torchdrivesim.rendering import RendererConfig, renderer_from_config
 from torchdrivesim.utils import Resolution
@@ -31,14 +33,9 @@
 class TorchDriveGymEnvConfig:
     simulator: TorchDriveConfig = TorchDriveConfig()
     visualize_to: Optional[str] = None
-    driving_surface_mesh_path: str = os.path.join(
-        os.path.dirname(os.path.realpath(__file__)), "../resources/maps/carla/meshes/Town03_driving_surface_mesh.pkl"
-    )
-    location: str = 'Town03'
+    map_name: str = 'carla_Town03'
     res: int = 1024
     fov: float = 200
-    center: Tuple[float, float] = (0, 0)
-    left_handed: bool = True
     agent_count: int = 5
     steps: int = 20
 
@@ -148,13 +145,18 @@ class IAIGymEnv(GymEnv):
     """
     def __init__(self, cfg: TorchDriveGymEnvConfig):
         device = torch.device('cuda')
-        driving_surface_mesh = BirdviewMesh.unpickle(cfg.driving_surface_mesh_path).to(device)
-        simulator_cfg = TorchDriveConfig(left_handed_coordinates=cfg.left_handed,
-                                         renderer=RendererConfig(left_handed_coordinates=cfg.left_handed))
+        map_cfg = find_map_config(cfg.map_name)
+        if map_cfg is None:
+            raise RuntimeError(f'Map {cfg.map_name} not found')
+        driving_surface_mesh = map_cfg.road_mesh.to(device)
+        simulator_cfg = TorchDriveConfig(
+            left_handed_coordinates=map_cfg.left_handed_coordinates,
+            renderer=RendererConfig(left_handed_coordinates=map_cfg.left_handed_coordinates)
+        )
+        iai_location = map_cfg.iai_location_name
 
-        iai_location = f'carla:{":".join(cfg.location.split("_"))}'
         agent_attributes, agent_states, recurrent_states = \
-            iai_initialize(location=iai_location, agent_count=cfg.agent_count, center=tuple(cfg.center))
+            iai_initialize(location=iai_location, agent_count=cfg.agent_count, center=tuple(map_cfg.center))
         agent_attributes, agent_states = agent_attributes.unsqueeze(0), agent_states.unsqueeze(0)
         agent_attributes, agent_states = agent_attributes.to(device).to(torch.float32), agent_states.to(device).to(
             torch.float32)

diff --git a/examples/initialize_simulation.py b/examples/initialize_simulation.py
@@ -8,31 +8,24 @@
 
 import numpy as np
 import imageio
-import lanelet2
 import torch
 from omegaconf import OmegaConf
 
 from torchdrivesim.behavior.iai import iai_initialize
 from torchdrivesim.behavior.heuristic import heuristic_initialize
 from torchdrivesim.kinematic import KinematicBicycle
-from torchdrivesim.lanelet2 import load_lanelet_map, road_mesh_from_lanelet_map, lanelet_map_to_lane_mesh
-from torchdrivesim.mesh import BirdviewMesh
-from torchdrivesim.rendering import renderer_from_config
+from torchdrivesim.map import find_map_config
+from torchdrivesim.rendering import renderer_from_config, RendererConfig
 from torchdrivesim.simulator import TorchDriveConfig, Simulator
 from torchdrivesim.utils import Resolution
 
 
 @dataclass
 class InitializationVisualizationConfig:
-    maps_path: str
-    driving_surface_mesh_path: str = os.path.join(
-        os.path.dirname(os.path.realpath(__file__)), "../resources/maps/carla/meshes/Town03_driving_surface_mesh.pkl"
-    )
-    map_name: str = "Town03"
+    map_name: str = "carla_Town03"
     res: int = 1024
     fov: float = 200
     center: Optional[Tuple[float, float]] = None
-    map_origin: Tuple[float, float] = (0, 0)
     orientation: float = np.pi / 2
     save_path: str = './initialization.png'
     method: str = 'iai'
@@ -43,17 +36,17 @@ class InitializationVisualizationConfig:
 def visualize_map(cfg: InitializationVisualizationConfig):
     device = 'cuda'
     res = Resolution(cfg.res, cfg.res)
-    map_path = os.path.join(cfg.maps_path, f'{cfg.map_name}.osm')
-
-    driving_surface_mesh = BirdviewMesh.unpickle(cfg.driving_surface_mesh_path).to(device)
-    simulator_cfg = TorchDriveConfig(left_handed_coordinates=cfg.left_handed)
+    map_cfg = find_map_config(cfg.map_name)
+    driving_surface_mesh = map_cfg.road_mesh.to(device)
+    simulator_cfg = TorchDriveConfig(left_handed_coordinates=map_cfg.left_handed_coordinates,
+                                     renderer=RendererConfig(left_handed_coordinates=map_cfg.left_handed_coordinates))
 
     if cfg.method == 'iai':
-        location = f'carla:{":".join(cfg.map_name.split("_"))}'
-        agent_attributes, agent_states, _ = iai_initialize(location=location, agent_count=cfg.agent_count, center=tuple(cfg.center) if cfg.center is not None else None)
+        location = map_cfg.iai_location_name
+        agent_attributes, agent_states, _ = iai_initialize(location=location, agent_count=cfg.agent_count,
+                                                           center=tuple(cfg.center) if cfg.center is not None else None)
     elif cfg.method == 'heuristic':
-        lanelet_map = load_lanelet_map(map_path, origin=cfg.map_origin)
-        agent_attributes, agent_states = heuristic_initialize(lanelet_map, agent_num=cfg.agent_count)
+        agent_attributes, agent_states = heuristic_initialize(map_cfg.lanelet_map, agent_num=cfg.agent_count)
     else:
         raise ValueError(f'Unrecognized initialization method: {cfg.method}')
     agent_attributes, agent_states = agent_attributes.to(device).to(torch.float32), agent_states.to(device).to(torch.float32)

diff --git a/examples/replay.py b/examples/replay.py
@@ -13,9 +13,7 @@
 import torch
 from omegaconf import OmegaConf
 
-from torchdrivesim.behavior.iai import iai_initialize, iai_drive
 from torchdrivesim.behavior.replay import interaction_replay, ReplayWrapper
-from torchdrivesim.behavior.heuristic import heuristic_initialize
 from torchdrivesim.kinematic import KinematicBicycle, TeleportingKinematicModel
 from torchdrivesim.lanelet2 import load_lanelet_map, road_mesh_from_lanelet_map, lanelet_map_to_lane_mesh
 from torchdrivesim.mesh import BirdviewMesh

diff --git a/examples/show_map.py b/examples/show_map.py
@@ -1,5 +1,6 @@
 """
 Simple demonstration for how to grab a map and visualize it.
+Downloading maps from the IAI API requires IAI_API_KEY to be set.
 """
 import os
 import sys
@@ -8,24 +9,20 @@
 
 import numpy as np
 import imageio
-import lanelet2
 from omegaconf import OmegaConf
 
-from torchdrivesim.lanelet2 import load_lanelet_map, road_mesh_from_lanelet_map, lanelet_map_to_lane_mesh
-from torchdrivesim.mesh import BirdviewMesh
-from torchdrivesim.rendering import BirdviewRenderer, renderer_from_config, RendererConfig
-from torchdrivesim.simulator import TorchDriveConfig
+from torchdrivesim.map import find_map_config, download_iai_map
+from torchdrivesim.rendering import  renderer_from_config, RendererConfig
+from torchdrivesim.traffic_controls import traffic_controls_from_map_config
 from torchdrivesim.utils import Resolution
 
 
 @dataclass
 class MapVisualizationConfig:
-    driving_surface_mesh_path: str = os.path.join(
-        os.path.dirname(os.path.realpath(__file__)), "../resources/maps/carla/meshes/Town03_driving_surface_mesh.pkl"
-    )
-    map_name: str = "Town03"
+    map_name: str = "carla_Town03"
+    iai_location_to_download: Optional[str] = None
     res: int = 1024
-    fov: float = 1000
+    fov: float = 400
     center: Optional[Tuple[float, float]] = None
     map_origin: Tuple[float, float] = (0, 0)
     orientation: float = 0
@@ -35,11 +32,19 @@ class MapVisualizationConfig:
 def visualize_map(cfg: MapVisualizationConfig):
     device = 'cuda'
     res = Resolution(cfg.res, cfg.res)
-    driving_surface_mesh = BirdviewMesh.unpickle(cfg.driving_surface_mesh_path).to(device)
-    renderer_cfg = RendererConfig(left_handed_coordinates=True)
+    if cfg.iai_location_to_download is not None:
+        download_iai_map(cfg.iai_location_to_download, save_path=f'{cfg.map_name}')
+    map_cfg = find_map_config(cfg.map_name)
+    driving_surface_mesh = map_cfg.road_mesh.to(device)
+    renderer_cfg = RendererConfig(left_handed_coordinates=map_cfg.left_handed_coordinates)
     renderer = renderer_from_config(
         renderer_cfg, device=device, static_mesh=driving_surface_mesh
     )
+
+    traffic_controls = traffic_controls_from_map_config(map_cfg)
+    controls_mesh = renderer.make_traffic_controls_mesh(traffic_controls).to(renderer.device)
+    renderer.add_static_meshes([controls_mesh])
+
     map_image = renderer.render_static_meshes(res=res, fov=cfg.fov)
     os.makedirs(os.path.dirname(cfg.save_path), exist_ok=True)
     imageio.imsave(

diff --git a/examples/simulate.py b/examples/simulate.py
@@ -14,39 +14,37 @@
 
 from torchdrivesim.behavior.iai import iai_initialize, IAIWrapper
 from torchdrivesim.kinematic import KinematicBicycle
-from torchdrivesim.lanelet2 import load_lanelet_map, road_mesh_from_lanelet_map, lanelet_map_to_lane_mesh
-from torchdrivesim.mesh import BirdviewMesh
+from torchdrivesim.map import find_map_config
 from torchdrivesim.rendering import renderer_from_config, RendererConfig
 from torchdrivesim.simulator import TorchDriveConfig, Simulator, HomogeneousWrapper
+from torchdrivesim.traffic_controls import traffic_controls_from_map_config
 from torchdrivesim.utils import Resolution
 
 
 @dataclass
 class SimulationConfig:
-    driving_surface_mesh_path: str = os.path.join(
-        os.path.dirname(os.path.realpath(__file__)), "../resources/maps/carla/meshes/Town03_driving_surface_mesh.pkl"
-    )
-    map_name: str = "Town03"
+    map_name: str = "carla_Town03"
     res: int = 1024
     fov: float = 200
     center: Optional[Tuple[float, float]] = None
-    map_origin: Tuple[float, float] = (0, 0)
     orientation: float = np.pi / 2
     save_path: str = './simulation.gif'
     method: str = 'iai'
-    left_handed: bool = True
     agent_count: int = 5
     steps: int = 20
 
 
 def simulate(cfg: SimulationConfig):
     device = 'cuda'
     res = Resolution(cfg.res, cfg.res)
-    driving_surface_mesh = BirdviewMesh.unpickle(cfg.driving_surface_mesh_path).to(device)
-    simulator_cfg = TorchDriveConfig(left_handed_coordinates=cfg.left_handed,
-                                     renderer=RendererConfig(left_handed_coordinates=cfg.left_handed))
+    map_cfg = find_map_config(cfg.map_name)
+    if cfg.center is None:
+        cfg.center = map_cfg.center
+    driving_surface_mesh = map_cfg.road_mesh.to(device)
+    simulator_cfg = TorchDriveConfig(left_handed_coordinates=map_cfg.left_handed_coordinates,
+                                     renderer=RendererConfig(left_handed_coordinates=map_cfg.left_handed_coordinates))
 
-    location = f'carla:{":".join(cfg.map_name.split("_"))}'
+    location = map_cfg.iai_location_name
     agent_attributes, agent_states, recurrent_states =\
         iai_initialize(location=location, agent_count=cfg.agent_count, center=tuple(cfg.center) if cfg.center is not None else None)
     agent_attributes, agent_states = agent_attributes.unsqueeze(0), agent_states.unsqueeze(0)
@@ -55,12 +53,13 @@ def simulate(cfg: SimulationConfig):
     kinematic_model.set_params(lr=agent_attributes[..., 2])
     kinematic_model.set_state(agent_states)
     renderer = renderer_from_config(simulator_cfg.renderer, static_mesh=driving_surface_mesh)
+    traffic_controls = traffic_controls_from_map_config(map_cfg)
 
     simulator = Simulator(
         cfg=simulator_cfg, road_mesh=driving_surface_mesh,
         kinematic_model=dict(vehicle=kinematic_model), agent_size=dict(vehicle=agent_attributes[..., :2]),
         initial_present_mask=dict(vehicle=torch.ones_like(agent_states[..., 0], dtype=torch.bool)),
-        renderer=renderer,
+        renderer=renderer, traffic_controls=traffic_controls,
     )
     simulator = HomogeneousWrapper(simulator)
     npc_mask = torch.ones(agent_states.shape[-2], dtype=torch.bool, device=agent_states.device)

diff --git a/torchdrivesim/__init__.py b/torchdrivesim/__init__.py
@@ -1,5 +1,13 @@
 __version__ = "0.1.0"
 
+import os
+
+if 'TDS_RESOURCE_PATH' in os.environ:
+    _resource_path = os.environ['TDS_RESOURCE_PATH'].split(':')
+else:
+    _resource_path = []
+_resource_path += [os.path.join(x, 'resources/maps') for x in __path__]
+
 
 def assert_pytorch3d_available():
     from torchdrivesim.rendering.pytorch3d import is_available as pytorch3d_available