Source-linked AI summary
Waymax: An Accelerated, Data-Driven Simulator for Large-Scale Autonomous Driving Research
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 Yang, Brandyn White, Aleksandra Faust, Rowan McAllister, Dragomir Anguelov, Benjamin Sapp
TL;DR
Autonomous-driving simulation must combine realistic multi-agent behavior with the speed needed for large-scale training and evaluation. Waymax addresses this gap with a differentiable, hardware-accelerated simulator initialized from diverse real-world driving data and equipped with routes, behavior models, and closed-loop benchmarks. The evaluations highlight route conditioning and show that reinforcement learning can overfit to simulated agents.
Problem
Autonomous-driving research needs simulation that is both high-throughput and realistic in diverse, interactive multi-agent behavior, while simulation-to-real transfer remains difficult.
Method
Waymax is a differentiable, hardware-accelerated multi-agent simulator using real-world driving data, configurable state and routes, behavior models, and closed-loop training and evaluation workflows.
Results
The benchmarks show that route-conditioned agents follow routes better with lower off-road, collision, and log ADE rates, while RL trained against IDM agents can overfit to their easier behavior.
Takeaways & Limitations
Waymax provides flexible infrastructure and strong baselines for large-scale autonomous-driving planning research across imitation learning, reinforcement learning, and interactive simulated agents.
Takeaways & Limitations
Simulation-to-real transfer remains an important limitation because performance and safety in Waymax are not guaranteed to carry over to the real world.
Abstract
from arXiv · showhide
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.
1 Introduction
Waymax addresses the need for autonomous-driving simulation that is both realistic in multi-agent behavior and fast enough for large-scale machine-learning workflows. It uses diverse real-world driving data, hardware acceleration, and closed-loop benchmarks to support planning research.
- Motivation: Simulation must balance speed for large-scale training and evaluation with realistic, diverse vehicle behavior that reduces the sim-to-real gap.The paper identifies speed and realism as the two primary simulator challenges.
- Motivation: Existing simulators still face difficulty generating diverse scenarios and realistic behavior for other agents while supporting high-throughput machine-learning workflows.The stated need is for realistic closed-loop behavior together with high speed and throughput.
- Waymax: Waymax is a differentiable, hardware-accelerated, multi-agent simulator built from real-world driving data, including over 250 hours in the Waymo Open Motion Dataset.It uses JAX to run simulation on GPUs and TPUs and supports urban obstacles and high-level routes.
- Contributions: Waymax provides hardware acceleration, real-data features and routes, diverse scenarios, simulated agents, and closed-loop benchmarks for imitation and reinforcement learning.The benchmark suite covers different training setups and interactive agents.
2 Related Work
Related work spans photorealistic and data-driven autonomous-driving simulators, alongside imitation- and reinforcement-learning methods for driving agents. Waymax emphasizes accelerated, differentiable, multi-agent benchmarking with real expert data and simulated-agent models.
- Driving simulators: Waymax differs from related simulators by combining hardware-accelerated training, in-graph agent models, differentiable simulation, real WOMD data, routes, and reactive agents.The comparison highlights support for pedestrians, cyclists, traffic lights, and goal-conditioned planning.
- Learning-based agents: Behavior cloning and reinforcement learning address autonomous-driving policy learning through demonstrations, feedback, or hand-designed rewards, but large-scale real-world benchmarking remains challenging.The paper positions Waymax as infrastructure for standard open- and closed-loop training and evaluation.
3 Simulator Features
Waymax provides configurable, real-data-initialized multi-agent simulation with route-aware state representations, interchangeable dynamics and metrics, and reactive behavior models. Its interface supports closed-loop evaluation of safety, correctness, comfort, and progress.
- Simulator Features: Waymax initializes scenarios from real driving logs and supports arbitrary numbers of controllable vehicles, pedestrians, cyclists, and traffic lights.The simulator is designed for complex interactions and allows components to be modified or replaced.
- State and Observation spaces: The state includes dynamic object poses, velocities, dimensions, and traffic-light colors plus static roadgraph and on-route or off-route paths.User-defined observation functions can induce partial observability; a default function uses agent-relative coordinates and roadgraph subsampling.
- Routes: Feasible paths are computed by depth-first traversal of the roadgraph, with on-route paths identified by agreement with the logged trajectory’s road.The resulting paths describe legal driving options from the AV’s initial position.
- Dynamics and Actions: Users can choose dynamics and action representations, including delta actions for all objects and bicycle actions using acceleration and steering curvature for vehicles.The dynamics model determines expected actions and state evolution.
- Metrics: Closed-loop metrics evaluate safety, traffic-rule correctness, comfort, progress, route adherence, collisions, kinematic feasibility, and displacement error.Metrics are computed by running agents in simulation rather than evaluating independent logged timesteps.
- Simulated Agents: Waymax includes an IDM-based reactive vehicle model that adjusts speed to avoid collisions and accelerate on free roads.The framework also permits learned behavior models for simulated objects.
4 Software API
Waymax exposes modular, functionally pure interfaces for multi-agent and single-agent autonomous-driving simulation. Its JAX-based design supports accelerator execution, in-graph training, and compatibility with stateful workflows.
- Software components: Waymax is organized as interoperable libraries for data structures, distributed data loading, simulator components, and a Gym-like environment interface.Each simulator component can be modified, replaced, or used independently.
- Environment interface: The simulator uses pure reset and step functions, enabling JAX compilation, functional workflows, and algorithms requiring backtracking.Adapters are also provided for stateful codebases.
- Accelerated execution: Waymax supports GPU and TPU acceleration and in-graph training, keeping simulation and training on the accelerator without host-communication bottlenecks.The implementation uses JAX operations compiled through XLA.
- Simulation workflows: The MultiAgentEnvironment accepts simultaneous actions and a control mask for selected scene objects.This interface targets multi-agent and sim-agent problems.
- Simulation workflows: The PlanningAgentEnvironment controls only the ego vehicle while user-specified sim agents or log playback control the remaining agents.Waymax therefore supports a single-agent planning workflow alongside multi-agent simulation.
5 Experiments
Waymax is evaluated for computational performance and for closed-loop planning with imitation, reinforcement-learning, and expert agents. The experiments show strong runtime scaling, benefits from route conditioning, and sensitivity to the choice of simulated agents.
- Runtime Benchmark: GPU execution exceeds 1000Hz for batched Step at batch size 1 and reaches over 5000Hz per example with batch size 16.Step computes both state transition and reward; batch size 16 takes 2.86ms, equivalent to over 5000Hz per example.
- Runtime Benchmark: GPU Rollout provides a consistent 2 orders of magnitude speedup, and evaluates 44K scenarios on 8 V100s in under 2 minutes.Rollout covers an entire 80-step episode and is faster than running Step 80 times.
- Planning Benchmark Results: The benchmark evaluates expert, behavior-prediction, behavior-cloning, and DQN agents using route conditioning and standard planning metrics.The environment controls the autonomous vehicle while simulated agents or log playback control other agents.
- Planning Benchmark Results: Expert agents achieve low off-road and collision rates, while discrete bicycle actions perform comparably to other expert action spaces.The nominal values capture noise in bounding boxes and logged data and serve as a lower bound; comparable performance indicates sufficiently fine discretization.
- Planning Benchmark Results: Discrete action spaces perform best for open-loop imitation and outperform the adapted Wayformer model.The authors suggest discrete actions may more easily model multimodal behavior and use the result as a check on Waymax training data.
- Planning Benchmark Results: Route conditioning improves route following while lowering off-road rate, collision rate, and log ADE, whereas RL agents can overfit easier IDM agents.An IDM-trained agent evaluated against logged agents has a collision rate over 4x higher than when evaluated against IDM agents.
6 Conclusion
Waymax provides diverse real-data scenarios, hardware acceleration, distributed training, and interoperable components for autonomous-driving research. It also enables differentiable simulation and benchmarks planning methods, while sim-to-real transfer remains an important limitation.
- 6 Conclusion: Waymax provides diverse scenarios from real driving data and supports hardware acceleration and distributed training for efficient machine-learned model training.Its interoperable libraries cover data loading, metric computation, and simulation across research problems.
- 6 Conclusion: Waymax benchmarks common planning approaches with ablations over dynamics and action representations, providing baselines for future work.
- 6 Conclusion: Differentiable simulation is enabled by assembling the entire simulator within a single JAX computation graph.The authors identify pass-through gradients as a promising direction for reducing policy-optimization gradient-estimate variance.
- 6 Conclusion: Sim-to-real transfer remains difficult because simulated performance and safety do not guarantee the same outcomes in the real world.Waymax uses real-world data to seed scenarios, but the authors identify transfer as an important limitation of simulation-based frameworks.
A.1 Dynamics Definitions
Waymax defines delta and bicycle action spaces for updating vehicle states and recovering actions or dynamics from consecutive states. The bicycle formulation uses acceleration and steering curvature to approximate vehicle motion.
- Delta Action Space: The delta action space represents state changes with (∆x, ∆y, ∆θ) and updates the agent state over a timestep.The state includes position, yaw, and velocities in the x and y directions.
- Delta Action Space: Inverse kinematics recovers behavior-cloning actions as the differences between consecutive positions and yaw angles.The action components are ∆x = x′ − x, ∆y = y′ − y, and ∆θ = θ′ − θ.
- Bicycle Action Space: The bicycle action space approximates vehicle dynamics by minimizing the discrepancy between predicted and recorded vehicle coordinates.Vehicle state includes global position, yaw angle, and x- and y-direction speeds.
- Bicycle Action Space: Forward dynamics use acceleration, steering curvature, and timestep length to calculate the vehicle’s next state.Inverse kinematics estimates acceleration and steering curvature from two consecutive states.
- Bicycle Action Space: The steering wheel angle determines steering curvature through the steering ratio, with axle length also defined for the vehicle model.The relation explicitly uses θwheel/STEER_RATIO to obtain the front-wheel steer angle.
A.2 Training Details
The benchmarks train behavior-cloning and DQN agents with specified neural architectures, action distributions, optimization settings, and action-space hyperparameter searches.
- Behavior Cloning: Behavior cloning uses a Wayformer encoder, a four-layer residual MLP, and likelihood maximization for expert actions.Continuous actions use a 10-component Gaussian Mixture Model Tanh-squashed head, while discrete actions use softmax probabilities.
- DQN: DQN uses prioritized replay double DQN with the discrete behavior-cloning architecture interpreted as a Q-network.Training uses γ = 0.99, learning rate 5 ∗ 10^-5, 1-step Q-learning, and 30 million actor steps.
- Hyperparameter Selection: Hyperparameters are selected by grid search over learning rates and action spaces for behavior cloning, and action spaces for DQN.Behavior cloning searches four action-space variants; DQN searches Bicycle-Discrete and Delta-Discrete.
A.3 Ablation Study: Runtime and Memory with Number of Objects
The ablation study measures runtime and memory as the number of simulated objects increases on CPU and GPU configurations. Runtime includes reset, transition, and metrics, while memory has a fixed cost followed by roughly linear scaling.
- Experimental Setup: The study examines runtime and memory against the number of simulated objects using AMD EPYC CPU and Nvidia V100 GPU configurations.The CPU machine has 64GB RAM.
- Results: Table 6 reports runtime in milliseconds and peak memory in megabytes for different object counts and batch sizes.BS-1 denotes batch size 1 and BS-16 denotes batch size 16.
- Runtime: Figure 4 plots runtime in milliseconds against simulated objects, with runtime defined as Reset + Transition + Metrics.CPU runtime scales linearly, while GPU performance is not saturated under the experimental parameters.
- Memory: Figure 5 plots rollout-sampled memory usage in megabytes against simulated objects, showing a fixed cost followed by roughly linear scaling.