diff --git a/pufferlib/ocean/benchmark/evaluate_imported_trajectories.py b/pufferlib/ocean/benchmark/evaluate_imported_trajectories.py index ceb95cff1..8e741bf4b 100644 --- a/pufferlib/ocean/benchmark/evaluate_imported_trajectories.py +++ b/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() diff --git a/pufferlib/ocean/benchmark/evaluator.py b/pufferlib/ocean/benchmark/evaluator.py index 003bfad27..4940ae439 100644 --- a/pufferlib/ocean/benchmark/evaluator.py +++ b/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"] diff --git a/pufferlib/ocean/drive/drive.h b/pufferlib/ocean/drive/drive.h index 4df3bf021..f34539337 100644 --- a/pufferlib/ocean/drive/drive.h +++ b/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; diff --git a/pufferlib/ocean/drive/drive.py b/pufferlib/ocean/drive/drive.py index e85eb7322..16c4bbb1b 100644 --- a/pufferlib/ocean/drive/drive.py +++ b/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)) diff --git a/pufferlib/ocean/env_binding.h b/pufferlib/ocean/env_binding.h index c6b766776..b20afa07f 100644 --- a/pufferlib/ocean/env_binding.h +++ b/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,7 +887,7 @@ 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++) { @@ -890,7 +895,7 @@ static PyObject *vec_get_road_edge_polylines(PyObject *self, PyObject *args) { 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; } diff --git a/pufferlib/pufferl.py b/pufferlib/pufferl.py index c797f3244..27795e0b9 100644 --- a/pufferlib/pufferl.py +++ b/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)