Merged commit includes the following changes:

592002277  by Waymo Research:

    Increment PIP package version number 1.6.0 -> 1.6.1

--
592002084  by Waymo Research:

    Internal change.

--
583186773  by Waymo Research:

    Internal Code Change

--
579275350  by Waymo Research:

    Internal Code Change

--
576198969  by Waymo Research:

    Modifying the behaviour of collision/offroad checking when invalid.

--
568168016  by Waymo Research:

    Fixing a bug in validity masking for collision and offroad.

--

PiperOrigin-RevId: 592002277

Merged commit includes the following changes:
592617521  by Waymo Research:

    Internal change.

--

PiperOrigin-RevId: 592617521

README.md

@@ -26,6 +26,12 @@ compliance with the Waymo Dataset License Agreement for Non-Commercial Use. See
 [wdl_limited/sim_agents_metrics/](src/waymo_open_dataset/wdl_limited/sim_agents_metrics),
 respectively, for details.
 
+## December 2023 Update
+We released v1.6.1 version of the pip package with fixes for the WOSAC metrics:
+- Fixing a bug in validity checking for collision and offroad.
+- Modifying the behaviour of collision/offroad checking when invalid.
+
+
 ## August 2023 Update
 We released a large-scale object-centric asset dataset containing over 1.2M images and lidar observations of two major categories (vehicles and pedestrians) from the Perception Dataset (v2.0.0).
 

src/waymo_open_dataset/pip_pkg_scripts/BUILD

@@ -199,7 +199,7 @@ py_wheel(
         "visu3d==1.5.1",
         "dacite==1.8.1",
     ],
-    version = "1.6.0",
+    version = "1.6.1",
     deps = [
         ":all_python_modules",
     ],
@@ -215,7 +215,7 @@ py_wheel(
         "numpy==1.21.5",
         "tensorflow==2.11",
     ],
-    version = "1.6.0",
+    version = "1.6.1",
     deps = [
         "@deeplab2",
     ],

src/waymo_open_dataset/utils/sim_agents/BUILD

@@ -90,6 +90,7 @@ py_test(
     srcs = ["converters_test.py"],
     python_version = "PY3",
     deps = [
+        requirement("absl-py"),
         requirement("tensorflow"),
         ":converters",
         ":submission_specs",

src/waymo_open_dataset/utils/sim_agents/converters.py

@@ -75,6 +75,7 @@ def scenario_to_joint_scene(
 def joint_scene_to_trajectories(
     joint_scene: sim_agents_submission_pb2.JointScene,
     scenario: scenario_pb2.Scenario,
+    use_log_validity: bool = False
 ) -> trajectory_utils.ObjectTrajectories:
   """Converts a JointScene and the relative Scenario into `ObjectTrajectories`.
 
@@ -86,16 +87,19 @@ def joint_scene_to_trajectories(
       (`CURRENT_TIME_INDEX` when 0-indexed). The history steps from this
       scenario are also prepended to the returned `Trajectories`, resulting in a
       total length of `submission_specs.N_FULL_SCENARIO_STEPS`.
+    use_log_validity: If True, copies the validity mask from the original
+      scenario instead of assuming all the steps are valid. This is used to
+      compute features for the logged scenario.
 
   Returns:
     An `ObjectTrajectories` containing the trajectories of all simulated objects
     in a scenario, with prepended history steps and inferred static dimensions
     and object types.
   """
-  logged_trajectories = trajectory_utils.ObjectTrajectories.from_scenario(
+  logged_trajectories_full = trajectory_utils.ObjectTrajectories.from_scenario(
       scenario
   )
-  logged_trajectories = logged_trajectories.slice_time(
+  logged_trajectories_history = logged_trajectories_full.slice_time(
       start_index=0, end_index=submission_specs.CURRENT_TIME_INDEX + 1
   )
   # Extract states and object IDs from the simulated scene.
@@ -110,28 +114,52 @@ def joint_scene_to_trajectories(
   sim_x, sim_y, sim_z, sim_heading, sim_ids = map(
       tf.convert_to_tensor, [sim_x, sim_y, sim_z, sim_heading, sim_ids])
   # Align objects from the logged scenario to the simulated one.
-  logged_trajectories = logged_trajectories.gather_objects_by_id(sim_ids)
+  logged_trajectories_history = (
+      logged_trajectories_history.gather_objects_by_id(sim_ids))
   # Prepare the missing tensors: validity and box sizes.
-  sim_valid = tf.fill(sim_x.shape, True)
+  if use_log_validity:
+    # When copying validity from the log, make sure objects are in the same
+    # order as `sim_ids`.
+    logged_trajectories_full = logged_trajectories_full.gather_objects_by_id(
+        sim_ids)
+    logged_trajectories_future = logged_trajectories_full.slice_time(
+        start_index=submission_specs.CURRENT_TIME_INDEX + 1
+    )
+    sim_valid = logged_trajectories_future.valid
+  else:
+    sim_valid = tf.fill(sim_x.shape, True)
+
   sim_length = tf.repeat(
-      logged_trajectories.length[:, -1, tf.newaxis], sim_x.shape[-1], axis=-1
+      logged_trajectories_history.length[:, -1, tf.newaxis],
+      sim_x.shape[-1],
+      axis=-1,
   )
   sim_width = tf.repeat(
-      logged_trajectories.width[:, -1, tf.newaxis], sim_x.shape[-1], axis=-1
+      logged_trajectories_history.width[:, -1, tf.newaxis],
+      sim_x.shape[-1],
+      axis=-1,
   )
   sim_height = tf.repeat(
-      logged_trajectories.height[:, -1, tf.newaxis], sim_x.shape[-1], axis=-1
+      logged_trajectories_history.height[:, -1, tf.newaxis],
+      sim_x.shape[-1],
+      axis=-1,
   )
-  # Concatenate and return.
+  # Concatenate history and logged/simulated future.
   return trajectory_utils.ObjectTrajectories(
-      x=tf.concat([logged_trajectories.x, sim_x], axis=-1),
-      y=tf.concat([logged_trajectories.y, sim_y], axis=-1),
-      z=tf.concat([logged_trajectories.z, sim_z], axis=-1),
-      heading=tf.concat([logged_trajectories.heading, sim_heading], axis=-1),
-      length=tf.concat([logged_trajectories.length, sim_length], axis=-1),
-      width=tf.concat([logged_trajectories.width, sim_width], axis=-1),
-      height=tf.concat([logged_trajectories.height, sim_height], axis=-1),
-      valid=tf.concat([logged_trajectories.valid, sim_valid], axis=-1),
+      x=tf.concat([logged_trajectories_history.x, sim_x], axis=-1),
+      y=tf.concat([logged_trajectories_history.y, sim_y], axis=-1),
+      z=tf.concat([logged_trajectories_history.z, sim_z], axis=-1),
+      heading=tf.concat(
+          [logged_trajectories_history.heading, sim_heading], axis=-1
+      ),
+      length=tf.concat(
+          [logged_trajectories_history.length, sim_length], axis=-1
+      ),
+      width=tf.concat([logged_trajectories_history.width, sim_width], axis=-1),
+      height=tf.concat(
+          [logged_trajectories_history.height, sim_height], axis=-1
+      ),
+      valid=tf.concat([logged_trajectories_history.valid, sim_valid], axis=-1),
       object_id=sim_ids,
-      object_type=logged_trajectories.object_type,
+      object_type=logged_trajectories_history.object_type,
   )

src/waymo_open_dataset/utils/sim_agents/converters_test.py

@@ -14,14 +14,16 @@
 # =============================================================================
 """Tests for waymo_open_dataset.utils.sim_agents.converters."""
 
+from absl.testing import parameterized
+
 import tensorflow as tf
 
 from waymo_open_dataset.utils import test_utils
 from waymo_open_dataset.utils.sim_agents import converters
 from waymo_open_dataset.utils.sim_agents import submission_specs
 
 
-class ConvertersTest(tf.test.TestCase):
+class ConvertersTest(parameterized.TestCase, tf.test.TestCase):
 
   def test_scenario_to_joint_scene(self):
     scenario = test_utils.get_womd_test_scenario()
@@ -34,11 +36,14 @@ def test_scenario_to_joint_scene(self):
     # Validate the joint scene.
     submission_specs.validate_joint_scene(joint_scene, scenario)
 
-  def test_joint_scene_to_trajectories(self):
+  @parameterized.named_parameters(
+      {'testcase_name': 'all_valid', 'use_log_validity': False},
+      {'testcase_name': 'log_valid', 'use_log_validity': True})
+  def test_joint_scene_to_trajectories(self, use_log_validity: bool):
     scenario = test_utils.get_womd_test_scenario()
     joints_scene = converters.scenario_to_joint_scene(scenario)
     trajectories = converters.joint_scene_to_trajectories(
-        joints_scene, scenario)
+        joints_scene, scenario, use_log_validity=use_log_validity)
     # Check shapes of the time-series fields. The test scenario contains a
     # total of 50 sim agents.
     time_fields = [
@@ -49,7 +54,14 @@ def test_joint_scene_to_trajectories(self):
     # Check shapes of the per-object fields.
     self.assertEqual(trajectories.object_id.shape, (50,))
     self.assertEqual(trajectories.object_type.shape, (50,))
-
+    # If `use_log_validity=True`, we want to check that the invalid objects
+    # in the test scenario are propagated as invalid. Otherwise, if
+    # `use_log_validity=False` we want to verify that all the objects are
+    # assigned a valid state.
+    self.assertEqual(
+        tf.reduce_all(
+            trajectories.valid[:, submission_specs.CURRENT_TIME_INDEX:]),
+        not use_log_validity)
 
 if __name__ == '__main__':
   tf.random.set_seed(42)

src/waymo_open_dataset/wdl_limited/sim_agents_metrics/BUILD

@@ -139,6 +139,8 @@ py_test(
     srcs = ["metric_features_test.py"],
     deps = [
         requirement("tensorflow"),
+        ":interaction_features",
+        ":map_metric_features",
         ":metric_features",
         "//waymo_open_dataset/utils:test_utils",
         "//waymo_open_dataset/utils/sim_agents:converters",

src/waymo_open_dataset/wdl_limited/sim_agents_metrics/interaction_features.py

@@ -26,7 +26,7 @@
 # Constant distance to apply when distances between objects are invalid. This
 # will avoid the propagation of nans and should be reduced out when taking the
 # minimum anyway.
-_EXTREMELY_LARGE_DISTANCE = 1e10
+EXTREMELY_LARGE_DISTANCE = 1e10
 # Collision threshold, i.e. largest distance between objects that is considered
 # to be a collision.
 COLLISION_DISTANCE_THRESHOLD = 0.0
@@ -147,7 +147,7 @@ def compute_distance_to_nearest_object(
   # Mask out self-distances.
   self_mask = tf.eye(
       num_eval_objects, num_objects, dtype=tf.float32)[:, :, tf.newaxis]
-  signed_distances = signed_distances + self_mask * _EXTREMELY_LARGE_DISTANCE
+  signed_distances = signed_distances + self_mask * EXTREMELY_LARGE_DISTANCE
 
   # Mask out invalid boxes. As with box coordinates, the validity mask needs to
   # be reshuffled to have the same ordering. This is necessary because the
@@ -164,7 +164,7 @@ def compute_distance_to_nearest_object(
   valid_mask = tf.logical_and(eval_validity[:, tf.newaxis, :],
                               all_validity[tf.newaxis, :, :])
   signed_distances = tf.where(
-      valid_mask, signed_distances, _EXTREMELY_LARGE_DISTANCE)
+      valid_mask, signed_distances, EXTREMELY_LARGE_DISTANCE)
   # Aggregate over the "all objects" dimension.
   return tf.reduce_min(signed_distances, axis=1)
 
@@ -296,7 +296,7 @@ def compute_time_to_collision_with_object_in_front(
   # `masked_long_distance` shape: (num_steps, num_eval_objects, num_objects)
   masked_long_distance = (
       long_distance
-      + (1.0 - tf.cast(valid_mask, tf.float32)) * _EXTREMELY_LARGE_DISTANCE
+      + (1.0 - tf.cast(valid_mask, tf.float32)) * EXTREMELY_LARGE_DISTANCE
   )
 
   # `box_ahead_index` shape: (num_steps, num_evaluated_objects)

src/waymo_open_dataset/wdl_limited/sim_agents_metrics/interaction_features_test.py

@@ -123,7 +123,7 @@ def test_distance_to_nearest_object_handles_invalid_objects(self, box2_valid):
       self.assertAllClose(distances, tf.fill(distances.shape, 10.0))
     else:
       self.assertAllClose(distances, tf.fill(
-          distances.shape, interaction_features._EXTREMELY_LARGE_DISTANCE))
+          distances.shape, interaction_features.EXTREMELY_LARGE_DISTANCE))
 
   def test_distance_to_nearest_object_selects_nearest(self):
     # Create 3 boxes, box1 centered on (0, 0), box2 at 5m from box1 and box3

src/waymo_open_dataset/wdl_limited/sim_agents_metrics/map_metric_features.py

@@ -24,7 +24,7 @@
 
 # Constant distance to apply when distances are invalid. This will avoid the
 # propagation of nans and should be reduced out when taking the maximum anyway.
-_EXTREMELY_LARGE_DISTANCE = 1e10
+EXTREMELY_LARGE_DISTANCE = 1e10
 # Off-road threshold, i.e. smallest distance away from the road edge that is
 # considered to be a off-road.
 OFFROAD_DISTANCE_THRESHOLD = 0.0
@@ -131,7 +131,7 @@ def compute_distance_to_road_edge(
   # Mask out invalid boxes.
   eval_validity = tf.gather(
       valid, tf.where(evaluated_object_mask)[:, 0], axis=0)
-  return tf.where(eval_validity, signed_distances, -_EXTREMELY_LARGE_DISTANCE)
+  return tf.where(eval_validity, signed_distances, -EXTREMELY_LARGE_DISTANCE)
 
 
 def _compute_signed_distance_to_polylines(

src/waymo_open_dataset/wdl_limited/sim_agents_metrics/metric_features.py

@@ -67,16 +67,15 @@ class MetricFeatures:
       in angular_speed. Shape: (n_samples, n_objects, n_steps).
     distance_to_nearest_object: Signed distance (in meters) to the nearest
       object in the scene. Shape: (n_samples, n_objects, n_steps).
-    collision_indication: Boolean value indicating whether the object collided,
-      at any point in time, with any other object.
-      Shape: (n_samples, n_objects).
+    collision_per_step: Boolean tensor indicating whether the object collided,
+      with any other object. Shape: (n_samples, n_objects, n_steps).
     time_to_collision: Time (in seconds) before the object collides with the
       object it is following (if it exists), assuming constant speeds.
       Shape: (n_samples, n_objects, n_steps).
     distance_to_road_edge: Signed distance (in meters) to the nearest road edge
       in the scene. Shape: (n_samples, n_objects, n_steps).
-    offroad_indication: Boolean value indicating whether the object went
-      off-road, at any point in time. Shape: (n_samples, n_objects).
+    offroad_per_step: Boolean tensor indicating whether the object went
+      off-road. Shape: (n_samples, n_objects, n_steps).
   """
   object_id: tf.Tensor
   valid: tf.Tensor
@@ -86,10 +85,10 @@ class MetricFeatures:
   angular_speed: tf.Tensor
   angular_acceleration: tf.Tensor
   distance_to_nearest_object: tf.Tensor
-  collision_indication: tf.Tensor
+  collision_per_step: tf.Tensor
   time_to_collision: tf.Tensor
   distance_to_road_edge: tf.Tensor
-  offroad_indication: tf.Tensor
+  offroad_per_step: tf.Tensor
 
 
 def compute_metric_features(
@@ -114,7 +113,7 @@ def compute_metric_features(
   # history from the original scenario. These composite trajectories are used to
   # compute dynamics features, which require a few steps of context.
   simulated_trajectories = converters.joint_scene_to_trajectories(
-      joint_scene, scenario)
+      joint_scene, scenario, use_log_validity=use_log_validity)
   # Extract `ObjectTrajectories` from the original scenario, used for
   # log-comparison metrics (i.e. displacement error). These also need to be
   # aligned to the simulated trajectories.
@@ -211,8 +210,6 @@ def compute_metric_features(
   # Shape: (n_evaluated_objects, n_steps).
   is_colliding_per_step = tf.less(
       distances_to_objects, interaction_features.COLLISION_DISTANCE_THRESHOLD)
-  # Shape: (n_evaluated_objects,)
-  is_colliding = tf.reduce_any(is_colliding_per_step, axis=1)
 
   times_to_collision = (
       interaction_features.compute_time_to_collision_with_object_in_front(
@@ -253,7 +250,6 @@ def compute_metric_features(
   is_offroad_per_step = tf.greater(
       distances_to_road_edge, map_metric_features.OFFROAD_DISTANCE_THRESHOLD
   )
-  is_offroad = tf.reduce_any(is_offroad_per_step, axis=1)
 
   # Pack into `MetricFeatures`, also adding a batch dimension of 1 (except for
   # `object_id`).
@@ -266,10 +262,10 @@ def compute_metric_features(
       angular_speed=angular_speed[tf.newaxis],
       angular_acceleration=angular_accel[tf.newaxis],
       distance_to_nearest_object=distances_to_objects[tf.newaxis],
-      collision_indication=is_colliding[tf.newaxis],
+      collision_per_step=is_colliding_per_step[tf.newaxis],
       time_to_collision=times_to_collision[tf.newaxis],
       distance_to_road_edge=distances_to_road_edge[tf.newaxis],
-      offroad_indication=is_offroad[tf.newaxis],
+      offroad_per_step=is_offroad_per_step[tf.newaxis],
   )
 
 

src/waymo_open_dataset/wdl_limited/sim_agents_metrics/metric_features_test.py

@@ -18,6 +18,8 @@
 from waymo_open_dataset.utils import test_utils
 from waymo_open_dataset.utils.sim_agents import converters
 from waymo_open_dataset.utils.sim_agents import test_utils as sim_agents_test_utils
+from waymo_open_dataset.wdl_limited.sim_agents_metrics import interaction_features
+from waymo_open_dataset.wdl_limited.sim_agents_metrics import map_metric_features
 from waymo_open_dataset.wdl_limited.sim_agents_metrics import metric_features
 
 
@@ -48,8 +50,8 @@ def test_features_from_log_scenario_trajectories(self):
                      (batch_size, num_agents, num_steps))
     self.assertEqual(simulated_features.distance_to_nearest_object.shape,
                      (batch_size, num_agents, num_steps))
-    self.assertEqual(simulated_features.collision_indication.shape,
-                     (batch_size, num_agents))
+    self.assertEqual(simulated_features.collision_per_step.shape,
+                     (batch_size, num_agents, num_steps))
     self.assertEqual(
         simulated_features.time_to_collision.shape,
         (batch_size, num_agents, num_steps),
@@ -58,9 +60,8 @@ def test_features_from_log_scenario_trajectories(self):
         simulated_features.distance_to_road_edge.shape,
         (batch_size, num_agents, num_steps),
     )
-    self.assertEqual(
-        simulated_features.offroad_indication.shape, (batch_size, num_agents)
-    )
+    self.assertEqual(simulated_features.offroad_per_step.shape,
+                     (batch_size, num_agents, num_steps))
     # Since this joint scene is coming directly from logs, we should expect a
     # zero ADE.
     ade = simulated_features.average_displacement_error
@@ -72,8 +73,32 @@ def test_features_from_log_scenario_copies_log_validity(self):
     simulated_features = (
         metric_features.compute_metric_features(
             scenario, joint_scene, use_log_validity=True))
+
     self.assertFalse(tf.reduce_all(simulated_features.valid))
 
+    # Check that collision and offroad are correctly filtered by validity.
+    distance_to_nearest_with_valids = tf.where(
+        simulated_features.valid,
+        simulated_features.distance_to_nearest_object,
+        interaction_features.EXTREMELY_LARGE_DISTANCE)
+    distance_to_road_edge_with_valids = tf.where(
+        simulated_features.valid,
+        simulated_features.distance_to_road_edge,
+        -map_metric_features.EXTREMELY_LARGE_DISTANCE)
+
+    collisions = tf.less(
+        tf.reduce_min(distance_to_nearest_with_valids, axis=2),
+        interaction_features.COLLISION_DISTANCE_THRESHOLD,
+    )
+    offroad = tf.greater(
+        tf.reduce_max(distance_to_road_edge_with_valids, axis=2),
+        map_metric_features.OFFROAD_DISTANCE_THRESHOLD,
+    )
+    self.assertAllEqual(collisions, tf.reduce_any(
+        simulated_features.collision_per_step, axis=2))
+    self.assertAllEqual(offroad, tf.reduce_any(
+        simulated_features.offroad_per_step, axis=2))
+
   def test_features_from_linear_extrapolation_test_submission(self):
     scenario = test_utils.get_womd_test_scenario()
     submission = sim_agents_test_utils.load_test_submission()