Adding a New Model to ReactSim-Bench

This guide describes the work needed to evaluate a new behavior world model on ReactSim-Bench.

Data Cache

ReactSim-Bench provides a shared preprocessed cache for all baselines. You can use the cache when adding a new model and write a data adapter for your model instead of extracting data from the raw nuPlan dataset.

In the data cache, each .pkl file is one scenario, and its name is <map_name>_<scene_token>_<sample_token>.pkl. Each file contains a dict with the fields below:

Top-level fields:

Field	Shape / type	Meaning
schema_version	int	Cache schema version.
scenario_id	str	<log_name>_<sample_token> identifier.
scenario_token	str	nuPlan scenario token.
sample_token	str	Initial lidar sample token.
scene_token	str	nuPlan scene token.
scenario_type	str	nuPlan scenario type, or unknown.
log_name	str	Source nuPlan log name.
map_name	str	nuPlan map name.
current_time_index	int	Current frame index in the 91-frame sequence.
timestamps_us	(T,) int64	Frame timestamps in microseconds.
timestamps_seconds	(T,) float32	Relative frame timestamps in seconds.
sdc_track_index	int	Ego vehicle index in agents.
agents	dict	Agent trajectories and metadata.
map	dict	Local map geometry around the scenario.
traffic_lights	dict	Traffic-light states and stop points.

agents fields:

Field	Shape / type	Meaning
ids	list[str], length A	Agent track ids. The ego id is ego:<scenario_id>.
types	(A,) uint8	Agent type: 0 vehicle, 1 pedestrian, 2 bicycle, 3 other.
role	(A, 3) bool	Role flags: ego, prediction target, observed agent.
position	(A, T, 2) float64	Global x/y position in meters.
heading	(A, T) float32	Heading angle in radians.
velocity	(A, T, 2) float32	Global x/y velocity in meters per second.
size	(A, 3) float32	Agent length, width, and height in meters.
valid	(A, T) bool	Whether each agent state is valid at each frame.
ego_index	int	Ego vehicle index. Same value as sdc_track_index.

map fields:

Field	Shape / type	Meaning
ids	list[str], length M	nuPlan map element ids.
kind	(M,) uint8	Map type: 0 lane, 1 lane connector, 2 road edge, 3 crosswalk, 4 stop line.
point_range	(M, 2) int32	Start/end indices into points for each map element.
points	(P, 2) float64	Flattened global x/y map points in meters.
speed_limit_mps	(M,) float32	Speed limit in meters per second, or -1 if unavailable.
has_stop_line	(M,) bool	Whether the lane or connector has an associated stop line.

traffic_lights fields:

Field	Shape / type	Meaning
lane_ids	list[str], length L	Lane connector ids controlled by traffic lights.
states	(T, L) uint8	State: 0 none, 1 unknown, 2 red, 3 yellow, 4 green.
valid	(T, L) bool	Whether the traffic-light state is observed at each frame.
stop_points	(L, 2) float64	Global x/y stop point for each controlled lane connector.
polyline_index	(L,) int32	Index of the controlled lane connector in map, or -1 if absent.

Notation:

• A: number of agents.
• T: number of frames. T=91, with 11 history frames and 80 future frames.
• M: number of map elements.
• P: total number of stored map points.
• L: number of traffic-light lane connectors.
• current_time_index: current frame index, fixed to 10.

Evaluation

Different baselines use very different model input formats during inference, so ReactSim-Bench does not enforce a strict closed-loop inference interface. Instead, it provides a small set of shared utilities for building a model-specific closed-loop script.

Current model entrypoints:

MTR/tools/closed_loop_test.py
CTG/scripts/closed_loop_test.py
VBD/script/closed_loop_test.py
SMART_CATK_TrajTok/tools/closed_loop_test.py

Useful shared components:

Component	Location	Use
parse_closed_loop_args	reactsim_data_pipeline/closed_loop/cli_args.py	Builds a consistent command-line entry for scenario paths, model paths, rollout I/O, replanning, and visualization.
load_input_sample	reactsim_data_pipeline/closed_loop/io.py	Loads a scenario cache and converts it to the shared closed-loop format when needed.
extract_custom_ego_trajectory	reactsim_data_pipeline/closed_loop/io.py	Loads the collected custom ego behavior for each scenario.
save_scene_rollout, load_scene_rollout	reactsim_data_pipeline/closed_loop/io.py	Saves or reloads per-scenario simulated trajectories.
batch_update_sample	reactsim_data_pipeline/closed_loop/state.py	Advances the shared scene state by one step using model predictions and the custom ego state.
NuPlanVisualizer	reactsim_data_pipeline/visualization/nuplan_visualizer.py	Draws cache samples, closed-loop image frames, and rollout videos.
build_video_task, render_videos	reactsim_data_pipeline/closed_loop/io.py	Renders videos from saved or in-memory rollouts.
process_scene_output	reactsim_data_pipeline/closed_loop/metrics.py	Saves rollouts if requested and sends one scenario to metric evaluation.

The usual closed-loop flow is:

args = parse_closed_loop_args(...)  # [tools] CLI helper
model = load_model(args.model_path, args.config_file)  # model-specific
evaluator, aggregator = build_metric_evaluator(...)  # [tools] metrics helper

for scenario_batch in scenario_batches:
    scenes = [load_input_sample(path) for path in scenario_batch]  # [tools] shared I/O
    ego_traj = extract_custom_ego_trajectory(...)  # [tools] shared I/O
    rollout = []

    for step in range(args.num_steps):
        if step % args.replan_interval == 0:
            model_input = build_model_input(scenes)  # model-specific
            model_output = model(model_input)  # model-specific
            pred_horizon = decode_to_global_states(model_output)  # model-specific

        pred_step = pred_horizon[:, :, step % args.replan_interval]
        scenes = batch_update_sample(scenes, pred_step, ego_traj, step)  # [tools] state helper or model-specific
        rollout.append(pred_step)
        visualize_step(...)  # [tools] NuPlanVisualizer

    scene_traj = stack_rollout(rollout)
    process_scene_output(...)  # saves rollout and evaluates one scenario

render_videos(...)  # [tools]
save_metric_outputs(aggregator, metrics_dir)  # [tools]

Metrics

The main metric entry is:

reactsim_data_pipeline/closed_loop/metrics.py

Use build_metric_evaluator once before simulation, process_scene_output once per scenario, and save_metric_outputs once after all scenarios:

from reactsim_data_pipeline.closed_loop import metrics as shared_closed_loop_metrics

evaluator, aggregator = shared_closed_loop_metrics.build_metric_evaluator(
    scenario_dir=scenario_dir,
    map_root=map_root,
    device=device,
)

shared_closed_loop_metrics.process_scene_output(
    scene_name=scene_name,
    scene_traj=scene_traj,
    scene_orig=scene_orig,
    evaluator=evaluator,
    aggregator=aggregator,
    metrics_dir=metrics_dir,
    rollout_save_path=rollout_save_path,
    save_rollout=rollout_save_path is not None,
)

shared_closed_loop_metrics.save_metric_outputs(aggregator, metrics_dir)

process_scene_output expects:

Input	Shape / fields	Meaning
scene_name	str	Scenario file stem, used to infer the map name and save outputs.
scene_traj	(80, N, 5)	Simulated rollout in global coordinates. State order is x, y, heading, vx, vy.
scene_orig["agents_history"]	(N, 11, 8)	History states. Columns include x, y, heading, vx, vy, length, width, height.
scene_orig["agents_future"]	(N, 81, 5)	Logged future states from the current frame onward.
scene_orig["agents_interested"]	(N,)	Nonzero agents are evaluated.
scene_orig["agents_id"]	(N,)	Agent ids used to align rollout and logged agents.
scene_orig["agents_type"]	(N,)	Agent type ids.

Visualization

The main visualization entries are:

reactsim_data_pipeline/visualize_samples.py
reactsim_data_pipeline/visualization/nuplan_visualizer.py

Use visualize_samples.py for checking preprocessed cache files:

cd reactsim_data_pipeline
python visualize_samples.py \
  --input ../data/nuplan_processed/val/processed_scenarios/<scenario>.pkl \
  --output_dir outputs/cache_vis

Use NuPlanVisualizer inside closed-loop scripts for rollout visualization:

from reactsim_data_pipeline.visualization.nuplan_visualizer import NuPlanVisualizer

visualizer = NuPlanVisualizer(vis_dir)
visualizer.visualize_batch_step(batch_data, pred_trajs, step, batch_names)
visualizer.visualize_rollout_frame(scene_orig, scene_traj, scene_name, step)
visualizer.create_gif_from_trajectory(scene_orig, scene_traj, scene_name, fps=video_fps)

The closed-loop scripts control visualization with:

--visualize_mode  none, image, video, or both
--vis_freq        save one image every N simulation steps
--video_fps       rendered video frame rate
--video_workers   number of video rendering worker processes

Core visualization inputs:

Input	Shape / fields	Meaning
sample for visualize_samples.py	ReactSim cache dict	Raw preprocessed scenario cache.
batch_data["agents_history"]	(B, N, 11, 8)	Current batched scene history.
batch_data["agents_future"]	(B, N, 81, 5)	Logged future states, used as reference if available.
batch_data["agents_type"]	(B, N)	Agent type ids for drawing boxes.
batch_data["polylines"]	(B, M, P, C)	Map polylines in shared closed-loop format.
batch_data["polylines_valid"]	(B, M)	Valid map polyline mask.
batch_data["traffic_light_points"]	(B, L, 3)	Traffic-light stop points and current state.
pred_trajs	(B, N, H, 5)	Predicted future trajectories in global coordinates.
scene_traj	(80, N, 5)	Saved or in-memory closed-loop rollout.

If rollout_load_path is set, the model is not run; saved rollouts can still be evaluated and visualized.