SlowGuess/NeurIPS_2023_OCR huggingface dataset:
In Tab. 2, we present the runtime performance of Waymax using a CPU (Intel Xeon W-2135@3.7GHz) and a GPU (Nvidia-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 to 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  $1000\mathrm{Hz}$  for Step function and over  $2000\mathrm{Hz}$  if only considering the Transition. More importantly, as Waymax supports batching, Step only takes 2.86ms using a batch size of 16. Note this is much faster than running batch size one for 16 times and gives an equivalent runtime of over  $5000\mathrm{Hz}$  per example (i.e. closer to 500 times faster than using a 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 a slow operation.
Agent	Action Space	Train Sim Agent	Off-Road Rate (%)	Collision Rate (%)	Kinematic Infeasibility (%)	Log (m)	ADE	Route 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 (Discrete)	-	0.41	0.67	0.00	0.09	100.00
Wayformer	Delta	-	7.89	10.68	5.40	2.38	123.58
BC	Delta	-	4.14±2.04	5.83±1.09	0.18±0.16	6.28±1.93	79.58±24.98
BC	Delta (Discrete)	-	4.42±0.19	5.97±0.10	66.25±0.22	2.98±0.06	98.82±3.46
BC	Bicycle	-	13.59±12.71	11.20±5.34	0.00±0.00	3.60±1.11	137.11±33.78
BC	Bicycle (Discrete)	-	1.11±0.20	4.59±0.06	0.00±0.00	2.26±0.02	129.84±0.98
DQN	Bicycle (Discrete)	IDM	3.74±0.90	6.50±0.31	0.00±0.00	9.83±0.48	177.91±5.67
DQN	Bicycle (Discrete)	Playback	4.31±1.09	4.91±0.70	0.00±0.00	10.74±0.53	215.26±38.20
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 a percentage of episodes where the metric is 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 an 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 ExpertActor that derives growth-truth actions from logged trajectory. It is faster than running 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 a 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 to 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 a point of reference, we adapt the state-of-the-art Wayformer behavior prediction model [35] to the planning setting. Originally, the Wayformer predicts multiple 8-second future trajectories given a 1-second context history. To adapt it to the planning setting, we autoregressively feed in its predictions as the context history and choose the most likely trajectory. We found that making predictions at a lower frequency than the environment frequency improved performance, so we predict 5-step long trajectories and only replan every 5 steps.
Behavior Cloning We re-use the encoder portion of Wayformer [35] followed by a 4-layer residual MLP to maximize the log likelihood of the expert actions. For continuous actions, we used a 6-component Gaussian Mixture Model. For discrete actions, we used a softmax layer to compute action probabilities.
Model-Free Reinforcement Learning - DQN We used the Acme [19] implementation of prioritized replay double DQN [45].
We used the same architecture as in discrete BC for the Q-network, interpreting the logits of the model as Q-values.
For simplicity, we use a sparse reward penalizing collisions and off-road events:  $r_t = -\mathbb{I}_{\mathrm{collision}}(t) - \mathbb{I}_{\mathrm{off-road}}(t)$ .
5.3 Planning Benchmark Results
To showcase the flexibility of our environment, we trained a number of baselines on 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 off-road and collision rates. The nominal values represent noise in the bounding boxes and logged data and serve as a 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 is easier to model multi-modal behavior. Furthermore, it outperforms the adapted Wayformer model, likely due to the fact that it is trained explicitly for this task. This serves as a check that the Waymax environment is 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 Rate (%)	Log (m)	ADE	Route Progress Ratio (%)
Expert (Bicycle Discrete)	0.41	0.67	0.00	0.00	100.00
BC (Bicycle Discrete)	1.45±0.05	4.92±0.24	2.31±0.12	2.41±0.06	128.73±1.50
BC (Bicycle Discrete) + Route	1.11±0.20	4.59±0.06	0.96±0.09	2.26±0.02	129.84±0.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 RL agent tends to overfit or exploit the behavior of 'easier' IDM agents. Since IDM will stop for the SDC to avoid collisions, the RL agent does not have as much incentive to learn how to avoid collisions itself. We can see that when an IDM-trained agent is evaluated against logged agents, the collision rate is over  $4\mathrm{x}$  higher than when evaluated against IDM agents.
Train Agent	Eval Agent	Collision Rate (%)	Offroad Rate (%)	Progress Ration (%)
Playback	Playback	8.67±0.97	5.16±1.00	193.56±26.82
Playback	IDM	4.91±0.70	4.31±1.09	215.26±38.20
IDM	Playback	25.15±0.76	4.91±0.41	163.11±4.51
IDM	IDM	6.50±0.31	3.74±0.90	177.91±5.67
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 is more difficult to train effective planning agents using interactive sim agents.
We have presented Waymax, a 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-learned models. It is 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 of 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 a set of strong baselines for benchmarking future work.
In addition to 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 pass-through gradient to reduce the variance of the gradient estimate. We believe that this is a promising line of future work to be explored.
As mentioned previously, the problem of sim-to-real transfer is a critical issue in autonomous driving, as it is cheap and desirable to evaluate in simulation but difficult to guarantee that the same performance and level of safety will carry over to the real world. While in 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-based framework. A 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].
[45] Tom Schaul, John Quan, Ioannis Antonoglou, and David Silver. Prioritized experience replay. arXiv preprint 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, 3
[47] Qiao Sun, Xin Huang, Brian C Williams, and Hang Zhao. Intersim: Interactive traffic simulation via explicit relation modeling. arXiv preprint arXiv:2210.14413, 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 Pattern Recognition, pages 8248–8258, 2022. 2
[49] Charlie Tang and Russ R Salakhutdinov. Multiple futures prediction. Advances in Neural Information Processing Systems, 32, 2019. 3
[50] Josh Tobin, Rachel Fong, Alex Ray, Jonas Schneider, Wojciech Zaremba, and Pieter Abbeel. Domain randomization for transferring deep neural networks from simulation to the real world. In 2017 IEEE/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 in empirical observations and microscopic simulations. Physical review E, 62(2):1805, 2000. 5
[52] Eugene Vinitsky, Nathan Lichtle, Xiaomeng Yang, Brandon Amos, and Jakob Foerster. Nocturne: a scalable driving benchmark for bringing multi-agent learning one step closer to the real world. arXiv preprint arXiv:2206.09889, 2022. 2, 3, 6
[53] Matt Vitelli, Yan Chang, Yawei Ye, Ana Ferreira, Maciej Wolczyk, Błajej Osiński, Moritz Niendorf, Hugo Grimmett, Qiangui Huang, Ashesh Jain, et al. Safetynet: Safe planning for real-world self-driving vehicles using machine-learned policies. In 2022 International Conference on Robotics and Automation (ICRA), pages 897–904. IEEE, 2022. 3
[54] Pin Wang, Ching-Yao Chan, and Arnaud de La Fortelle. A reinforcement learning based approach for automated lane change maneuvers. In 2018 IEEE Intelligent Vehicles Symposium (IV), pages 1379-1384. IEEE, 2018. 3
[55] Bernhard Wymann, Eric Espie, Christophe Guionneau, Christos Dimitrakakis, Rémi Coulom, and Andrew Sumner. Torcs, the open racing car simulator. Software available at 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, et al. Smarts: Scalable multi-agent reinforcement learning training school for autonomous driving. arXiv preprint arXiv:2010.09776, 2020. 3