huawei-noah · saulfield · Mar 17, 2023 · Mar 10, 2023 · Mar 13, 2023 · Mar 13, 2023
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -10,10 +10,12 @@ Copy and pasting the git commit messages is __NOT__ enough.
 
 ## [Unreleased]
 ### Added
+- Added `waymo_open_dataset` as a module at the SMARTS repo level, to be able to load waymo scenarios without any external packages
 ### Changed
 ### Deprecated
 ### Fixed
 ### Removed
+- Removed `waymo-open-dataset-tf-2-4-0` package as a dependency
 ### Security
 
 ## [1.0.8] # 2023-03-10

diff --git a/cli/cli.py b/cli/cli.py
@@ -27,6 +27,7 @@
 from cli.envision import envision_cli
 from cli.run import run_experiment
 from cli.studio import scenario_cli
+from cli.waymo import waymo_cli
 from cli.zoo import zoo_cli
 
 
@@ -44,24 +45,6 @@ def scl():
 scl.add_command(scenario_cli)
 scl.add_command(zoo_cli)
 scl.add_command(run_experiment)
-
-try:
-    from cli.waymo import waymo_cli
-except (ModuleNotFoundError, ImportError):
-
-    @click.group(
-        name="waymo",
-        invoke_without_command=True,
-        help="The `scl waymo` command requires `[waymo]`.",
-    )
-    @click.pass_context
-    def waymo_cli(ctx):
-        click.echo(
-            "The `scl waymo` command is unavailable. To enable, pip install the missing dependencies.\n"
-            "pip install pathos==0.2.8 tabulate>=0.8.10 waymo-open-dataset-tf-2-4-0"
-        )
-
-
 scl.add_command(waymo_cli)
 scl.add_command(diagnostic_cli)
 

diff --git a/scenarios/waymo/README.md b/scenarios/waymo/README.md
@@ -9,10 +9,7 @@ The SMARTS repository contains tools to assist with using the Waymo Motion Datas
 - `gen_sumo_map.py`: a command-line program that converts the map from a Waymo scenario to a SUMO map
 
 ## Setup
-Follow the SMARTS setup instructions in the main [README](https://github.com/huawei-noah/SMARTS/). Then install the `[waymo]` dependencies.
-```bash
-$ pip install -e .[waymo]
-```
+Follow the SMARTS setup instructions in the main [README](https://github.com/huawei-noah/SMARTS/).
 
 Download the dataset files from the [Waymo Motion Dataset](https://waymo.com/open/download/). It is recommended to download the dataset files from the `uncompressed/scenario/training_20s` folder as they have the full traffic capture for each scenario. Note: Waymo provides 2 different formats for the dataset files. SMARTS expects the `Scenario protos` format (not the `tf.Example protos` format). It is also recommended to use version 1.1 of the dataset, which includes enhanced map information.
 

diff --git a/setup.cfg b/setup.cfg
@@ -115,7 +115,6 @@ train =
 waymo = 
     pathos==0.2.8
     tabulate>=0.8.10
-    waymo-open-dataset-tf-2-4-0
 opendrive = 
     opendrive2lanelet>=1.2.1
 

diff --git a/smarts/core/default_map_builder.py b/smarts/core/default_map_builder.py
@@ -127,14 +127,7 @@ def get_road_map(map_spec) -> Tuple[Optional[RoadMap], Optional[str]]:
         map_class = OpenDriveRoadNetwork
 
     elif map_type == _WAYMO_MAP:
-        try:
-            from smarts.core.waymo_map import WaymoMap  # pytype: disable=import-error
-        except (ImportError, ModuleNotFoundError):
-            print(sys.exc_info())
-            print(
-                "You may not have installed the [waymo] dependencies required to build and use WaymoMap Scenario. Install them first using the command `pip install -e .[waymo]` at the source directory."
-            )
-            return None, None
+        from smarts.core.waymo_map import WaymoMap
         map_class = WaymoMap
     else:
         return None, None

diff --git a/smarts/core/waymo_map.py b/smarts/core/waymo_map.py
@@ -924,7 +924,7 @@ def parse_source_to_scenario(source: str):
                     f"Dataset file does not contain scenario with id: {scenario_id}"
                 )
             parsed_scenario = scenario_pb2.Scenario()
-            parsed_scenario.ParseFromString(bytearray(record))
+            parsed_scenario.ParseFromString(bytes(record))
             WaymoMap._scenario_cache[parsed_scenario.scenario_id] = parsed_scenario
             if parsed_scenario.scenario_id == scenario_id:
                 return parsed_scenario

diff --git a/smarts/sstudio/genhistories.py b/smarts/sstudio/genhistories.py
@@ -42,17 +42,8 @@
 )
 from smarts.sstudio import types
 from smarts.waymo.waymo_utils import WaymoDatasetError
-
-try:
-    # pytype: disable=import-error
-    from waymo_open_dataset.protos import scenario_pb2
-    from waymo_open_dataset.protos.map_pb2 import TrafficSignalLaneState
-
-    # pytype: enable=import-error
-except ImportError:
-    print(
-        "You may not have installed the [waymo] dependencies required to use the waymo replay simulation. Install them first using the command `pip install -e .[waymo]` at the source directory."
-    )
+from waymo_open_dataset.protos import scenario_pb2
+from waymo_open_dataset.protos.map_pb2 import TrafficSignalLaneState
 
 METERS_PER_FOOT = 0.3048
 DEFAULT_LANE_WIDTH = 3.7  # a typical US highway lane is 12ft ~= 3.7m wide
@@ -810,7 +801,7 @@ def _get_scenario(self):
         dataset = read_tfrecord_file(self._dataset_spec["input_path"])
         for record in dataset:
             parsed_scenario = scenario_pb2.Scenario()
-            parsed_scenario.ParseFromString(bytearray(record))
+            parsed_scenario.ParseFromString(bytes(record))
             if parsed_scenario.scenario_id == scenario_id:
                 return parsed_scenario
         raise ValueError(

diff --git a/smarts/sstudio/scenario_construction.py b/smarts/sstudio/scenario_construction.py
@@ -57,10 +57,13 @@ def build_scenario(
                     f"from smarts.core import seed as smarts_seed; smarts_seed({seed});\n"
                 )
                 c.write(o.read())
-
             c.flush()
+
+            # Set the python path to the base dir so the script can find modules properly
+            env = os.environ.copy()
+            env["PYTHONPATH"] = Path(__file__).parents[2]
             subprocess.check_call(
-                [sys.executable, Path(c.name).name], cwd=scenario_root
+                [sys.executable, Path(c.name).name], cwd=scenario_root, env=env
             )
 
 

diff --git a/smarts/waymo/gen_sumo_map.py b/smarts/waymo/gen_sumo_map.py
@@ -28,7 +28,7 @@
 from pathlib import Path
 from typing import Callable, Dict, List, Tuple
 
-from waymo_open_dataset.protos import map_pb2
+from smarts.waymo import map_pb2
 
 from smarts.core.waymo_map import WaymoMap
 
@@ -98,7 +98,7 @@ def _create_edge(
         edge_id: str,
         start_id: str,
         end_id: str,
-        shape_str: Tuple[List[float], List[float]],
+        shape_str: str,
         width: float = 5,
     ):
         edge = ET.SubElement(self.edges_root, "edge")

diff --git a/smarts/waymo/waymo_browser.py b/smarts/waymo/waymo_browser.py
@@ -1026,7 +1026,7 @@ def get_scenario_and_tag_dict(
     dataset = read_tfrecord_file(tfrecord_file)
     for record in dataset:
         scenario = scenario_pb2.Scenario()
-        scenario.ParseFromString(bytearray(record))
+        scenario.ParseFromString(bytes(record))
         scenario_id = scenario.scenario_id
         scenario_dict[scenario_id] = [scenario, None, None]
         tags_per_scenario[scenario_id] = []

diff --git a/smarts/waymo/waymo_utils.py b/smarts/waymo/waymo_utils.py
@@ -272,7 +272,7 @@ def get_tfrecord_info(tfrecord_file: str) -> Dict[str, Dict[str, Any]]:
     records = read_tfrecord_file(tfrecord_file)
     for record in records:
         scenario = scenario_pb2.Scenario()
-        scenario.ParseFromString(bytearray(record))
+        scenario.ParseFromString(bytes(record))
         scenario_id = scenario.scenario_id
         num_vehicles = 0
         num_pedestrians = 0

diff --git a/waymo_open_dataset/README.md b/waymo_open_dataset/README.md
@@ -0,0 +1,22 @@
+# Waymo Open Dataset
+
+This module is intended to eliminate dependence on the [`waymo_open_dataset`](https://github.com/waymo-research/waymo-open-dataset) package in SMARTS. We only include the `scenario` and `map` protobuf definitions, as well as their compiled Python files. The current Python files were generated using version `3.19.6` of the protobuf compiler. See the below instructions if you encounter any errors related to protobuf versions that require regenerating the Python files.
+
+# Compilation Instructions
+
+To regenerate the `scenario_pb2.py` and `map_pb2.py` files:
+
+1. Install the protobuf compiler: see https://grpc.io/docs/protoc-installation/
+2. (Optional) Download the proto files if necessary
+
+```sh
+curl https://raw.githubusercontent.com/waymo-research/waymo-open-dataset/master/waymo_open_dataset/protos/scenario.proto -o waymo_open_dataset/protos/scenario.proto
+
+curl https://raw.githubusercontent.com/waymo-research/waymo-open-dataset/master/waymo_open_dataset/protos/map.proto -o waymo_open_dataset/protos/map.proto
+```
+
+3. Compile the proto files
+
+```sh
+protoc -I=. --python_out=. waymo_open_dataset/protos/scenario.proto && protoc -I=. --python_out=. waymo_open_dataset/protos/map.proto
+```
diff --git a/waymo_open_dataset/__init__.py b/waymo_open_dataset/__init__.py