Waymax: An Accelerated, Data-Driven Simulator for Large-Scale Autonomous Driving Research
Authors: Cole Gulino
Justin Fu
Wenjie Luo
George Tucker
Eli Bronstein
Yiren Lu
Jean Harb
Xinlei Pan
Yan Wang
Xiangyu Chen
John D Co-Reyes
Rishabh Agarwal
Rebecca Roelofs
Yao Lu
Nico Montali
Paul Mougin
Zoey Zeyu Yang
Brandyn White
Aleksandra Faust
Rowan Thomas McAllister
Dragomir Anguelov
Benjamin Sapp
Venue: NeurIPS 2023 Datasets and Benchmarks Poster (2023)
Keywords: simulator
autonomous driving
planning
behavior prediction
reinforcement learning
imitation learning
TL;DR: An autonomous driving simulator for behavior problems with hardware acceleration and intelligent out-of-the-box agents.

Abstract: Simulation is an essential tool to develop and benchmark autonomous vehicle planning software in a safe and cost-effective manner. However, realistic simulation requires accurate modeling of multi-agent interactive behaviors to be trustworthy, behaviors which can be highly nuanced and complex. To address these challenges, we introduce Waymax, a new data-driven simulator for autonomous driving in multi-agent scenes, designed for large-scale simulation and testing.  Waymax uses publicly-released, real-world driving data (e.g., the Waymo Open Motion Dataset) to initialize or play back a diverse set of multi-agent simulated scenarios.   It runs entirely on hardware accelerators such as TPUs/GPUs and supports in-graph simulation for training, making it suitable for modern large-scale, distributed machine learning workflows. To support online training and evaluation, Waymax includes several learned and hard-coded behavior models that allow for realistic interaction within simulation. To supplement Waymax, we benchmark a suite of popular imitation and reinforcement learning algorithms with ablation studies on different design decisions, where we highlight the effectiveness of routes as guidance for planning agents and the ability of RL to overfit against simulated agents.

Waymax: An Accelerated, Data-Driven Simulator for Large-Scale Autonomous Driving Research
Cole Gulino
Justin Fu +f
Wenjie Luo *f
George Tucker *1
Eli Bronstein
Yiren Lu t
Jean Harb
Xinlei Pan 1
Yan Wang f
Xiangyu Chen
John Co-Reyes t
Rishabh Agarwal
t
Rebecca Roelofs :
Yao Lu
Nico Montali
Paul Mougin
Zoey Yang
Brandyn White f
Aleksandra Faust
Rowan McAllister
Dragomir Anguelov
Benjamin Sapp
Equal Contribution
Waymo Research
Google DeepMind
Abstract
Simulation is essential tool to develop and benehmark autonomous vehicle planning software a safe and cost-effective manner. However, realistic simula- tion requires accurate modeling of nuanced and complex multi-agent interactive behaviors. To address these challenges, we introduce Waymax, a new data-driven simulator for autonomous driving in multi-agent scenes, designed for large-scale simulation and testing. Waymax uses publicly-released real-world driving data (e.g., the Waymo Open Motion Dataset [15)) to initialize play back diverse set of multi-agent simulated scenarios. It runs entirely on hardware accelerators such TPUs/GPUs and supports in-graph simulation for training, making it suitable for modern large-scale, distributed machine learning workflows. To support online training and evaluation, Waymax includes several learned and hard-coded behav- ior models that allow for realistic interaction within simulation. To supplement Waymax, we benchmark suite of popular imitation and reinforcement learning algorithms with ablation studies on different design decisions, where we highlight the effectiveness of routes as guidance for planning agents and ability of RL. overfit against simulated agents.
Introduction
Due to the cost and risk of deploying autonomous vehicles (AVs) in the real world, simulation is a crucial tool in the research and development autonomous driving software. The two primary challenges of simulator
speed and realism: wish for simulator to be fast in order to cost-effectively train/evaluate on many hours of synthetic driving experience, and we wish for a simulator to be diverse and realistic in terms vehicle behavior in order to minimize the sim-to-real gap [50, 37], such that performance in the simulator correlates with real-world performance.
Existing work in simulation for autonomous driving has made significant progress in recent years. Simulators such as CARLA [14], Sim4CV [33] and SUMMIT [9] focus on photo-realistic rendering of driving scenarios, enabling users to train and evaluate driving solutions. However, major
37th Conference on Neural Information Processing Systems (NeurIPS 2023) Track on Datasets and Benchmarks.

---

(a) Waiting for turn into oncoming traffic.
(b) Navigating a 4-way intersection.
Figure I: Two examples demonstrating the types of interactive, urban driving scenarios available in Waymax. (a) shows a vehicle waiting for oncoming traffic to pass before turning into a narrow street. (b) shows an agent performing left turn at 4-way intersection while following a route (boundaries highlighted in green).
simulation challenge still remains in the generation of diverse scenarios and realistic behavior for other agents (such vehicles and pedestrians) in the scene, and the driving field has matured, behavior challenges have been shown to be a significant bottleneck to scaling [31]. To this end, there is still need for simulation tools that provide (a) realistic, closed-loop simulation of agent behavior, and (b) high speed and throughput to support modern trends in machine learning that use large models and datasets.
To address these challenges, we propose Waymax- a differentiable, hardware-accelenn agen! simulator that is built using real-world driving data from the Waymo Open Dataset. Waymax aims to provide, within simulation, a faithful reproduction of the data and types of challenges a real autonomous driving agent would face, such those shown in Fig. 1. Waymax simulates challenging obstacles present in urban driving, such as pedestrians and cyclists, and provides high-level route information for the ego vehicle follow. To optimize runtime speed and facilitate rapid development, Waymax is written using JAX [5], which allows simulation to be run entirely on accelerators such as graphics and tensor processing units (GPUs and TPUs). To provide better simulation realism, Waymax uses diverse scenarios initialized from the Waymo Open Motion Dataset (WOMD) [15], which contains over 250 hours of real driving data collected in dense urban environments. Waymax data loading and processing can be extended to other popular datasets without loss of generality.
and multi-
Our contributions two-fold. First, we introduce the Waymax simulator, which is multi-agent simulator for autonomous driving that is (a) hardware-accellertted (b) provides features and routes information from real driving data, and (c) constructs scenarios upon a large and diverse dataset of real-world driving. Our second contribution is in providing set of common benchmarks and simulated agents that allow researchers score and benchmark their autonomous planning methods in closed-loop. We show-case the flexibility of Waymax by training behavior algorithms for an autonomous vehicle in different setups (imitation, on and off policy RL, etc) against range of different interactive agents.
2
Related Work
Simulators for Autonomous Driving. Waymax is a differentiable, hardwwrre-uceelertte and multi-agent simulator that is built on top of real-world driving data. We compare our work to other related publicty available simulators in Table 1. The closest works ours multi-agent autonomous driving simulators which use real driving data to initialize scenarios and logged behavior, such Nocturne [52], MetaDrive [25], and nuPlan [8]. In comparison, our simulator is designed to support hardware-acceleree acc.-ereeeereed training, the agent models can be connected in-graph for training and inference, and the simulation is differentiable. We provide full set of features including pedestrians, cyclists, and traffic lights available in WOMD, and the inferred routes for coal-conditioned policy and progress metrics. Additionally reactive sim agent models are provided facilitate realistic simulation. TorchDriveSim [46] is the only public simulator that supports differentiable simulation for in graph acceleration. In comparison, we provide rich and diverse real-world human expert driving data from WOMD. It is worth noting that the driving policy modeling problem is primarily focused on behavior than perception, thereby, we do not intend to support sensor simulation, such as [14, 28, 48, 9, 2]. which represents another important research area for autonomous driving perception.

---

(a) The route the union of logged trajecto- ries and driveable futures.
(b) Reactive simulated agents stopping to avoid collision.
The routes given to an agent (all areas highlighted in color) are computed by combining the logged future trajectory of the agent with all possible future routes after the logged trajectory. b): Waymax is bundled with reactive simulated agents. Here, agent #5 (circled in red) is stopped in front of an intersection, causing the IDM- controlled agents (#1, 2, and to brake in order to avoid collision.
Figure 2: A sample of features available in Waymax. a):
Multi-agent Accel.
Sensor Sim Expert Data Sim-agents Real Data
Routes/Goals
TORCS [55]
GTA V [29]
CARLA [14]
Waypoints
Highway-eny [24]
Sim4CV [33]
Directions
SUMMIT [9]
V 400)
MACAD [38]
Goal Goint
DeepDrive-Zero [40]
SMARTS [57]
Waypoints Goal point
MADRaS [44]
/C 10)
DriverGym [23]
VISTA [2]
nuPlan [8]
Waypoints Goal point
Nocturne [52]
V 50)
MetaDrive [25]
Intersim [47]
Goal point
TorchDriveSim [46]
tbsim [56]
Goal point Waypoints
Waymax (ours)
/C> 128)
Table A comparison of related driving simulators (chronological order). Multi-agent: Simulating multiple agents. Supported number of agents are in the parentheses (if available). Accel:: In orgrraph compilation for hardware (GPU/TPU) acceleration. Sensor Sim: Sensors (e.g. camera, lidar & radar) input simulation. Expert Data: Human demonstrations rollout trajectories collected with an expert policy. Sim-agents: agent models for simulated objects (e.g. other vehicles). Real data: Real world driving data. Routes/Goals: "_" means routes goals are provided; "Waypoints" means positions sampled from trajectory: "Directions" means discrete driving directions including left, straight, and right; "Goal point" means the goal position.
Learning-based Driving Agents. In the context of autonomous driving. open-loop imitation learn- ing (IL), known as behavior cloning (BC), has been applied to predict driving behaviors of other road users [3, 10, 49, 26, 36, 17, 34] well the ego vehicle [39, 4, 41, 12, 53, 11]. is widely known that BC suffers from covariate shift [42] and causal confusion [13]. Closed-loop methods including
adversarial imitation learning [20, 7] and reinforcement learning (RL) methods [22, 21, 54, 27] have been proposed to address these challenges by learning from feedback and explicit hand designed rewards the simulator. Despite the excitement of machine learning as an avenue towards devising autonomous driving policies, benchmarking different learned policies large scale real world datasets remains challenge. To enable accelerated and effective autonomous driving agents research, Waymax enables standard training and evaluation workflows and reliable benchmarking in both open

---

and closed-loo settings; we also provide the implementation of a representative set of IL and RL baselines and report their performance against standard set of metrics on Waymax as references.
3 Simulator Features
In this section, we give an overview of the features of Waymax and its interface to the user. Waymax is simulator that supports controlling arbitrary number of objects in scene. A primary goal of Waymax is to initialize from real- weal-world driving scenarios to model complex interactions between vehicles, pedestrians, and traffic lights, and following goal route provided by high-level planner. Additionally, Waymax is designed to be both fast and flexible each component discussed this section can easily be modified or replaced by an user to suit their own project needs. We discuss the scenarios and datasets in Sec. 3.1, state representation in Sec. 3.2, and the dynamics and action representation in Sec. 3.3. Waymax includes suite of common metrics described in Sec. 3.4, and several options for modeling the behavior of dynamic objects (vehicles and pedestrians) in the scene, outlined in Sec. 3.5.
3.1 Scenarios and Datasets
In contrast to simulators that generate synthetic scenarios (e.g. CARLA [14]), Waymax utilizes real-world driving logs to instantiate driving scenarios, and runs for a fixed number of steps. We provide default support for the Waymo Open Motion Dataset (WOMD) [15]. which includes over 100, 000 trajectories snippets and 7.64 million unique objects to interact or control. Each trajectory snippet is 9 seconds recorded with 0.1 Hz. The trajectory contains pose and velocity information for all objects in a scene, including the autonomous vehicle (AV), other vehicles, pedestrians, and cyclists. For each scenario, we take the static information such as the road graph and initialize dynamic objects using the first second of logged information. Then, agent models (described in Sec. 3.5) will be used to control the dynamic objects such as pedestrians and the other vehicles through the simulation steps. Note importantly, users can inject multiple agent models and dynamics model to Waymax environment where each model can control multiple objects.
3.2 State and Observation spaces
The first component of defining autonomous driving as a sequential control problem is defining the state space. We include two types of data in the state: dynamic data which can change over the course of an episode and across scenarios, and static data which remains the same during an episode but varies across scenarios. The dynamic data in the state consists of the position, rotation, velocity, and bounding box dimensions for all vehicles, cyclists, and pedestrians in a scene, along with the color of traffic light signals (red, yellow, green). The static data includes the road and lane boundaries sampled as a 3D point-cloud (known as the "roadgraph"). as well as on-route and off-route paths for the ego vehicle. Each agent views the simulator state through user-defined observation function, which can induce partial observability. We provide a default observation function that transforms the location of all other vehicles the agent's 's own coordinate frame, and sub-samples the roadgraph via distance.
We augment each scenario with feasible paths that the AV could take from its initial position. A path is represented as sequence of points, which are subset of the roadgraph points. Each path is computed by performing a depth- depth-first-seede first-search traversal of the roadgraph from the starting position. Together, these paths describe all the ways in which the AV can legally drive in the scenario. Similar to the "road-route" in [7]. a path is considered on-route if it follows the same road the AV's logged trajectory. The remainder of the paths that are not on-route are deemed to be oofteo Fig. 2a gives an example of on-route paths. These paths are useful for computing metries as well as developing coal-conditioned planning and interactive agents.
On-Route and Off-Route Paths
3.3 Object Dynamics
The object dynamics defines what 'actions' an object would expect and how its state would evolve given an action. Waymax allows the user to define a dynamics model and provides several pre-defined options for controlling the physical dynamics of vehicles in simulation: 1) the delta action space.

---

which is suitable for all types of objects, uses position difference (the delta term Az, Ay. AB)) between two consecutive states; and the bicycle action space ((a, k),, which is only for vehicles, uses acceleration and steering curvature). The equations defining these dynamics can be found in Appendix A.1.
3.4 Metrics
Waymax provides a set of intuitive metrics evaluate the ego vehicle as well as simulated agents for safety and correctness of behavior (such as obeying traffic rules, not colliding), as well as comfort and progress, All metrics in Waymax are computed closed-loop, meaning that they are computed by running the agent in simulation, rather than in open-loop, where metrics are computed on per-timestep basis without feedback from simulation. The metrics available are as follows:
The route progress ratio measures how far the ego vehicle drives along the goal route compared to the logged trajectory. At time step this metric associates the vehicle's position to the closest point x() in an on-route path. It then computes the distance along the path from the start of the path to (t)., denoted as dz(t). The route progress ratio is then defined as 4z0 d,-dp- where dp and dy are the distances along the path to the initial and final positions of the vehicle's logged trajectory, respectively. Since the vehicle can continue driving after reaching its destination, this ratio could be greater than
Route Progress Ratio
Off-Route
The off-route metric binary value indicating if the vehicle is following an on-route path. If the vehicle is sufficiently closer to an off-route path than on-route path or it is far enough away from an on-route path, is considered off-route.
Off-Road The off-road metric triggers if vehicle drives off the road. This is measured relative the oriented roadgraph points. If a vehicle is on the left side of an oriented road edge, it is considered on the road: otherwise the vehicle is considered off-road.
The collision metric is a binary metric that measures if the vehicle is in collision with another object in the scene. For each pair of objects, if the 2D top-down view of their bounding boxes overlap in the same timestep, they are considered in collision.
Collision
Kinematic Infeasibility Metric The kinematic infeasibility metric computes binary value of whether a transition is kinematically feasible for the vehicle. Given two consecutive states, we first estimate the acceleration and steering curvature using the inverse kinematics defined in Appendix A.I, and check if the values are out of bounds. We empirically set the limit of acceleration magnitude to be m// s2 and the steering curvature magnitude to be 0.3 m In order to determine these empirically, we fit the logged trajectories of the ego agent with our steering and acceleration action space. We then chose the limits be roughly the maximum (rounding up for some slack) of the values we observed in the logs.
Displacement Error The average displacement error metric (ADE) measures how far the simula- tion deviates from logged behavior. It is defined as the L2 distance between each object's current XY position and the corresponding position recorded in the logs at the current timestep, averaged across all timesteps.
3.5 Simulated Agent Behavior
An important part of constructing a simulator for autonomous driving is realistic behavior for simulated agents other than the AV. Waymax, as a multi-agent simulator, gives the user the ability to control the behavior of all objects in simulation. This allows the user to control agents with any model of choice, such as learned behavior models. However, to support training AV agents out-of-the-box, Waymax also includes a rule-based reactive agent model based on the intelligent driver model (IDM) [51]. IDM describes rule for updating the acceleration of a vehicle to avoid collisions based on the proximity and relative velocity of the vehiele to the object directly in front of the vehicle, demonstrated in Fig. 2b. The IDM agent in Waymax follows the logged path that is
5

---

recorded in the data, but uses IDM to adjust the speed profile to avoid collisions and accelerate on free roads.
4 Software API
We now outline the Waymax software components and interfaces. In order to support a wide variety of research workflows, Waymax designed as collection of inter-operable libraries while maintaining fast simulation speed. The main libraries comprise of (1) set of common distributed data-loading library, (3) simulator components such as metrics and dynamics, and (4) Gym-like environment interface. Each component of the simulator can be modified, replaced, or used standalone by the user. this manner, users who only need one component of Waymax (e.g. only metrics, or only data loading), who wish to significantly modify the behavior of the simulator (such as generating synthetic scenerios) can easily so through Waymax's 's APIs.
(2)
4.1 Environment Interface
Users primarily interact with Waymax as purtially-observable stochastic game. The Waymax interface follows the the Brax [16] design to only define functionally pure initialization and transition functions. This stateless design enables efficient optimization through JAX's [5] JIT compiler and functional libraries, and easily allows users implement control algorithms that require backtracking, such search. In contrast with stateful simulators, such as OpenAl Gym [6] and DM Control [32]. Waymax users need to maintain the simulator state within a simulation loop and interact with the simulator primarily through two functions:
The reset (scenario) function takes as input a raw scenario, performs any initialization necessary such populating the simulation history, and returns the initial state object.
The step (state, action) function takes as input the current state, the actions for all agents, and computes the successor state as well as the new observation and metrics. The actions argument is data structure that contains a data tensor of actions for each agent, as well as a validity mask which denotes which agents the user wishes to control. step then returns these results in
timestep object.
Using these two functions, a user can run simple, but complete simulation of stochas- tic game between multiple agents, such as in the following pseudocode example:
Run one episode until termination. state env. (next (dataset))
while not done:
In addition, we do provide adapters to convert the functionally pure Waymax simulator into stateful one to support existing codebases.
action policy (env.observeeeeeeeoeeeeeeeeeeeee
observe observe(state))
state env sey.stee(steee action)
4.2 Hardware Acceleration and In-graph training.
Waymax supports both hardware acceleration on GPUs and TPUs, as well as combining training and simulation within the same computation graph (referred to as "in-graph"" training), which allows training and simulation to happen entirely the accelerator without communication bottlenecks through the host machine. These features are possible because Waymax is written entirely using the JAX [5] library, which converts operations into XLA [43], linear algebra instruction set and optimizing compiler which supports execution on CPU, GPU, or TPU. In-graph training requires the modeling and training code to be written using XLA-compatible frontend such as JAX [5], or Tensorflow [1]. The XLA compiler can then optimize combined training and simulation program to produce single computation graph that can be entirely hardware accelerators, without communication costs between the accelerator and host device.
4.3 Single and Multi-agent Simulation
While the base multi-agent environment allows us to do sim-agents anulti-agent) learning similar as in Nocturne [52], MetaDrive [25], the ultimate goal of the autonomous driving problem train an AV planning agent. Thus, Waymax supports both multi-agent simulation that allows users to control

---

Device
BS-1 BS-16
Reset
Step
Transition
Metries
RolloutExpert
CPU
1.09
131
0.90
112
1.0x10*
Single-
CPU
12.2
1.7x10a
10.9
1.69x10'
1.4x10$
Agent
GPU-v100
/
0.58
0.75
0.47
0.21
56.2
Env
GPU-v100
V
0.67
2,48
0.52
2.27
279
CPU
6.23
129
1.01
112
x.1
Multi-
CPU
49.8
1.1x103
14.3
1.72x103
1.6x
Agent
GPU-v100
/
0.64
0.92
0.53
0.19
73.3
Env
GPU-v100
0.81
2.86
0.51
2.24
OOM
Table 2: Runtime benchmark in milliseconds: the environment controls all objects in the scene (up to 128 as defined in WOD).
arbitrary objects within the scenario, as well as single-agent workflow where single AV agent is trained using learned or rule-based models to control the other vehicles in the scene.
While it might be possible to put multiple poli-
Mill agent Interbace
cies one environment directly, it is certainly not flexible way as it is hard to coordinate different policies or change policies. Waymax provides two interfaces for different use-cases:
The MultiAgentEEnvinen
provides an in- terface for multi-agent and sim-agent problems. The user provides simultaneous actions for all controlled objects in the scene, as well as a mask indicate which objects should be controlled.
Single-apnt Interface
Figure 3: An illustration of simulation rollout using reactive simulated agents to control non-AV agents, and a user-defined policy control the AV.
The PlanningAgentEnviionn
exposes an interface for controlling only the ego vehicle in the scene. All other agents are controlled by user-specified sim agents or log playback (Fig. 3).
5 Experiments
We now evaluate both Waymax as simulator and the performance of several reference agents simulated using Waymax. We first evaluate the computational performance of Waymax in Sec. 5.1 under various configurations. Second, we perform an empirical study of several benchmark agents for planning Sec. 5.3, where we compare the performance of several broad categories of learned planning algorithms (such as imitation learning and reinforcement learning) against both logged agents and reactive simulated agents. For the second part, our goals was showcase potential options for using Waymax, so we opted for simple design choices and a breadth of configurations, and we expect that the performance the baseline agents could be significantly improved in future work.
5.1 Runtime Benchmark
In Tab. 2, we present the runtime performance of Waymax using a CPU (Intel Xeon W-2135@33 3.7GHz) and GPU Evidia-V100), We evaluate the performance of both multi-agent and the single-agent environment with different batch size. All functions are jit compiled and runtime is reported in millisecond. Following WOMD, the environment controls up 128 objects in one scene. Note that the Step function computes both the state transition and the reward. While users specify customized reward function, for this runtime evaluation, we use the negative sum of all metrics in 3.4 as the reward, which measures the effect of computing all metrics. When considering batch size 1 and using a GPU, Waymax achieves over 1000Hz for Step function and over 2000Hz only considering the Transition on. More importantly, as Waymax supports batching, Step only takes 2.86ms using a batch size of 16. Note this much faster than running batch size one for 16 times and gives an equivalent runtime of over 5000Hz per example (i.e. closer to 500 times faster than using CPU). Noticeably the Metrics function consumes more computation then the Transition function because the Off Road metric needs to find nearby roadgraph points, which is slow operation.
7

---

Agent
Action
Train Sim
Off-Road Rate (%)
Collision Rate (%)
Kinematic Infeasibility
Log (m)
ADE Route
Space
Agent
Progress
(%)
Ratio (%)
Expert
Delta
0.32
0.61
4.33
0.00
100.00
Expert
Bicycle
0.34
0.62
0.00
0.04
100.00
Expert
Bicycle
0,41
0.67
0.00
0.09
100.00
(Discrete)
Wayformer
Delta
7.89
10.68
5.40
2.38
123.58
BC
Delta
4.14+2.04
5.83=1.099
0.18+
6.28=1.93
79.58=24.98
BC
Delta (Dis-
0.4
5.97=0.10
66.25=0.22
2.98
98.82+33.46
crete)
BC
Bicycle
12.71
11.20+-5.34
0.00++0.00
3.60+11111
1.3.11 1:33378
BC
Bicycle
1.11+0020
4.59
0.00+00.00
2.26+-0.02
129.84+0098
(Discrete)
DQN
Bicycle
IDM
3.74
6.501:031
0.00:10.00
9.83:10488
177.91
(Discrete)
DQN
Bicycle
Playback
4.31+1.09
4.91++0.70
0.00+00.00
10.74+0.53
215.26++38.00
(Discrete)
Table 3: Baseline agent performance evaluated against IDM sim agents with route conditioning. Models trained ourselves (BC and DQN) report mean and standard deviation over 3 seeds. Off-Road, Collision, and Kinematic Infeasibility are reported as percentage of episodes where the metric flagged at any timestep. Action spaces are continuous unless noted otherwise. By construction the bicycle action space does not violate the comfort metric.
Rollout We also benchmark a Rollout function which rolls out the environment given an Actor for entire episode (i.e., 80 steps for WOD). This is especially useful to provide faster inference and evaluation. In the last column of Tab. 2, we show the runtime of Rollout with an Expert ExpertActon that derives grouth-truth actions from logged trajectory. It is faster than rumning Step function 80 times. More importantly, we can see that running on GPU has a consistent 2 orders of magnitude speedup. As a point of reference, evaluating the full WOD evaluation dataset (44K scenarios) with 8-V100 machine takes less than 2min.
5.2 Baseline Planning Agents
Expert We provide number of expert agent models to provide groundtruth actions for open-loop training. Each agent uses the inverse function of the action spaces defined in Section 3.3 fit an action to the logged trajectory. For discrete action spaces, the inverse is computed by discretizing the continuous inverse.
Behavior Prediction Model (Wayformer) As point of reference, we adapt the state-of- trate-oftthe-wee the-arrt Wayformer behavior prediction model [35] to the planning setting. Originally, the Wayformer predicts multiple 8-second future trajectories given 1-second context history. To adapt it to the planning setting, we autoregressively feed its predictions as the context history and choose the most likely trajectory. We found that making predictions at lower frequency than the environment frequency improved performance, so we predict 5-step long trajectories and only replan every 5 steps.
We re-use the encoder portion of Wayformer [35] followed by a 4-layer residual MLP maximize the log likelihood of the expert actions. For continuous actions, we used 6- component Gaussian Mixture Model. For discrete actions, we used softmax layer to compute action probabilities.
Behavior Cloning
Model-Free Reinforcement Learning DQN tized replay double DQN [45].
We used the Acme [19] implementation of priori-
We used the same architecture as in discrete BC for the Q-network, interpreting the logits of the model Q-values.
For simplicity, we use a sparse reward penalizing collisions and off-road events: r, -Teollision (1) Loft-road (t).

---

5.3 Planning Benchmark Results
To showcase the flexibility of our environment, trained number of baselines different action spaces and algorithms as shown in Table 3 and evaluated the metrics defined in Section 3.4. We evaluated each agent against the IDM sim agent and conditioned it on the route by adding the points from all the on-route paths as an additional input group to the Wayformer [35] encoder. These points represent the on-route subset of the roadgraph points. All agents are trained for the planning agent task and thus only provide predictions for the autonomous vehicle. See Appendix A.2 for training details.
As expected, the expert agents have low oof-road and collision rates. The nominal values represent noise in the bounding boxes and logged data and serve lower bound for performance. The expert using the discrete bicycle action space has comparable performance to the other experts, confirming that the discretization is sufficiently fine.
For open-loop imitation, the discrete action space performs best, possibly because it easier model multi-modal behavior. Furthermore, it outperforms the adapted Wayformer model, likely due to the fact that trained explicitly for this task. This serves as a check that the Waymax environment producing the correct training data.
Route Conditioning Ablation To showcase the utility of route conditioning, we compare the performance of route conditioned versus non-route conditioned behavior cloning agents Table 4 shows that the route conditioned agent is substantially better at following the route, while also achieving a lower off-road rate, collision rate, and log ADE. These results indicate that the route provides a strong signal for the planning task.
Agent (Action Space)
Off-Road Rate (%)
Collision Rate (%)
Off-Route
Log ADE (m)
Route
Rate (%)
Progress
Ratio (%)
Expert (Bicyele Discrete)
0.41
0.67
0.00
0.00
100.00
BC (Bicycle Discrete)
1.45=0005
4.92:10.24 2.31 12 2.411:0.06
128.73+
BC (Bicycle Discrete) + Route
1.11:1..0
4.59
0.96 2.26-10.02 1..98410.98
Table 4: Experimental ablation comparing performance with and without route conditioning.
Sim Agent Ablation In Table 5, we show the effect of training and evaluating an imitation agent against IDM sim agents versus playing back logged trajectories. As expected, evaluating with the IDM agent produces fewer collisions than evaluating with log playback. However, training an RL agent with IDM agents was less effective than training against logged agents. We believe this is because the RI. agent tends to overfit or exploit the behavior of 'easier' IDM agents. Since IDM will stop for the SDC avoid collisions, the RI agent does not have as much incentive to learn how avoid collisions itself. We can see that when an IDM-trained agent is evaluated against logged agents, the collision rate is over 4x higher than when evaluated against IDM agents.
Train
Eval
Collision
Offroad
Progress Ra-
Agent
Agent
Rate (%)
Rate (%)
tio (%)
Playback
Playback
8.67:10.97
5.16:1.00
193.56
Playback
IDM
4.91-10.0
4.31+1.09
318.26-38.20
IDM
Playback
25.15=0..76
4.91+0.41
163.11
IDM
IDM
6.50+00.1
3.74=0.90
177.91
Table 5: Experimental ablation over different configurations of train/evaluation sim agents with
the DQN algorithm using the Bicycle dynamics model. Using interactive sim agents such as IDM can reduce the rate of unrealistic collisions during evaluation. However, it more difficult train effective planning agents using interactive sim agents.

---

Conclusion
We have presented Waymax, multi-agent simulator for autonomous driving. Waymax provides diverse scenarios drawn from real driving data, and supports hardware acceleration and distributed training for efficient and cost-effective training of machine-leamed models. also designed with flexibility in mind Waymax is written as a collection of inter-operable libraries for data loading, metric computation, and simulation, which can support a wide variety research problems that are not limited to just the planning evaluations presented in this work. We conclude by benchmarking several common approaches to planning with ablation studies over different dynamics and action representations, which provide set of strong baselines for benchmarking future work.
In addition hardware acceleration, Waymax also enables the exploration of methods utilizing differentiable simulation, as the entire simulation can be assembled within a single JAX computation graph. Prior work[30, 20] has shown that differentiable simulation can improve the efficiency of policy optimization methods as they can rely on a "reparameterized" or tassuhh gradient to reduce the variance of the gradient estimate. We believe that this is promising line of future work to be explored.
As mentioned previously, the problem of sim-to-real transfer is critical issue in autonomous driving, as it is cheap and desirable to evaluate simulation but difficult to guarantee that the same performance and level of safety will carry over the real world. While Waymax we have made design decisions to minimize this gap (such as using real-world data to seed scenarios), this remains an important limitation for any simulation-basel framework. fruitful line of future work is to close the gap between simulated and real-world performance, potentially using techniques such as domain randomization [50] or combining real and synthetic data[37, 18].
10

---

References
[1] Martín Abadi. TensorFlow: learning functions at scale. In Proceedings of the 21st ACM SIGPLAN
International Conference on Functional Programming, pages 1-1, 2016. 6
[2] Alexander Amini, Tsun-Hsuan Wang, Igor Gilitschenski, Wilko Schwarting, Zhijian Liu, Song Han, Sertac Karaman, and Daniela Rus. Vista 2.0: An open, data-driven simulator for multimodal sensing and policy learning for autonomous vehicles. In 2022 International Conference on Robotics and Automation (ICRA),
pages 2419 2426. IEEE, 2022. 3
[3] Mayank Bansal, Alex Krizhevsky, and Abhijit Ogale. ChauffeurNet: Learning to drive by imitating the best and synthesizing the worst. In Robotics: Science and Systems (RSS), 2019. 3
[4] Mariusz Bojanski, Davide Del Testa, Daniel Dworakowski, Bernhard Fimer,, Beat Flepp, Prasoon Goyal, Lawrence D Jackel, Mathew Monfort, Urs Muller, Jiakai Zhang, et al. End to end learning for self-driving cars. arXiv preprint 1004.0:664.0733
2016. 3
[5] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclau- rin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Millen and Qiao Zhang. JAX: composable transformations of Python+NumPPy programs, 2018. 6,
[6] Greg Brockman, Vicki Cheung. Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojeiech Zaremba. Openai gym, 2016.
[7] Eli Bronstein, Mark Palatucci, Dominik Notz, Brandyn White, Alex Kuefler, Yiren Lu, Supratik Paul, Payam Nikdel, Paul Mougin, Hongge Chen, Justin Fu, Austin Abrams, Punit Shah, Evan Racah, Benjamin Frenkel, Shimon Whiteson, and Dragomir Anguelov. Hierarchical model-based imitation learning for planning in autonomous driving. In 2022 IEFF/RSI international conference on intelligent rohots and systems (IROS), pages 8652-8659. IEEE, 2022. 3,4
[8] Holger Caesar, Juraj Kabzan, Kok Seang Tan, Whye Kit Fong. Eric Wolff, Alex Lang, Luke Fletcher, Oscar Beijbom, and Sammy Omari. nuplan: A closed-loop ml-based planning benchmark for autonomous vehicles. arXiv preprint arXiv:2106. 11810, 2021. 2,
[9] Panpan Cai, Yiyuan Lee. Yuanfu Luo, and David Hsu. Summit: A simulator for urban driving in massive mixed traffic. In 2020 IEEE International Conference on Robotics and Automation (ICRA), pages 4023- 4029. IEEE, 2020. 2.
[10] Yuning Chai, Benjamin Sapp, Mayank Bansal, and Dragomir Anguclov. Multipath: Multiple probabilistic anchor trajectory hypotheses for behavior prediction. arXiv preprint arXivvI910.0544 2019.
[11] Dian Chen, Brady Zhou, Vladlen Koltun, and Philipp Krähenbühh. Learning by cheating. In Conference
Robot Learning, pages 66-75. PMLR, 2020. 3
[12] Felipe Codevilla, Matthias Müller, Antonio López, Vladlen Koltun, and Alexey Dosovitskiy. End-to- end driving via conditional imitation learning. In 2018 IEEE international conference robotics and
andomation (ICRA), pages 4693-4700. IEEE, 2018. 3
[13] Pim De Haan, Dinesh Jayaraman, and Sergey Levine. Causal confusion in imitation learning, Advances in Neural Information Processing Systems, 32, 2019. 3
[14] Alexey Dosovitskiy, German Ros, Felipe Codevilla, Antonio Lopez, and Vladlen Koltun. Carla: An open urban driving simulator. In Conference on Robot Learning, pages 1-16. PMLR, 2017. 2,
[15] Scott Ettinger. Shuyang Cheng, Benjamin Caine, Chenxi Liu. Hang Zhao, Sabeek Pradhan, Yuning Chai, Ben Supp, Charles Qi, Yin Zhou, Zoey Yang, Aurelien Chouard, Pei Sun, Jiquan Ngiam, Vijay Vasudevan, Alexander McCauley, Jonathon Shlens, and Dragomir Anguelov. Large scale interactive motion forecasting for autonomous driving: : The waymo open motion datasct. arXiv, 2021. 1, 2, 4
[16] C. Daniel Freeman, Frik Frey, Anton Raichuk, Sertan Girgin, Igor Mordatch, and Olivier Bachem. Brax
differentiable physics engine for large scale rigid hody simulation, 2021. 6
[17] Junru Gu, Chen Sun, and Hang Zhao. Densetnt: End-to-end trajectory prediction from dense goal In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 15303-15312, 2021.
[18] Alexander Herzog. Kanishka Rao, Karol Hausman, Yao Lu, Paul Wohlhart, Mengyuan Yan, Jessica Lin, Montserrat Gonzalez Arenas, Ted Xiao, Daniel Kappler, al. Deep at scale: Sorting waste in office buildings with a flect of mobile manipulators. arXiv preprint arXiv:2305.03277 032700. 2023. 10
[19] Matthew W. Hoffman, Bobak Shahriari, John Aslanides, Gabriel Barth-Maron, Nikola Momehev, Danila Sinopalnikov. Piotr Stanczyk, Sabela Ramos, Anton Raichuk, Damien Vincent, Léonard Hussenot, Robert Dadashi, Gabriel Dulac-Arnold, Manu Orsini, Alexis Jacq, Johan Ferret, Nino Vieillard, Seyed Kam- yar Seyed Ghasemipour, Sertan Girgin, Olivier Pietquin, Feryal Behbahani, Tamara Norman, Abbas Abdolmaleki, Albin Cassirer, Fan Yang, Kate Baumli, Sarah Henderson, Abe Friesen, Ruba Haroun, Alex Novikov, Sergio Gómez. Colmenarejo, Serkan Cabi, Caglar Gulcehre, Tom Le Paine, Srivatsan Srinivasan, Audrew Cowie, Ziyu Wang. Bilal Piot, and Nando de Freitas. Acme: A research framework for distributed reinforcement learning. arXiv preprint arXiv:2006, 00979, 2020. 8, 15
[20] Maximilian Igl. Daewoo Kim, Alex Kuefler, Paul Mougin, Punit Shah, Kyriacos Shiarlis, Dragomir Anguelov, Mark Palatucci, Brandyn White, and Shimon Whiteson. Symphony: Learning realistic and diverse agents for autonomous driving simulation. arXiv preprint 220.:22005003195. 2022. 3, 10
11

---

[21] David Isele. Reza Rahimi, Akansel Cosgun. Kaushik Subramanian, and Kikuo Fujimura. Navigating occluded intersections with autonomous vehicles using deep reinforcement learning. In 2018 IEEE
International Conference on Robotics and Automation (ICRA), pages 2034-2039. IEEE, 2018. 3
[22] Alex Kendall, Hawke, David Janz, Przemyslaw Mazur, Daniele Reda. John-Mark Allen, Vinh-Dieu Lam, Alex Bewley, and Amar Shah. Learning to drive in day. In 2019 International Conference on Roboties and Automation (ICRA), pages 8248 8254, IEEE, 2019.
[23] Parth Kothari, Christian Perone, Luca Bergamini, Alexandre Alahi, and Peter Ondruska. Drivergym: Democratising reinforcement learning for autonomous driving. arXiv preprint arXiv:21 2021. 3 [24] Edouard Leurent. An environment for autonomous driving sscision-making https:/ /ttp://////////////////gg thub. com/
eleur warreentthhhaay. env, 2018. 3
[25] Quanyi Li, Zhenghao Peng, Lan Feng, Qihang Zhang, Zhenghai Xuc, and Bolei Zhou. Metadrive: Composing diverse driving scenarios for generalizable reinforcement learning. IEEE transactions on pattern amalysis and muchine intelligence, 2022. 2, 3,6 6
[26] Ming Liang, Bin Yang. Rui Hu, Yun Chen, Renjie Liao, Song Feng, and Raquel Urtasun. Learning lane graph representations for motion forecasting. In European Conference on Computer Vision, pages 541-556.
Springer, 2020. 3
[27] Yiren Lu, Justin Fu, George Tucker, Xinlei Pan, Eli Bronstein, Becca Roelofs, al. Imitation is not
enough: Robustifying imitation with reinforcement learning for challenging driving scenarios. In NeurIPS
2022 Machine Learning for Autonomous Driving Workshop. 2022. 3
[28] Sivabalan Manivasagam, Shenlong Wang. Wei-Chiu Ma, Kelvin Ka Wing Wong, Wenyuan Zeng, and Raquel Urtasun. Systems and methods for generating synthetic sensor data via machine learning, Sept. 24
2020. US Patent App. 16/826,990. 2
[29] Mark Martinez, Chawin Sitawarin, Kevin Finch, Lennart Meineke, Alex Yablonski, and Alain Korhauser. Beyond grand theft auto for training, testing and enhancing deep learning in self driving cars, arXiv
preprint arXiv: 1277999. 2017.
[30] Miguel Angel Zamora Mora, Momchil Peychev, Sehoon Ha, Martin Vechev, and Stelian Coros. Pods: Policy optimization via differentiable simulation. In International Conference on Machine Learning, pages
1805-7817, PMLR, 2021. 10
[31] Khan Muhammad, Amin Ullah, Jaime Lloret, Javier Del Ser, and Vietor Hugo C de Albuquerque. Deep learning for safe autonomous driving: Current challenges and future directions. IEEE Transactions on Intelligent Transportation Systems, 22(()):311 0-366, 2020. 2
[32] Alistair Muldal, Yotam Doron, John Aslanides, Tim Harley, Tom Ward, and Siqi Liu. dm_env: A python interface for reinforcement learning environments, 2019.
[33] Matthias Müller, Vincent Casser, Jean Lahoud, Neil Smith, and Bernard Ghanem. Sim4cv: A rhoto-rcalistic simulator for computer vision applications. International Journal of Computer Vision, 126(9):902-99)
2018. 3,3
[34] Nigamaa Nayakanti, Rami Al-Rfou, Aurick Zhou, Kratarth Goel, Khaled Refaat, and Benjamin Sapp. Wayformer: Motion forecasting via simple & efficient attention networks. arXiv preprint 2rXiv:2007.05844,
2022. 3
[35] Nigamaa Nayakanti, Rami AI-Rfou, Aurick Zhou, Kratarth Goel, Khaled S. Refaat, and Benjamin Sapp.
Wayformer: Motion forecasting via simple and efficient attention networks, 2022. 8, 9, 15
[36] Jiquan Ngiam, Benjamin Caine, Vijay Vasudevan, Zhengdong Zhang, Hao-Tien Lewis Chiang, Jeffrey Ling, Rebecca Roelofs, Alex Bewley, Chenxi Liu, Ashish Venugopal, et al. Scene transformer: unified architecture for predicting multiple agent trajectories. arXiv preprint arXiv:2106.0844
2021. 3
[37] Blažej Osiński, Adam Jakubowski, Pawel Ziecina, Piotr Miloš, Christopher Galias, Silviu Homoceanu, and Henryk Michalewski. Simulation sanuuetion-based reinforcement learning for real world autonomous driving. In 2020 IEEE International Conference on Robotics and Automation (ICRA), pages 6411 6411-6418 6418. IEEE, 2020.
1, 10
[38] Praveen Palanisamy. Multi-agent connected autonomous driving using deep reinforcement learning. In
2020 International Joint Conference on Neural Networks (IJCNN), pages 1-7. IERE, 2020. 3
[39] Dean A Pomerleau. Alvinn: An autonomous land vehicle in neural network. Advances in neural information processing systems, 1, 1988. 3
[40] Craig Quiter. Deepdrive zero, June 2020.
[41] Nicholas Rhinehart, Rowan McAllister, Kris M. Kitani, and Sergey Levine. PRECOG: prediction condi-
tioned on goals visual multi-agent settings. CORR, ahs/1105.00 00000.12200 296., 2019. 3
[42] Stéphane Ross, Geoffrey Gordon, and Drew Bagnell. A reduction of imitation learning and structured prediction to no-regret online learning. In Proceedings of the fourteenth international conference on artificial intelligence and statistics, pages 627 627-635. JMLR Workshop and Conference Proceedings. 2011.
[43] Amit Sabne. Xla: Compiling machine learning for peak performance, 2020. 6
[44] Anirban Santara, Sohan Rudra, Sree Aditya Buridi, Meha Kaushik, Abhishek Naik, Bharat Kaul, and Balaraman Ravindran. Madras: Multi agent driving simulator. Journal of Artificial Intelligence Research, 10:1517-1555. 2021. 3
12

---

[45] Tom Schaul, John Quan, Ioannis Antonoglou, and David Silver. Prioritized experience replay. arXiv preprini arXiv:1511.05952 2015. 8, 15
[46] Adam Scibior, Vasileios Lioutas, Daniele Reda, Peyman Bateni, and Frank Wood. Imagining the road ahead: Multi-agent trajectory prediction via differentiable simulation. In 2021 IEEE International Intelligent Transportation Systems Conference (ITSC), pages 720-725, 2021. 2,
[47] Qiao Sun, Xin Huang. Brian C Williams, and Hang Zhao. Intersim: Interactive traffic simulation via explicit relation modeling, arXiv preprint 0rXiv:22/0.144133 2022. 3
[48] Matthew Tancik, Vincent Casser, Xinchen Yan, Sabeek Pradhan, Ben Mildenhall, Pratul P Srinivasan. Jonathan T Barron, and Henrik Kretzschmar, Block-nerf: Scalable large scene neural view synthesis. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pages 8248-8258
2022. 2
[49] Charlie Tang and Russ R Salakhutdinov. Multiple futures prediction. Advances in Neural Information Processing Systems, 32, 2019.
[50] Josh Tobin, Rachel Fong, Alex Ray, Jonas Schneider, Wojciech Zaremba, Pieter Abbeel. Domain randomization for transferring deep neural networks from simulation to the real world. In 2017 JEEE/RSJ international conference on intelligent robots and systems (IROS), pages 23-30. IEEE, 2017. 1, 10
[51] Martin Treiber, Ansgar Hennecke, and Dirk Helbing. Congested traffic states empirical observations and microscopic simulations. Physical review E, 12(2))1805,, 2000.
[52] Eugene Vinitsky, Nathan Lichtlé, Xiaomeng Yang, Brandon Amos, and Jakob Foerster. Nocturne scalable driving benchmark for bringing multi-agent learning one step closer to the real world. arXiv preprint or Xiv:2206. 09889, 2022. 2, 3,6
[53] Matt Vitelli, Yan Chang, Yawei Yc, Ana Ferreira, Maciej Wolzzyk, Blažej Osifiski, Moritz Niendorf, Hugo Grimmett, Qiangui Huang, Ashesh Jain, Safetynet: Safe planning real-world self-driving vehicles using machine-lcanee Icarned policies. 2022 International Conference on Robotics and Automation (ICRA),
pages 897-904. IEEE, 2022. 3
[54] Pin Wang, Ching- Yao Chan, and Amnad La Fortellc. A reinforeement lcarning based approach for automated lane change mancuvers. 2018 IEEE Intelligent Vehicles Symposium (/V), pages 1379-1384.
IEEE, 2018. 3
[55] Bernhard Wymann, Eric Espié, Christophe Guionneau, Christos Dimitrakakis, Rémi Coulom, and Andrew Sumner. Tores, the open racing car simulator. Software available http://torcs. sourceforge. net, 4(6):2, 2000. 3
[56] Danfei Xu, Yuxiao Chen, Boris Ivanovic, and Marco Pavone, Bits: Bi-level imitation for traffic simulation. arXiv preprint arXiv:2208. 12403, 2022. 3
[57] Ming Zhou, Jun Luo, Julian Villella, Yaodong Yang, David Rusu, Jiayu Miao, Weinan Zhang, Montgomery Alban, Iman Fadakar, Zheng Chen, al. Smarts: Scalable multi-agent reinforcement learning training school for autonomous driving. arXiv preprint urXiv:2010.09776, 2020.
13


Reviews:
[{"review_id": "mU75mJO3vN", "reviewer": "Reviewer_DvxP", "summary": "", "questions": "", "limitations": "Yes", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "The main strength of this submission is the introduction of DriveMax, a versatile and efficient autonomous driving simulator that allows various configurations of the learning agent and the environment. This contributes significantly to the field of autonomous driving research by providing a tool for researchers to design and evaluate different learning and planning algorithms in a controllable and reproducible environment. The research appears to be of high quality, as evidenced by comprehensive benchmarks and thoughtful experimental designs. The simulator, combined with detailed benchmarks and baseline planning agents, can serve as a valuable resource for the research and development of safer and more efficient self-driving technologies. ", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ""}, {"review_id": "IUFaPIVF40", "reviewer": "Reviewer_ngUf", "summary": "", "questions": "", "limitations": "Limitations should be discussed as self-driving is a safety-critical application. This could be emphasizing that the policy trained in the simulator can not be deployed on a real car due to the sim2real gap. ", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "In recent years, much progress in the robotics field are made with IsaacGym which can run large-scale parallel physics simulations to expedite data sampling and policy training. I am glad to see that the autonomous driving community finally has our GPU-based simulator now. It can facilitate traversing large-scale real-world data quickly and would definitely benefit ML-based planner research. In addition, the differentiable system dynamics open new research opportunities for trajectory prediction and policy learning. For example, the policy can be optimized without estimating the value function as the derivate of the reward function can serve as a gradient for updating the policy network.\nThe experiment section proves that the simulator is ready for being a new testbed where users can develop/test new data-driven planning algorithms and enjoy the efficiency. ", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ""}, {"review_id": "OpKe7TQd5U", "reviewer": "Reviewer_1Xpz", "summary": "", "questions": "", "limitations": "The authors clearly note that DriveMax does not include any kind of sensor simulation. Moreover, the authors provide an extensive set of benchmarking studies in Section 5 that help to illustrate the potential weaknesses of the model.", "rating": 7, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "Although I am not an expert in the area of autonomous driving research, the authors have articulated the need for DriveMax quite well through a comparison to other comparable simulators. It is clear that a large amount of effort has been put towards making this simulator useful for researchers. This includes making it fully differentiable by programming it in JAX, providing both stateful and stateless versions of the main interface, and by designing DriveMax as a set of interoperable libraries, enabling researchers to pick and choose components that are most useful for them. In addition, the authors directly demonstrate the utility of their simulator through several direct benchmarks and ablation studies. ", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ""}]

Author rebuttal:

Meta-review: This paper presents a novel simulator, DriveMax, for autonomous driving aimed at facilitating extensive planning tasks. The research appears to be of high quality to the whole autonomous driving field. All reviewers give positive comments. Extensive studies are provided.The area chair thus decides to accept it.
Decision: Accept (Poster)