More complete and Robust evaluation of WOSAC rollouts

pufferlib/ocean/benchmark/evaluate_imported_trajectories.py

@@ -1,49 +1,30 @@
 import sys
 import pickle
 import numpy as np
-from scipy.spatial import cKDTree
 import pufferlib.pufferl as pufferl
 from pufferlib.ocean.benchmark.evaluator import WOSACEvaluator
 
 
-def align_trajectories_by_initial_position(simulated, ground_truth, tolerance=1e-4):
-    """
-    If the trajectories where generated using the same dataset, then regardless of the algorithm the initial positions should be the same.
-    We use this information to align the trajectories for WOSAC evaluation.
-
-    Ideally we would not have to use a tolerance, but the preprocessing in SMART shifts some values by around 2-e5 for some agents.
-
-    Also, the preprocessing in SMART messes up some heading values, so I decided not to include heading.
-
-    Idea of this script, use a nearest neighbor algorithm to associate all initial positions in gt to positions in simulated,
-    and check that everyone matching respects the tolerance and there are no duplicates.
-    """
-
-    sim_pos = np.stack([simulated["x"][:, 0, 0], simulated["y"][:, 0, 0], simulated["z"][:, 0, 0]], axis=1).astype(
-        np.float64
-    )
+def align_trajectories(simulated, ground_truth):
+    # Idea is to use the (scenario_id, id) pair to reindex simulated_trajectories in order to align it with GT
+    gt_scenario_ids = ground_truth["scenario_id"][:, 0]
+    sim_scenario_ids = simulated["scenario_id"][:, 0, 0]
 
-    gt_pos = np.stack(
-        [ground_truth["x"][:, 0, 0], ground_truth["y"][:, 0, 0], ground_truth["z"][:, 0, 0]], axis=1
-    ).astype(np.float64)
+    gt_ids = ground_truth["id"][:, 0]
+    sim_ids = simulated["id"][:, 0, 0]
 
-    tree = cKDTree(sim_pos)
+    lookup = {(s_id, a_id): idx for idx, (s_id, a_id) in enumerate(zip(sim_scenario_ids, sim_ids))}
 
-    dists, indices = tree.query(gt_pos, k=1)
+    try:
+        indices = [lookup[(s, i)] for (s, i) in zip(gt_scenario_ids, gt_ids)]
+        indices = np.array(indices, dtype=int)
+    except KeyError:
+        print("An agent present in the GT is missing in your simulation")
+        raise
 
-    tol_check = dists <= tolerance
+    sim_traj = {k: v[indices] for k, v in simulated.items()}
 
-    if not np.all(tol_check):
-        max_dist = np.max(dists)
-        raise ValueError(f"Didn't find a match for {np.sum(~tol_check)} agents, tolerance broken by {max_dist}m.")
-
-    if len(set(indices)) != len(indices):
-        raise ValueError("Duplicate matching found, I am sorry but this likely indicates that your data is wrong")
-
-    reordered_sim = {}
-    for key, val in simulated.items():
-        reordered_sim[key] = val[indices]
-    return reordered_sim
+    return sim_traj
 
 
 def check_alignment(simulated, ground_truth, tolerance=1e-4):
@@ -97,30 +78,26 @@ def evaluate_trajectories(simulated_trajectory_file, args):
 
     print(f"Number of scenarios: {len(np.unique(gt_trajectories['scenario_id']))}")
     print(f"Number of controlled agents: {num_agents_gt}")
-    print(f"Number of evaluated agents: {np.sum(gt_trajectories['id'] >= 0)}")
+    print(f"Number of evaluated agents: {gt_trajectories['is_track_to_predict'].sum()}")
 
     print(f"Loading simulated trajectories from {simulated_trajectory_file}...")
     with open(simulated_trajectory_file, "rb") as f:
         sim_trajectories = pickle.load(f)
 
-    if sim_trajectories["x"].shape[0] != gt_trajectories["x"].shape[0]:
-        print("\nThe number of agents in simulated and ground truth trajectories do not match.")
-        print("This is okay if you are running this script on a subset of the val dataset")
-        print("But please also check that in drive.h MAX_AGENTS is set to 256 and recompile")
-
-    if not check_alignment(sim_trajectories, gt_trajectories):
-        print("\nTrajectories are not aligned, trying to align them, if it fails consider changing the tolerance.")
-        sim_trajectories = align_trajectories_by_initial_position(sim_trajectories, gt_trajectories)
-        assert check_alignment(sim_trajectories, gt_trajectories), (
-            "There might be an issue with the way you generated your data."
-        )
-        print("Alignment successful")
-    else:
-        sim_trajectories = {k: v[:num_agents_gt] for k, v in sim_trajectories.items()}
-
-    # Evaluator code expects to have matching ids between gt and sim trajectories
-    # Since alignment is checked it is safe to do that
-    sim_trajectories["id"][:] = gt_trajectories["id"][..., None]
+    num_agents_sim = sim_trajectories["x"].shape[0]
+    assert num_agents_sim >= num_agents_gt, (
+        "There is less agents in your simulation than in the GT, so the computation won't be valid"
+    )
+
+    if num_agents_sim > num_agents_gt:
+        print("If you are evaluating on a subset of your trajectories it is fine.")
+        print("\n Else, you should consider changing the value of MAX_AGENTS in drive.h and compile")
+
+    sim_trajectories = align_trajectories(sim_trajectories, gt_trajectories)
+
+    assert check_alignment(sim_trajectories, gt_trajectories), (
+        "There might be an issue with the way you generated your data."
+    )
 
     agent_state = vecenv.driver_env.get_global_agent_state()
     road_edge_polylines = vecenv.driver_env.get_road_edge_polylines()

pufferlib/ocean/benchmark/evaluator.py

@@ -192,7 +192,7 @@ def compute_metrics(
             "Agent IDs don't match between simulated and ground truth trajectories"
         )
 
-        eval_mask = ground_truth_trajectories["id"][:, 0] >= 0
+        eval_mask = ground_truth_trajectories["is_track_to_predict"][:, 0]
 
         # Extract trajectories
         sim_x = simulated_trajectories["x"]

pufferlib/ocean/drive/drive.h

@@ -321,7 +321,7 @@ struct Drive {
     int goal_behavior;
     float goal_target_distance;
     char *ini_file;
-    char *scenario_id;
+    char scenario_id[16];
     int collision_behavior;
     int offroad_behavior;
     int sdc_track_index;
@@ -386,6 +386,9 @@ Entity *load_map_binary(const char *filename, Drive *env) {
     if (!file)
         return NULL;
 
+    // Read scenario_id
+    fread(env->scenario_id, sizeof(char), 16, file);
+
     // Read sdc_track_index
     fread(&env->sdc_track_index, sizeof(int), 1, file);
 
@@ -1662,16 +1665,16 @@ void move_dynamics(Drive *env, int action_idx, int agent_idx) {
     return;
 }
 
-static inline int get_track_id_or_placeholder(Drive *env, int agent_idx) {
+static inline int is_in_track_to_predicts(Drive *env, int agent_idx) {
     if (env->tracks_to_predict_indices == NULL || env->num_tracks_to_predict == 0) {
-        return -1;
+        return 0;
     }
     for (int k = 0; k < env->num_tracks_to_predict; k++) {
         if (env->tracks_to_predict_indices[k] == agent_idx) {
-            return env->tracks_to_predict_indices[k];
+            return 1;
         }
     }
-    return -1;
+    return 0;
 }
 
 void c_get_global_agent_state(Drive *env, float *x_out, float *y_out, float *z_out, float *heading_out, int *id_out,
@@ -1685,20 +1688,24 @@ void c_get_global_agent_state(Drive *env, float *x_out, float *y_out, float *z_o
         y_out[i] = agent->y + env->world_mean_y;
         z_out[i] = agent->z;
         heading_out[i] = agent->heading;
-        id_out[i] = get_track_id_or_placeholder(env, agent_idx);
+        id_out[i] = agent->id;
         length_out[i] = agent->length;
         width_out[i] = agent->width;
     }
 }
 
 void c_get_global_ground_truth_trajectories(Drive *env, float *x_out, float *y_out, float *z_out, float *heading_out,
-                                            int *valid_out, int *id_out, int *is_vehicle_out, int *scenario_id_out) {
+                                            int *valid_out, int *id_out, bool *is_vehicle_out,
+                                            bool *is_track_to_predict_out, char *scenario_id_out) {
     for (int i = 0; i < env->active_agent_count; i++) {
         int agent_idx = env->active_agent_indices[i];
         Entity *agent = &env->entities[agent_idx];
-        id_out[i] = get_track_id_or_placeholder(env, agent_idx);
+        id_out[i] = agent->id;
         is_vehicle_out[i] = agent->type == VEHICLE;
-        scenario_id_out[i] = agent->scenario_id;
+        is_track_to_predict_out[i] = is_in_track_to_predicts(env, agent_idx);
+
+        // The scenario_id is an array of 16 char
+        memcpy(scenario_id_out + (i * 16), env->scenario_id, 16);
 
         for (int t = env->init_steps; t < agent->array_size; t++) {
             int out_idx = i * (agent->array_size - env->init_steps) + (t - env->init_steps);
@@ -1724,13 +1731,16 @@ void c_get_road_edge_counts(Drive *env, int *num_polylines_out, int *total_point
     *total_points_out = points;
 }
 
-void c_get_road_edge_polylines(Drive *env, float *x_out, float *y_out, int *lengths_out, int *scenario_ids_out) {
+void c_get_road_edge_polylines(Drive *env, float *x_out, float *y_out, int *lengths_out, char *scenario_ids_out) {
     int poly_idx = 0, pt_idx = 0;
     for (int i = env->num_objects; i < env->num_entities; i++) {
         Entity *e = &env->entities[i];
         if (e->type == ROAD_EDGE) {
             lengths_out[poly_idx] = e->array_size;
-            scenario_ids_out[poly_idx] = e->scenario_id;
+
+            char *scenario_id_ptr = scenario_ids_out + poly_idx * 16;
+            memcpy(scenario_id_ptr, env->scenario_id, 16);
+
             for (int j = 0; j < e->array_size; j++) {
                 x_out[pt_idx] = e->traj_x[j] + env->world_mean_x;
                 y_out[pt_idx] = e->traj_y[j] + env->world_mean_y;

pufferlib/ocean/drive/drive.py

@@ -326,8 +326,9 @@ def get_ground_truth_trajectories(self):
             "heading": np.zeros((num_agents, self.episode_length - self.init_steps), dtype=np.float32),
             "valid": np.zeros((num_agents, self.episode_length - self.init_steps), dtype=np.int32),
             "id": np.zeros(num_agents, dtype=np.int32),
-            "is_vehicle": np.zeros(num_agents, dtype=np.int32),
-            "scenario_id": np.zeros(num_agents, dtype=np.int32),
+            "is_vehicle": np.zeros(num_agents, dtype=bool),
+            "is_track_to_predict": np.zeros(num_agents, dtype=bool),
+            "scenario_id": np.zeros(num_agents, dtype="S16"),
         }
 
         binding.vec_get_global_ground_truth_trajectories(
@@ -339,12 +340,15 @@ def get_ground_truth_trajectories(self):
             trajectories["valid"],
             trajectories["id"],
             trajectories["is_vehicle"],
+            trajectories["is_track_to_predict"],
             trajectories["scenario_id"],
         )
 
         for key in trajectories:
             trajectories[key] = trajectories[key][:, None]
 
+        trajectories["scenario_id"] = trajectories["scenario_id"].astype(str)
+
         return trajectories
 
     def get_road_edge_polylines(self):
@@ -360,7 +364,7 @@ def get_road_edge_polylines(self):
             "x": np.zeros(total_points, dtype=np.float32),
             "y": np.zeros(total_points, dtype=np.float32),
             "lengths": np.zeros(num_polylines, dtype=np.int32),
-            "scenario_id": np.zeros(num_polylines, dtype=np.int32),
+            "scenario_id": np.zeros(num_polylines, dtype="S16"),
         }
 
         binding.vec_get_road_edge_polylines(
@@ -371,6 +375,8 @@ def get_road_edge_polylines(self):
             polylines["scenario_id"],
         )
 
+        polylines["scenario_id"] = polylines["scenario_id"].astype(str)
+
         return polylines
 
     def render(self):
@@ -439,6 +445,11 @@ def save_map_binary(map_data, output_file, unique_map_id):
         sdc_track_index = metadata.get("sdc_track_index", -1)  # -1 as default if not found
         tracks_to_predict = metadata.get("tracks_to_predict", [])
 
+        # Write original scenario_id
+        # NOTE: Think about how it might break CARLA ? (can add a default value)
+        scenario_id = map_data["scenario_id"]
+        f.write(struct.pack("16s", scenario_id.encode("utf-8")))
+
         # Write sdc_track_index
         f.write(struct.pack("i", sdc_track_index))
 

pufferlib/ocean/env_binding.h

@@ -727,8 +727,8 @@ static PyObject *vec_get_global_agent_state(PyObject *self, PyObject *args) {
 }
 
 static PyObject *get_ground_truth_trajectories(PyObject *self, PyObject *args) {
-    if (PyTuple_Size(args) != 8) {
-        PyErr_SetString(PyExc_TypeError, "get_ground_truth_trajectories requires 8 arguments");
+    if (PyTuple_Size(args) != 9) {
+        PyErr_SetString(PyExc_TypeError, "get_ground_truth_trajectories requires 9 arguments");
         return NULL;
     }
 
@@ -747,11 +747,12 @@ static PyObject *get_ground_truth_trajectories(PyObject *self, PyObject *args) {
     PyObject *valid_arr = PyTuple_GetItem(args, 5);
     PyObject *id_arr = PyTuple_GetItem(args, 6);
     PyObject *is_vehicle_arr = PyTuple_GetItem(args, 7);
-    PyObject *scenario_id_arr = PyTuple_GetItem(args, 8);
+    PyObject *is_track_to_predict_arr = PyTuple_GetItem(args, 8);
+    PyObject *scenario_id_arr = PyTuple_GetItem(args, 9);
 
     if (!PyArray_Check(x_arr) || !PyArray_Check(y_arr) || !PyArray_Check(z_arr) || !PyArray_Check(heading_arr) ||
         !PyArray_Check(valid_arr) || !PyArray_Check(id_arr) || !PyArray_Check(is_vehicle_arr) ||
-        !PyArray_Check(scenario_id_arr)) {
+        !PyArray_Check(is_track_to_predict_arr) || !PyArray_Check(scenario_id_arr)) {
         PyErr_SetString(PyExc_TypeError, "All output arrays must be NumPy arrays");
         return NULL;
     }
@@ -762,18 +763,19 @@ static PyObject *get_ground_truth_trajectories(PyObject *self, PyObject *args) {
     float *heading_data = (float *)PyArray_DATA((PyArrayObject *)heading_arr);
     int *valid_data = (int *)PyArray_DATA((PyArrayObject *)valid_arr);
     int *id_data = (int *)PyArray_DATA((PyArrayObject *)id_arr);
-    int *is_vehicle_data = (int *)PyArray_DATA((PyArrayObject *)is_vehicle_arr);
-    int *scenario_id_data = (int *)PyArray_DATA((PyArrayObject *)scenario_id_arr);
+    bool *is_vehicle_data = (bool *)PyArray_DATA((PyArrayObject *)is_vehicle_arr);
+    bool *is_track_to_predict_data = (bool *)PyArray_DATA((PyArrayObject *)is_track_to_predict_arr);
+    char *scenario_id_data = (char *)PyArray_DATA((PyArrayObject *)scenario_id_arr);
 
     c_get_global_ground_truth_trajectories(drive, x_data, y_data, z_data, heading_data, valid_data, id_data,
-                                           is_vehicle_data, scenario_id_data);
+                                           is_vehicle_data, is_track_to_predict_data, scenario_id_data);
 
     Py_RETURN_NONE;
 }
 
 static PyObject *vec_get_global_ground_truth_trajectories(PyObject *self, PyObject *args) {
-    if (PyTuple_Size(args) != 9) {
-        PyErr_SetString(PyExc_TypeError, "vec_get_global_ground_truth_trajectories requires 9 arguments");
+    if (PyTuple_Size(args) != 10) {
+        PyErr_SetString(PyExc_TypeError, "vec_get_global_ground_truth_trajectories requires 10 arguments");
         return NULL;
     }
 
@@ -790,11 +792,12 @@ static PyObject *vec_get_global_ground_truth_trajectories(PyObject *self, PyObje
     PyObject *valid_arr = PyTuple_GetItem(args, 5);
     PyObject *id_arr = PyTuple_GetItem(args, 6);
     PyObject *is_vehicle_arr = PyTuple_GetItem(args, 7);
-    PyObject *scenario_id_arr = PyTuple_GetItem(args, 8);
+    PyObject *is_track_to_predict_arr = PyTuple_GetItem(args, 8);
+    PyObject *scenario_id_arr = PyTuple_GetItem(args, 9);
 
     if (!PyArray_Check(x_arr) || !PyArray_Check(y_arr) || !PyArray_Check(z_arr) || !PyArray_Check(heading_arr) ||
         !PyArray_Check(valid_arr) || !PyArray_Check(id_arr) || !PyArray_Check(is_vehicle_arr) ||
-        !PyArray_Check(scenario_id_arr)) {
+        !PyArray_Check(is_track_to_predict_arr) || !PyArray_Check(scenario_id_arr)) {
         PyErr_SetString(PyExc_TypeError, "All output arrays must be NumPy arrays");
         return NULL;
     }
@@ -806,6 +809,7 @@ static PyObject *vec_get_global_ground_truth_trajectories(PyObject *self, PyObje
     PyArrayObject *valid_array = (PyArrayObject *)valid_arr;
     PyArrayObject *id_array = (PyArrayObject *)id_arr;
     PyArrayObject *is_vehicle_array = (PyArrayObject *)is_vehicle_arr;
+    PyArrayObject *is_track_to_predict_array = (PyArrayObject *)is_track_to_predict_arr;
     PyArrayObject *scenario_id_array = (PyArrayObject *)scenario_id_arr;
 
     // Get base pointers to the arrays
@@ -815,8 +819,9 @@ static PyObject *vec_get_global_ground_truth_trajectories(PyObject *self, PyObje
     float *heading_base = (float *)PyArray_DATA(heading_array);
     int *valid_base = (int *)PyArray_DATA(valid_array);
     int *id_base = (int *)PyArray_DATA(id_array);
-    int *is_vehicle_base = (int *)PyArray_DATA(is_vehicle_array);
-    int *scenario_id_base = (int *)PyArray_DATA(scenario_id_array);
+    bool *is_vehicle_base = (bool *)PyArray_DATA(is_vehicle_array);
+    bool *is_track_to_predict_base = (bool *)PyArray_DATA(is_track_to_predict_array);
+    char *scenario_id_base = (char *)PyArray_DATA(scenario_id_array);
 
     // Get number of timesteps from array shape
     npy_intp *x_shape = PyArray_DIMS(x_array);
@@ -829,10 +834,10 @@ static PyObject *vec_get_global_ground_truth_trajectories(PyObject *self, PyObje
     for (int i = 0; i < vec->num_envs; i++) {
         Drive *drive = (Drive *)vec->envs[i];
 
-        c_get_global_ground_truth_trajectories(drive, &x_base[traj_offset], &y_base[traj_offset], &z_base[traj_offset],
-                                               &heading_base[traj_offset], &valid_base[traj_offset],
-                                               &id_base[agent_offset], &is_vehicle_base[agent_offset],
-                                               &scenario_id_base[agent_offset]);
+        c_get_global_ground_truth_trajectories(
+            drive, &x_base[traj_offset], &y_base[traj_offset], &z_base[traj_offset], &heading_base[traj_offset],
+            &valid_base[traj_offset], &id_base[agent_offset], &is_vehicle_base[agent_offset],
+            &is_track_to_predict_base[agent_offset], &scenario_id_base[agent_offset * 16]);
 
         // Move offsets forward
         agent_offset += drive->active_agent_count;
@@ -882,15 +887,15 @@ static PyObject *vec_get_road_edge_polylines(PyObject *self, PyObject *args) {
     float *x_base = (float *)PyArray_DATA((PyArrayObject *)x_arr);
     float *y_base = (float *)PyArray_DATA((PyArrayObject *)y_arr);
     int *lengths_base = (int *)PyArray_DATA((PyArrayObject *)lengths_arr);
-    int *scenario_ids_base = (int *)PyArray_DATA((PyArrayObject *)scenario_ids_arr);
+    char *scenario_ids_base = (char *)PyArray_DATA((PyArrayObject *)scenario_ids_arr);
 
     int poly_offset = 0, pt_offset = 0;
     for (int i = 0; i < vec->num_envs; i++) {
         Drive *drive = (Drive *)vec->envs[i];
         int np, tp;
         c_get_road_edge_counts(drive, &np, &tp);
         c_get_road_edge_polylines(drive, &x_base[pt_offset], &y_base[pt_offset], &lengths_base[poly_offset],
-                                  &scenario_ids_base[poly_offset]);
+                                  &scenario_ids_base[poly_offset * 16]);
         poly_offset += np;
         pt_offset += tp;
     }

pufferlib/pufferl.py

@@ -1067,7 +1067,7 @@ def eval(env_name, args=None, vecenv=None, policy=None):
 
         print(f"Number of scenarios: {len(np.unique(gt_trajectories['scenario_id']))}")
         print(f"Number of controlled agents: {gt_trajectories['x'].shape[0]}")
-        print(f"Number of evaluated agents: {np.sum(gt_trajectories['id'] >= 0)}")
+        print(f"Number of evaluated agents: {gt_trajectories['is_track_to_predict'].sum()}")
 
         # Roll out trained policy in the simulator
         simulated_trajectories = evaluator.collect_simulated_trajectories(args, vecenv, policy)