Driving in traffic - dependency version and type annotation

cli/studio.py

@@ -21,9 +21,9 @@
 import os
 import subprocess
 import sys
-from typing import Sequence
 from pathlib import Path
 from threading import Thread
+from typing import Sequence
 
 import click
 

examples/driving_in_traffic/Dockerfile

@@ -41,7 +41,7 @@ WORKDIR /src
 RUN pip install --no-cache-dir -e ./examples/driving_in_traffic
 
 # Build scenarios and maps.
-RUN scl scenario build-all --clean ./examples/driving_in_traffic/scenarios/loop
+RUN scl scenario build-all --clean ./examples/driving_in_traffic/scenarios
 
 # For Envision.
 EXPOSE 8081
@@ -50,6 +50,6 @@ EXPOSE 8081
 ENV TF_XLA_FLAGS="--tf_xla_enable_xla_devices --tf_xla_auto_jit=2"
 
 # Suppress message of missing /dev/input folder.
-RUN echo "mkdir -p /dev/input\n" >> ~/.bashrc
+RUN echo "mkdir -p /dev/input" >> ~/.bashrc
 
 SHELL ["/bin/bash", "-c", "-l"]

examples/driving_in_traffic/README.md

@@ -1,5 +1,5 @@
 # Driving in Traffic
-This example illustrates the training of an ego agent to drive, as fast and as far as possible, in traffic using the DreamerV2 (https://github.com/danijar/dreamerv2) reinforcement-learning algorithm.
+This example illustrates the training of an ego agent to drive, as fast and as far as possible, in traffic using the [DreamerV2](https://github.com/danijar/dreamerv2) reinforcement-learning algorithm.
 
 The ego agent earns reward for the distance travelled per-step and is penalised for colliding with other vehicles and for going off-road.
 

examples/driving_in_traffic/config.yaml

@@ -11,6 +11,7 @@ smarts:
     - loop
 
 dreamerv2:
+  steps: 1e8
   log_every: 1e4
   eval_every: 1e5
   prefill: 1e4

examples/driving_in_traffic/driving_in_traffic/env/action.py

@@ -1,9 +1,11 @@
+from typing import Callable, Tuple
+
 import gym
 import numpy as np
 
 
 class Action(gym.ActionWrapper):
-    def __init__(self, env):
+    def __init__(self, env: gym.Env):
         super().__init__(env)
         self._wrapper, self.action_space = continuous()
 
@@ -14,7 +16,7 @@ def action(self, act):
         return wrapped_act
 
 
-def continuous():
+def continuous() -> Tuple[Callable[[np.array], np.array], gym.Space]:
     space = gym.spaces.Box(low=-1.0, high=1.0, shape=(3,), dtype=np.float32)
 
     def wrapper(model_action):

examples/driving_in_traffic/driving_in_traffic/env/reward.py

@@ -1,9 +1,11 @@
 import gym
 import numpy as np
 
+from smarts.core.sensors import Observation
+
 
 class Reward(gym.Wrapper):
-    def __init__(self, env):
+    def __init__(self, env: gym.Env):
         super().__init__(env)
 
     def reset(self, **kwargs):
@@ -19,7 +21,7 @@ def step(self, action):
 
         return obs, wrapped_reward, done, info
 
-    def _reward(self, obs, env_reward):
+    def _reward(self, obs: Observation, env_reward: np.float64) -> np.float32:
         reward = 0
 
         # Penalty for driving off road

examples/driving_in_traffic/driving_in_traffic/env/single_agent.py

@@ -1,6 +1,7 @@
 from functools import partial
-from typing import Dict
+from typing import Any, Callable, Dict, Generator
 
+import gym
 from driving_in_traffic.env import action, reward
 
 from smarts.core import agent as smarts_agent
@@ -11,14 +12,16 @@
 from smarts.env.wrappers import single_agent as smarts_single_agent
 
 
-def gen_env(config: Dict, seed: int):
+def gen_env(
+    config: Dict[str, Any], seed: int
+) -> Generator[Callable[[str], gym.Env], None, None]:
     base_seed = seed
     while True:
         yield partial(make_env, config=config, seed=base_seed)
         base_seed += 1
 
 
-def make_env(config: Dict, seed: int, env_name: str = None):
+def make_env(config: Dict[str, Any], seed: int, env_name: str = None) -> gym.Env:
 
     vehicle_interface = smarts_agent_interface.AgentInterface(
         max_episode_steps=config["max_episode_steps"],

examples/driving_in_traffic/run.py

@@ -11,7 +11,9 @@
 import re
 import warnings
 from datetime import datetime
+from typing import Callable, Generator
 
+import gym
 import numpy as np
 import rich.traceback
 import tensorflow as tf
@@ -33,7 +35,7 @@
 seed(42)
 
 
-def main(args):
+def main(args: argparse.Namespace):
     # Load config file.
     config_file = yaml.load(
         (pathlib.Path(__file__).absolute().parent / "config.yaml").read_text()
@@ -42,7 +44,6 @@ def main(args):
     # Load SMARTS env config.
     config_env = config_file["smarts"]
     config_env["mode"] = args.mode
-    config_env["logdir"] = args.logdir
     config_env["headless"] = not args.head
     config_env["scenarios_dir"] = (
         pathlib.Path(__file__).absolute().parents[0] / "scenarios"
@@ -57,46 +58,24 @@ def main(args):
     tf.config.run_functions_eagerly(not config_dv2.jit)
 
     # Setup GPU.
-    gpus = tf.config.list_physical_devices("GPU")
-    if gpus:
-        try:
-            for gpu in gpus:
-                tf.config.experimental.set_memory_growth(gpu, True)
-            logical_gpus = tf.config.list_logical_devices("GPU")
-            print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
-        except RuntimeError as e:
-            print(e)
-    else:
-        warnings.warn(
-            f"Not configured to use GPU or GPU not available.",
-            ResourceWarning,
-        )
+    _setup_gpu()
 
     # Train or evaluate.
-    if config_env["mode"] == "train":
-        if not config_env["logdir"]:
-            # Begin training from scratch.
-            time = datetime.now().strftime("%Y_%m_%d_%H_%M")
-            logdir = pathlib.Path(__file__).absolute().parents[0] / "logs" / time
-        else:
-            # Begin training from a pretrained model.
-            logdir = config_env["logdir"]
-        config_dv2 = config_dv2.update(
-            {
-                "logdir": logdir,
-            }
-        )
+    if config_env["mode"] == "train" and not args.logdir:
+        # Train from scratch.
+        time = datetime.now().strftime("%Y_%m_%d_%H_%M")
+        logdir = pathlib.Path(__file__).absolute().parents[0] / "logs" / time
+    elif config_env["mode"] == "train" and args.logdir:
+        # Begin training from a pretrained model.
+        logdir = args.logdir
     elif config_env["mode"] == "evaluate":
-        config_dv2 = config_dv2.update(
-            {
-                "logdir": config_env["logdir"],
-                "eval_eps": 1e8,
-            }
-        )
+        logdir = args.logdir
+        config_dv2 = config_dv2.update({"eval_eps": 1e8})
     else:
         raise KeyError(
             f'Expected \'train\' or \'evaluate\', but got {config_env["mode"]}.'
         )
+    config_dv2 = config_dv2.update({"logdir": logdir})
 
     # Create SMARTS env.
     gen_env = single_agent.gen_env(config_env, config_env["seed"])
@@ -111,7 +90,24 @@ def _build_scenario():
     os.system(build_scenario)
 
 
-def _wrap_env(env, config):
+def _setup_gpu():
+    gpus = tf.config.list_physical_devices("GPU")
+    if gpus:
+        try:
+            for gpu in gpus:
+                tf.config.experimental.set_memory_growth(gpu, True)
+            logical_gpus = tf.config.list_logical_devices("GPU")
+            print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
+        except RuntimeError as e:
+            print(e)
+    else:
+        warnings.warn(
+            f"Not configured to use GPU or GPU not available.",
+            ResourceWarning,
+        )
+
+
+def _wrap_env(env: gym.Env, config: common.config.Config):
     env = dv2.common.GymWrapper(env)
     env = dv2.common.ResizeImage(env)
     if hasattr(env.act_space["action"], "n"):
@@ -122,12 +118,16 @@ def _wrap_env(env, config):
     return env
 
 
-def run(config, gen_env, mode: str):
+def run(
+    config: common.config.Config,
+    gen_env: Generator[Callable[[str], gym.Env], None, None],
+    mode: str,
+):
     logdir = pathlib.Path(config.logdir).expanduser()
     logdir.mkdir(parents=True, exist_ok=True)
     config.save(logdir / "config.yaml")
     print(config, "\n")
-    print("Logdir", logdir)
+    print("Logdir:", logdir)
 
     train_replay = dv2.common.Replay(logdir / "train_episodes", **config.replay)
     eval_replay = dv2.common.Replay(

examples/driving_in_traffic/setup.py

@@ -18,7 +18,7 @@
     python_requires=">=3.7",
     install_requires=[
         "setuptools>=41.0.0,!=50.0",
-        "smarts[camera-obs] @ git+https://[email protected]/huawei-noah/SMARTS@develop",
+        "smarts[camera-obs]==0.5.0",
         "dreamerv2==2.1.1",
         "tensorflow==2.4.0",
         "tensorflow-probability==0.12.2",

scenarios/od_merge/scenario.py

@@ -1,7 +1,7 @@
 from pathlib import Path
 
-from smarts.sstudio.genscenario import gen_scenario
 from smarts.sstudio import types as t
+from smarts.sstudio.genscenario import gen_scenario
 
 ego_missions = [t.Mission(t.Route(begin=("1_0_R", 1, 5), end=("1_2_R", 1, "max")))]
 

smarts/core/default_map_builder.py

@@ -20,12 +20,12 @@
 
 import os
 from dataclasses import replace
-from typing import NamedTuple, Optional, Tuple
 from pathlib import Path
+from typing import NamedTuple, Optional, Tuple
+
 from smarts.core.road_map import RoadMap
 from smarts.core.utils.file import file_md5_hash, path2hash
 
-
 _existing_map = None
 
 

smarts/core/lanepoints.py

@@ -28,6 +28,7 @@
 from dataclasses import dataclass
 from functools import lru_cache
 from typing import List, NamedTuple, Sequence, Tuple
+
 import numpy as np
 
 with warnings.catch_warnings():
@@ -39,7 +40,7 @@
         # aggressive
         from sklearn.neighbors import KDTree
 
-from smarts.core.coordinates import Heading, Pose, Point
+from smarts.core.coordinates import Heading, Point, Pose
 from smarts.core.road_map import RoadMap
 from smarts.core.utils.math import (
     fast_quaternion_from_angle,

smarts/core/opendrive_road_network.py

@@ -18,19 +18,22 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 # THE SOFTWARE.
 import heapq
-import os
 import logging
 import math
+import os
 import random
 import time
+from bisect import bisect
 from dataclasses import dataclass
 from functools import lru_cache
-from typing import Dict, List, Sequence, Set, Tuple, Optional
+from typing import Dict, List, Optional, Sequence, Set, Tuple
+
+import numpy as np
+import rtree
 import trimesh
 import trimesh.scene
-from trimesh.exchange import gltf
-import numpy as np
 from cached_property import cached_property
+from lxml import etree
 from opendrive2lanelet.opendriveparser.elements.geometry import Line as LineGeometry
 from opendrive2lanelet.opendriveparser.elements.opendrive import (
     OpenDrive as OpenDriveElement,
@@ -50,29 +53,27 @@
     PlanView as PlanViewElement,
 )
 from opendrive2lanelet.opendriveparser.parser import parse_opendrive
-from lxml import etree
 from shapely.geometry import Polygon
-import rtree
-from bisect import bisect
+from trimesh.exchange import gltf
 
 from smarts.core.road_map import RoadMap, Waypoint
-from smarts.sstudio.types import MapSpec
+from smarts.core.utils.geometry import generate_mesh_from_polygons
+from smarts.core.utils.key_wrapper import KeyWrapper
 from smarts.core.utils.math import (
     CubicPolynomial,
     constrain_angle,
     distance_point_to_polygon,
     get_linear_segments_for_range,
+    inplace_unwrap,
     offset_along_shape,
     position_at_shape_offset,
-    vec_2d,
-    inplace_unwrap,
     radians_to_vec,
+    vec_2d,
 )
+from smarts.sstudio.types import MapSpec
 
-from .lanepoints import LinkedLanePoint, LanePoints
-from .coordinates import BoundingBox, Point, Pose, RefLinePoint, Heading
-from smarts.core.utils.geometry import generate_mesh_from_polygons
-from smarts.core.utils.key_wrapper import KeyWrapper
+from .coordinates import BoundingBox, Heading, Point, Pose, RefLinePoint
+from .lanepoints import LanePoints, LinkedLanePoint
 
 
 def _convert_camera(camera):

smarts/core/sumo_road_network.py

@@ -21,24 +21,23 @@
 import os
 import random
 from functools import lru_cache
+from subprocess import check_output
 from typing import List, Optional, Sequence, Set, Tuple
 
 import numpy as np
 import trimesh
 import trimesh.scene
 from cached_property import cached_property
+from shapely.geometry import Point as shPoint
 from shapely.geometry import Polygon
-from shapely.ops import snap, nearest_points
+from shapely.ops import nearest_points, snap
 from trimesh.exchange import gltf
-from shapely.geometry import Point as shPoint
-from subprocess import check_output
-
 
 from smarts.sstudio.types import MapSpec
 
 from .coordinates import BoundingBox, Heading, Point, Pose, RefLinePoint
+from .lanepoints import LanePoints, LinkedLanePoint
 from .road_map import RoadMap, Waypoint
-from .lanepoints import LinkedLanePoint, LanePoints
 from .utils.geometry import buffered_shape, generate_mesh_from_polygons
 from .utils.math import inplace_unwrap, radians_to_vec, vec_2d
 

smarts/core/sumo_traffic_simulation.py

@@ -38,6 +38,7 @@
 from smarts.core.utils import networking
 from smarts.core.utils.logging import suppress_output
 from smarts.core.vehicle import VEHICLE_CONFIGS, VehicleState
+
 from smarts.core.utils.sumo import SUMO_PATH, traci  # isort:skip
 from traci.exceptions import FatalTraCIError, TraCIException  # isort:skip
 import traci.constants as tc  # isort:skip