# Waymax: An Accelerated, Data-Driven Simulator for Large-Scale Autonomous Driving Research

arXiv: https://arxiv.org/abs/2310.08710 · HTML: https://mlsanity.com/paper/2310.08710/

## Key facts

- Cited by 160 papers in the map (data through 14 Oct 2026).
- Placed in World models for autonomous driving, one of 11 directions in World models.
- Builds on 41 of its 55 references that are in the map.
- First posted on arXiv 12 Oct 2023.

## Placed in

- World models / [World models for autonomous driving](https://mlsanity.com/wm/world-models-for-autonomous-driving/)
- End-to-end autonomous driving / [Imitation-learning end-to-end planners](https://mlsanity.com/driving/imitation-learning-end-to-end-planners/)

## Builds on (most cited first, 30 of 41)

- [CARLA: An Open Urban Driving Simulator](https://mlsanity.com/paper/1711.03938/) (2017)
- [Domain Randomization for Transferring Deep Neural Networks from Simulation to the Real World](https://mlsanity.com/paper/1703.06907/) (2017)
- [Prioritized Experience Replay](https://mlsanity.com/paper/1511.05952/) (2015)
- [End to End Learning for Self-Driving Cars](https://mlsanity.com/paper/1604.07316/) (2016)
- [Block-NeRF: Scalable Large Scene Neural View Synthesis](https://mlsanity.com/paper/2202.05263/) (2022)
- [Large Scale Interactive Motion Forecasting for Autonomous Driving : The Waymo Open Motion Dataset](https://mlsanity.com/paper/2104.10133/) (2021)
- [NuPlan: A closed-loop ML-based planning benchmark for autonomous vehicles](https://mlsanity.com/paper/2106.11810/) (2021)
- [End-to-end Driving via Conditional Imitation Learning](https://mlsanity.com/paper/1710.02410/) (2017)
- [MultiPath: Multiple Probabilistic Anchor Trajectory Hypotheses for Behavior Prediction](https://mlsanity.com/paper/1910.05449/) (2019)
- [Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulation](https://mlsanity.com/paper/2106.13281/) (2021)
- [ChauffeurNet: Learning to Drive by Imitating the Best and Synthesizing the Worst](https://mlsanity.com/paper/1812.03079/) (2018)
- [Learning by Cheating](https://mlsanity.com/paper/1912.12294/) (2019)
- [Learning Lane Graph Representations for Motion Forecasting](https://mlsanity.com/paper/2007.13732/) (2020)
- [DenseTNT: End-to-end Trajectory Prediction from Dense Goal Sets](https://mlsanity.com/paper/2108.09640/) (2021)
- [MetaDrive: Composing Diverse Driving Scenarios for Generalizable Reinforcement Learning](https://mlsanity.com/paper/2109.12674/) (2021)
- [Causal Confusion in Imitation Learning](https://mlsanity.com/paper/1905.11979/) (2019)
- [Wayformer: Motion Forecasting via Simple & Efficient Attention Networks](https://mlsanity.com/paper/2207.05844/) (2022)
- [Learning to Drive in a Day](https://mlsanity.com/paper/1807.00412/) (2018)
- [PRECOG: PREdiction Conditioned On Goals in Visual Multi-Agent Settings](https://mlsanity.com/paper/1905.01296/) (2019)
- [Multiple Futures Prediction](https://mlsanity.com/paper/1911.00997/) (2019)
- [Scene Transformer: A unified architecture for predicting multiple agent trajectories](https://mlsanity.com/paper/2106.08417/) (2021)
- [Acme: A Research Framework for Distributed Reinforcement Learning](https://mlsanity.com/paper/2006.00979/) (2020)
- [Imitation Is Not Enough: Robustifying Imitation with Reinforcement Learning for Challenging Driving Scenarios](https://mlsanity.com/paper/2212.11419/) (2022)
- [SMARTS: Scalable Multi-Agent Reinforcement Learning Training School for Autonomous Driving](https://mlsanity.com/paper/2010.09776/) (2020)
- [VISTA 2.0: An Open, Data-driven Simulator for Multimodal Sensing and Policy Learning for Autonomous Vehicles](https://mlsanity.com/paper/2111.12083/) (2021)
- [BITS: Bi-level Imitation for Traffic Simulation](https://mlsanity.com/paper/2208.12403/) (2022)
- [SafetyNet: Safe planning for real-world self-driving vehicles using machine-learned policies](https://mlsanity.com/paper/2109.13602/) (2021)
- [Symphony: Learning Realistic and Diverse Agents for Autonomous Driving Simulation](https://mlsanity.com/paper/2205.03195/) (2022)
- [Nocturne: a scalable driving benchmark for bringing multi-agent learning one step closer to the real world](https://mlsanity.com/paper/2206.09889/) (2022)
- [Multi-Agent Connected Autonomous Driving using Deep Reinforcement Learning](https://mlsanity.com/paper/1911.04175/) (2019)

## Built on by (most cited first, 30 of 160)

- [NAVSIM: Data-Driven Non-Reactive Autonomous Vehicle Simulation and Benchmarking](https://mlsanity.com/paper/2406.15349/) (2024)
- [A Survey on Multimodal Large Language Models for Autonomous Driving](https://mlsanity.com/paper/2311.12320/) (2023)
- [Bench2Drive: Towards Multi-Ability Benchmarking of Closed-Loop End-To-End Autonomous Driving](https://mlsanity.com/paper/2406.03877/) (2024)
- [Understanding World or Predicting Future? A Comprehensive Survey of World Models](https://mlsanity.com/paper/2411.14499/) (2024)
- [End-to-End Driving with Online Trajectory Evaluation via BEV World Model](https://mlsanity.com/paper/2504.01941/) (2025)
- [Pseudo-Simulation for Autonomous Driving](https://mlsanity.com/paper/2506.04218/) (2025)
- [DriveArena: A Closed-loop Generative Simulation Platform for Autonomous Driving](https://mlsanity.com/paper/2408.00415/) (2024)
- [DrivingGPT: Unifying Driving World Modeling and Planning with Multi-modal Autoregressive Transformers](https://mlsanity.com/paper/2412.18607/) (2024)
- [Hydra-MDP++: Advancing End-to-End Driving via Expert-Guided Hydra-Distillation](https://mlsanity.com/paper/2503.12820/) (2025)
- [Robust Autonomy Emerges from Self-Play](https://mlsanity.com/paper/2502.03349/) (2025)
- [NeuroNCAP: Photorealistic Closed-loop Safety Testing for Autonomous Driving](https://mlsanity.com/paper/2404.07762/) (2024)
- [A Survey of World Models for Autonomous Driving](https://mlsanity.com/paper/2501.11260/) (2025)
- [LidarDM: Generative LiDAR Simulation in a Generated World](https://mlsanity.com/paper/2404.02903/) (2024)
- [RealGen: Retrieval Augmented Generation for Controllable Traffic Scenarios](https://mlsanity.com/paper/2312.13303/) (2023)
- [GPUDrive: Data-driven, multi-agent driving simulation at 1 million FPS](https://mlsanity.com/paper/2408.01584/) (2024)
- [Open-TI: Open Traffic Intelligence with Augmented Language Model](https://mlsanity.com/paper/2401.00211/) (2023)
- [PlanAgent: A Multi-modal Large Language Agent for Closed-loop Vehicle Motion Planning](https://mlsanity.com/paper/2406.01587/) (2024)
- [A Survey of Sim-to-Real Methods in RL: Progress, Prospects and Challenges with Foundation Models](https://mlsanity.com/paper/2502.13187/) (2025)
- [Foundation Models in Autonomous Driving: A Survey on Scenario Generation and Scenario Analysis](https://mlsanity.com/paper/2506.11526/) (2025)
- [VaViM and VaVAM: Autonomous Driving through Video Generative Modeling](https://mlsanity.com/paper/2502.15672/) (2025)
- [CaRL: Learning Scalable Planning Policies with Simple Rewards](https://mlsanity.com/paper/2504.17838/) (2025)
- [Multi-Agent Reinforcement Learning for Autonomous Driving: A Survey](https://mlsanity.com/paper/2408.09675/) (2024)
- [Hidden Biases of End-to-End Driving Datasets](https://mlsanity.com/paper/2412.09602/) (2024)
- [Towards Knowledge-driven Autonomous Driving](https://mlsanity.com/paper/2312.04316/) (2023)
- [The Role of World Models in Shaping Autonomous Driving: A Comprehensive Survey](https://mlsanity.com/paper/2502.10498/) (2025)
- [Scenario Dreamer: Vectorized Latent Diffusion for Generating Driving Simulation Environments](https://mlsanity.com/paper/2503.22496/) (2025)
- [SLEDGE: Synthesizing Driving Environments with Generative Models and Rule-Based Traffic](https://mlsanity.com/paper/2403.17933/) (2024)
- [Text2LiDAR: Text-guided LiDAR Point Cloud Generation via Equirectangular Transformer](https://mlsanity.com/paper/2407.19628/) (2024)
- [Improving Agent Behaviors with RL Fine-tuning for Autonomous Driving](https://mlsanity.com/paper/2409.18343/) (2024)
- [Choose Your Simulator Wisely: A Review on Open-source Simulators for Autonomous Driving](https://mlsanity.com/paper/2311.11056/) (2023)

## Abstract (arXiv)

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 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) 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.