Fixes: Align WOSAC metric calculation with original (#258)

* quick commit so you can read the code

* Improve naming of sampling argument to better describe its function.

* Ensure that initialization works with Carla maps or other, when sdc_index=-1.

* Simplify CARLA compatibility

* Replace num_maps by wosac_num_maps in all the eval scripts

* Comment about random baseline.

* Bug fix: do not reweight by the total weight (0.95).

* Add this back for now.

* Delete agent shrinking code.

* Ignore ttc metric when agents are not vehicles

* Update WOSAC weights to align with 2024 challenge since we don't have the traffic light metric.

* Update table

* Revert MAX_AGENTS to original.

* Update formatting.

---------

Co-authored-by: Wael Boumediene Doulazmi <wbd2016@cs713.hpc.nyu.edu>
Co-authored-by: Waël Doulazmi <waeldoulazmi@gmail.com>

Author: 3 people

docs/src/wosac.md

@@ -55,30 +55,18 @@ We provide baselines on a small curated dataset from the WOMD validation set wit
 
 | Method | Realism meta-score | Kinematic metrics | Interactive metrics | Map-based metrics | minADE | ADE |
 |--------|-------------------|-------------------|---------------------|-------------------|--------|------|
-| Ground-truth (UB) | 0.832 | 0.606 | 0.846 | 0.961 | 0 | 0 |
-| π_Base self-play RL | 0.737 | 0.319 | 0.789 | 0.938 | 10.834 | 11.317 |
-| [SMART-tiny-CLSFT](https://arxiv.org/abs/2412.05334) | 0.805 | 0.534 | 0.830 | 0.949 | 1.124 | 3.123 |
-| π_Random | 0.485 | 0.214 | 0.657 | 0.408 | 6.477 | 18.286 |
+| Ground-truth (UB) | 0.8179 | 0.6070 | 0.9590 | 0.8722 | 0 | 0 |
+| Self-play RL agent | 0.6750 | 0.2798 | 0.7966 | 0.7811 | 10.8057 | 11.4108 |
+| [SMART-tiny-CLSFT](https://arxiv.org/abs/2412.05334) | 0.7818 | 0.5200 | 0.8914 | 0.8378 | 1.1236 | 3.1231 |
+| Random | 0.4459 | 0.0506 | 0.7843 | 0.4704 | 23.5936 | 25.0097 |
 
 *Table: WOSAC baselines in PufferDrive on 229 selected clean held-out validation scenarios.*
 
-
-> ✏️ Download the dataset from [Hugging Face](https://huggingface.co/datasets/daphne-cornelisse/pufferdrive_wosac_val_clean) to reproduce these results or benchmark your policy.
-
-
-| Method | Realism meta-score | Kinematic metrics | Interactive metrics | Map-based metrics | minADE | ADE |
-| :--- | :--- | :--- | :--- | :--- | :--- | :--- |
-| Ground-truth (UB) | 0.833 | 0.574 | 0.864 | 0.958 | 0 | 0 |
-| π_Base self-play RL | 0.737 | 0.323 | 0.792 | 0.930 | 8.530 | 9.088 |
-| [SMART-tiny-CLSFT](https://arxiv.org/abs/2412.05334) | 0.795 | 0.504 | 0.832 | 0.932 | 1.182 | 2.857 |
-| π_Random | 0.497 | 0.238 | 0.656 | 0.430 | 6.395 | 18.617 |
-
-*Table: WOSAC baselines in PufferDrive on validation 10k dataset.*
-
-
-> ✏️ Download the dataset from [Hugging Face](https://huggingface.co/datasets/daphne-cornelisse/pufferdrive_womd_val) to reproduce these results or benchmark your policy.
+- **Random agent:** Following the [WOSAC 2023 paper](https://arxiv.org/abs/2305.12032), the random agent samples future trajectories by independently sampling (x, y, θ) at each timestep from a Gaussian distribution in the AV coordinate frame `(mu=1.0, sigma=0.1)`, producing uncorrelated random motion over the horizon of 80 steps.
+- **Goal-conditioned self-play RL agent**: An agent trained through self-play RL to reach the end point points ("goals") without colliding or going off-road. Baseline can be reproduced using the default settings in the `drive.ini` file with the Waymo dataset. We also open-source the weights of this policy, see `pufferlib/resources/drive/puffer_drive_weights` `.bin` and `.pt`.
 
 
+> ✏️ Download the dataset from [Hugging Face](https://huggingface.co/datasets/daphne-cornelisse/pufferdrive_wosac_val_clean) to reproduce these results or benchmark your policy.
 
 ## Evaluating trajectories
 

pufferlib/ocean/benchmark/evaluate_imported_trajectories.py

@@ -72,8 +72,8 @@ def evaluate_trajectories(simulated_trajectory_file, args):
     """
     env_name = "puffer_drive"
     args["env"]["map_dir"] = args["eval"]["map_dir"]
-    args["env"]["num_maps"] = args["eval"]["num_maps"]
-    args["env"]["use_all_maps"] = True
+    args["env"]["num_maps"] = args["eval"]["wosac_num_maps"]
+    args["env"]["sequential_map_sampling"] = True
     dataset_name = args["env"]["map_dir"].split("/")[-1]
 
     print(f"Running WOSAC realism evaluation with {dataset_name} dataset. \n")

pufferlib/ocean/benchmark/evaluator.py

@@ -48,9 +48,7 @@ def _compute_metametric(self, metrics: pd.Series) -> float:
             weight = self.metrics_config.getfloat(field_name, "metametric_weight")
             metric_score = metrics[likelihood_field_name]
             metametric += weight * metric_score
-
-        weight_sum = sum(self.metrics_config.getfloat(fn, "metametric_weight") for fn in _METRIC_FIELD_NAMES)
-        return metametric / weight_sum
+        return metametric
 
     def _get_histogram_params(self, metric_name: str):
         return (
@@ -121,6 +119,53 @@ def collect_simulated_trajectories(self, args, puffer_env, policy):
 
         return trajectories
 
+    def collect_wosac_random_baseline(self, puffer_env):
+        """
+        Random Baseline from Wosac 2023 paper
+        """
+        driver = puffer_env.driver_env
+        num_agents = puffer_env.observation_space.shape[0]
+
+        trajectories = {
+            "x": np.zeros((num_agents, self.num_rollouts, self.sim_steps), dtype=np.float32),
+            "y": np.zeros((num_agents, self.num_rollouts, self.sim_steps), dtype=np.float32),
+            "heading": np.zeros((num_agents, self.num_rollouts, self.sim_steps), dtype=np.float32),
+            "id": np.zeros((num_agents, self.num_rollouts, self.sim_steps), dtype=np.int32),
+        }
+
+        for rollout_idx in range(self.num_rollouts):
+            obs, info = puffer_env.reset()
+
+            # Do Initialization
+            agent_state = driver.get_global_agent_state()
+            trajectories["x"][:, rollout_idx, 0] = agent_state["x"]
+            trajectories["y"][:, rollout_idx, 0] = agent_state["y"]
+            trajectories["heading"][:, rollout_idx, 0] = agent_state["heading"]
+            trajectories["id"][:, rollout_idx, 0] = agent_state["id"]
+
+            # Update using Gaussian:
+            samples = np.random.normal(loc=1, scale=0.1, size=(num_agents, self.sim_steps, 3))
+            for time_idx in range(1, self.sim_steps):
+                dx, dy, d_heading = samples[:, time_idx, 0], samples[:, time_idx, 1], samples[:, time_idx, 2]
+                x, y, heading = (
+                    trajectories["x"][:, rollout_idx, time_idx - 1],
+                    trajectories["y"][:, rollout_idx, time_idx - 1],
+                    trajectories["heading"][:, rollout_idx, time_idx - 1],
+                )
+
+                cos_h = np.cos(heading)
+                sin_h = np.sin(heading)
+
+                x += dx * cos_h - dy * sin_h
+                y += dx * sin_h + dy * cos_h
+                heading += d_heading
+
+                trajectories["x"][:, rollout_idx, time_idx] = x
+                trajectories["y"][:, rollout_idx, time_idx] = y
+                trajectories["heading"][:, rollout_idx, time_idx] = heading
+
+        return trajectories
+
     def compute_metrics(
         self,
         ground_truth_trajectories: Dict,
@@ -159,6 +204,7 @@ def compute_metrics(
         ref_valid = ground_truth_trajectories["valid"]
         agent_length = agent_state["length"]
         agent_width = agent_state["width"]
+        is_vehicle = ground_truth_trajectories["is_vehicle"]
         scenario_ids = ground_truth_trajectories["scenario_id"]
 
         # We evaluate the metrics only for the Tracks to Predict.
@@ -172,6 +218,7 @@ def compute_metrics(
         eval_agent_length = agent_length[eval_mask]
         eval_agent_width = agent_width[eval_mask]
         eval_scenario_ids = scenario_ids[eval_mask]
+        eval_is_vehicle = is_vehicle[eval_mask]
 
         # Compute features
         # Kinematics-related features
@@ -337,60 +384,48 @@ def compute_metrics(
             sanity_check=False,
         )
 
-        speed_likelihood = np.exp(
-            metrics._reduce_average_with_validity(
-                linear_speed_log_likelihood,
-                speed_validity[:, 0, :],
-                axis=1,
-            )
+        speed_log_likelihood = metrics._reduce_average_with_validity(
+            linear_speed_log_likelihood,
+            speed_validity[:, 0, :],
+            axis=1,
         )
 
-        accel_likelihood = np.exp(
-            metrics._reduce_average_with_validity(
-                linear_accel_log_likelihood,
-                acceleration_validity[:, 0, :],
-                axis=1,
-            )
+        accel_log_likelihood = metrics._reduce_average_with_validity(
+            linear_accel_log_likelihood,
+            acceleration_validity[:, 0, :],
+            axis=1,
         )
 
-        angular_speed_likelihood = np.exp(
-            metrics._reduce_average_with_validity(
-                angular_speed_log_likelihood,
-                speed_validity[:, 0, :],
-                axis=1,
-            )
+        angular_speed_log_likelihood = metrics._reduce_average_with_validity(
+            angular_speed_log_likelihood,
+            speed_validity[:, 0, :],
+            axis=1,
         )
 
-        angular_accel_likelihood = np.exp(
-            metrics._reduce_average_with_validity(
-                angular_accel_log_likelihood,
-                acceleration_validity[:, 0, :],
-                axis=1,
-            )
+        angular_accel_log_likelihood = metrics._reduce_average_with_validity(
+            angular_accel_log_likelihood,
+            acceleration_validity[:, 0, :],
+            axis=1,
         )
 
-        distance_to_nearest_object_likelihood = np.exp(
-            metrics._reduce_average_with_validity(
-                distance_to_nearest_object_log_likelihood,
-                eval_ref_valid[:, 0, :],
-                axis=1,
-            )
+        distance_to_nearest_object_log_likelihood = metrics._reduce_average_with_validity(
+            distance_to_nearest_object_log_likelihood,
+            eval_ref_valid[:, 0, :],
+            axis=1,
         )
 
-        time_to_collision_likelihood = np.exp(
-            metrics._reduce_average_with_validity(
-                time_to_collision_log_likelihood,
-                eval_ref_valid[:, 0, :],
-                axis=1,
-            )
+        # TTC is computed only for vehicles
+        ttc_valid = eval_ref_valid & eval_is_vehicle[..., None]
+        time_to_collision_log_likelihood = metrics._reduce_average_with_validity(
+            time_to_collision_log_likelihood,
+            ttc_valid[:, 0, :],
+            axis=1,
         )
 
-        distance_to_road_edge_likelihood = np.exp(
-            metrics._reduce_average_with_validity(
-                distance_to_road_edge_log_likelihood,
-                eval_ref_valid[:, 0, :],
-                axis=1,
-            )
+        distance_to_road_edge_log_likelihood = metrics._reduce_average_with_validity(
+            distance_to_road_edge_log_likelihood,
+            eval_ref_valid[:, 0, :],
+            axis=1,
         )
 
         # Collision likelihood is computed by aggregating in time. For invalid objects
@@ -411,7 +446,6 @@ def compute_metrics(
             num_bins=2,
             use_bernoulli=True,
         )
-        collision_likelihood = np.exp(collision_log_likelihood)
 
         # Offroad likelihood (same pattern as collision)
         sim_offroad_indication = np.any(np.where(eval_ref_valid, sim_offroad_per_step, False), axis=2)
@@ -428,7 +462,6 @@ def compute_metrics(
             num_bins=2,
             use_bernoulli=True,
         )
-        offroad_likelihood = np.exp(offroad_log_likelihood)
 
         # Get agent IDs
         eval_agent_ids = ground_truth_trajectories["id"][eval_mask]
@@ -443,15 +476,15 @@ def compute_metrics(
                 "num_offroad_ref": ref_num_offroad.flatten(),
                 "ade": ade,
                 "min_ade": min_ade,
-                "likelihood_linear_speed": speed_likelihood,
-                "likelihood_linear_acceleration": accel_likelihood,
-                "likelihood_angular_speed": angular_speed_likelihood,
-                "likelihood_angular_acceleration": angular_accel_likelihood,
-                "likelihood_distance_to_nearest_object": distance_to_nearest_object_likelihood,
-                "likelihood_time_to_collision": time_to_collision_likelihood,
-                "likelihood_collision_indication": collision_likelihood,
-                "likelihood_distance_to_road_edge": distance_to_road_edge_likelihood,
-                "likelihood_offroad_indication": offroad_likelihood,
+                "likelihood_linear_speed": speed_log_likelihood,
+                "likelihood_linear_acceleration": accel_log_likelihood,
+                "likelihood_angular_speed": angular_speed_log_likelihood,
+                "likelihood_angular_acceleration": angular_accel_log_likelihood,
+                "likelihood_distance_to_nearest_object": distance_to_nearest_object_log_likelihood,
+                "likelihood_time_to_collision": time_to_collision_log_likelihood,
+                "likelihood_collision_indication": collision_log_likelihood,
+                "likelihood_distance_to_road_edge": distance_to_road_edge_log_likelihood,
+                "likelihood_offroad_indication": offroad_log_likelihood,
             }
         )
 
@@ -475,6 +508,10 @@ def compute_metrics(
             ]
         ].mean()
 
+        # Transform log-likelihoods to positive scores:
+        likelihood_cols = [c for c in scene_level_results.columns if "likelihood" in c]
+        scene_level_results[likelihood_cols] = np.exp(scene_level_results[likelihood_cols])
+
         scene_level_results["realism_meta_score"] = scene_level_results.apply(self._compute_metametric, axis=1)
         scene_level_results["num_agents"] = df.groupby("scenario_id").size()
         scene_level_results = scene_level_results[

pufferlib/ocean/benchmark/metrics_sanity_check.py

@@ -26,11 +26,11 @@ def replace_rollouts_with_gt(simulated_traj, gt_traj, num_replacements):
     return modified
 
 
-def run_validation_experiment(config, vecenv, policy):
+def run_validation_experiment(config, vecenv):
     evaluator = WOSACEvaluator(config)
 
     gt_trajectories = evaluator.collect_ground_truth_trajectories(vecenv)
-    simulated_trajectories = evaluator.collect_simulated_trajectories(config, vecenv, policy)
+    simulated_trajectories = evaluator.collect_wosac_random_baseline(vecenv)
     agent_state = vecenv.driver_env.get_global_agent_state()
     road_edge_polylines = vecenv.driver_env.get_road_edge_polylines()
 
@@ -90,19 +90,17 @@ def main():
     config["eval"]["enabled"] = True
     config["eval"]["wosac_num_rollouts"] = 32
     config["env"]["map_dir"] = config["eval"]["map_dir"]
-    config["env"]["num_maps"] = config["eval"]["num_maps"]
-    config["env"]["use_all_maps"] = True
+    config["env"]["num_maps"] = config["eval"]["wosac_num_maps"]
+    config["env"]["sequential_map_sampling"] = True
 
-    config["env"]["num_agents"] = config["eval"]["wosac_num_agents"]
     config["env"]["init_mode"] = config["eval"]["wosac_init_mode"]
     config["env"]["control_mode"] = config["eval"]["wosac_control_mode"]
     config["env"]["init_steps"] = config["eval"]["wosac_init_steps"]
     config["env"]["goal_behavior"] = config["eval"]["wosac_goal_behavior"]
 
     vecenv = load_env(args.env, config)
-    policy = load_policy(config, vecenv, args.env)
 
-    results = run_validation_experiment(config, vecenv, policy)
+    results = run_validation_experiment(config, vecenv)
     print("\n" + format_results_table(results))
 
 

pufferlib/ocean/benchmark/visual_sanity_check.py

@@ -104,8 +104,8 @@ def main():
     config["env"]["goal_behavior"] = config["eval"]["wosac_goal_behavior"]
 
     config["env"]["map_dir"] = config["eval"]["map_dir"]
-    config["env"]["num_maps"] = config["eval"]["num_maps"]
-    config["env"]["use_all_maps"] = True
+    config["env"]["num_maps"] = config["eval"]["wosac_num_maps"]
+    config["env"]["sequential_map_sampling"] = True
 
     vecenv = load_env(args.env, config)
     policy = load_policy(config, vecenv, args.env)

pufferlib/ocean/benchmark/wosac.ini

@@ -1,5 +1,6 @@
-# Taken from WOSAC 2025 Sim Agents Challenge configuration
-# Link: https://github.com/waymo-research/waymo-open-dataset/blob/master/src/waymo_open_dataset/wdl_limited/sim_agents_metrics/challenge_2025_sim_agents_config.textproto
+# Taken from WOSAC 2024 Sim Agents Challenge configuration
+# Note: 2024 challenge does not include traffic_light_violation metric (weight = 0.0)
+# Traffic light violation was added in 2025 challenge with weight 0.05
 
 [linear_speed]
 histogram.min_val = 0.0
@@ -51,7 +52,7 @@ histogram.max_val = 40
 histogram.num_bins = 10
 histogram.additive_smoothing_pseudocount = 0.1
 independent_timesteps = true
-metametric_weight = 0.05
+metametric_weight = 0.1
 
 [offroad_indication]
 bernoulli = true
@@ -67,4 +68,4 @@ metametric_weight = 0.1
 
 [traffic_light_violation]
 bernoulli = true
-metametric_weight = 0.05
+metametric_weight = 0.0

pufferlib/ocean/drive/binding.c

@@ -76,20 +76,21 @@ static PyObject *my_shared(PyObject *self, PyObject *args, PyObject *kwargs) {
     int init_steps = unpack(kwargs, "init_steps");
     int goal_behavior = unpack(kwargs, "goal_behavior");
     float goal_target_distance = unpack(kwargs, "goal_target_distance");
-    int use_all_maps = unpack(kwargs, "use_all_maps");
+    int sequential_map_sampling = unpack(kwargs, "sequential_map_sampling");
     clock_gettime(CLOCK_REALTIME, &ts);
     srand(ts.tv_nsec);
     int total_agent_count = 0;
     int env_count = 0;
-    int max_envs = use_all_maps ? num_maps : num_agents;
+    int max_envs = sequential_map_sampling ? num_maps : num_agents;
     int map_idx = 0;
     int maps_checked = 0;
     PyObject *agent_offsets = PyList_New(max_envs + 1);
     PyObject *map_ids = PyList_New(max_envs);
     // getting env count
-    while (use_all_maps ? map_idx < max_envs : total_agent_count < num_agents && env_count < max_envs) {
+    while (sequential_map_sampling ? map_idx < max_envs : total_agent_count < num_agents && env_count < max_envs) {
         char map_file[512];
-        int map_id = use_all_maps ? map_idx++ : rand() % num_maps;
+        // Take the next map in sequence or a random map
+        int map_id = sequential_map_sampling ? map_idx++ : rand() % num_maps;
         Drive *env = calloc(1, sizeof(Drive));
         env->init_mode = init_mode;
         env->control_mode = control_mode;
@@ -102,7 +103,7 @@ static PyObject *my_shared(PyObject *self, PyObject *args, PyObject *kwargs) {
 
         // Skip map if it doesn't contain any controllable agents
         if (env->active_agent_count == 0) {
-            if (!use_all_maps) {
+            if (!sequential_map_sampling) {
                 maps_checked++;
 
                 // Safeguard: if we've checked all available maps and found no active agents, raise an error
@@ -154,7 +155,7 @@ static PyObject *my_shared(PyObject *self, PyObject *args, PyObject *kwargs) {
     }
     // printf("Generated %d environments to cover %d agents (requested %d agents)\n", env_count, total_agent_count,
     // num_agents);
-    if (!use_all_maps && total_agent_count >= num_agents) {
+    if (!sequential_map_sampling && total_agent_count >= num_agents) {
         total_agent_count = num_agents;
     }
     PyObject *final_total_agent_count = PyLong_FromLong(total_agent_count);