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

Waymax is a library for simulating and evaluating agents using scenarios instantiated from the Waymo Open Motion Dataset. It is designed to support behavioral research (single and multi-agent planning, behavior prediction, sim agents) using the minimal representation of objects and bounding boxes adopted by the Waymo Open Motion Dataset, without the need for users to run their own perception system from raw sensor data.

Waymax is entirely written using JAX, allowing users to easily run simulations on hardware accelerators, write programs in a single computation graph, and take advantage of JAX’s functional transforms such as JIT compilation.

All software, tutorials, and documentation can be found at https://github.com/waymo-research/waymax.

Data Management

Waymax supports loading data from the Waymo Open Motion Dataset, either locally from disk or streaming remotely through Google Cloud. This requires users to register for access to the Open Motion Dataset through the Waymo website (https://waymo.com/open/data/motion/).

Metrics

Waymax implements a number of metrics used to evaluate behavior, including:

• Log Divergence: L2 distance between simulated and logged poses.
• Collision: An indicator metric determining whether the vehicle collides with another object.
• Offroad: An indicator metric on whether the vehicle drives off the road.
• Wrong-Way: An indicator metric on whether the vehicle is driving on the wrong side of the road.
• Kinematics Infeasibility: An indicator metric on whether the vehicle’s action results in a kinematically infeasible transition.

Dynamics Models

Waymax supports two models for vehicle simulation - direct state-based control, and control via the kinematic bicycle model. Simulation is by default assumed to occur at 10Hz (the same as used by the Waymo Open Motion Dataset), although this can be adjusted. The bicycle model is controlled via acceleration and steering, and simulates the vehicle as a frictionless, rigid body. The state-based control model allows users to directly set the position, orientation, and/or velocity regardless of kinematic constraints.

Sim Agents

Waymax supports two modes of operation for simulating the behavior of uncontrolled objects in the scene - log playback and an IDM-based route-following model. Log playback replays the behavior of objects as they are recorded in the logged dataset, without regard to the actual states realized in a closed-loop rollout. The IDM-based model assumes vehicles follow the same path as recorded in the logged dataset, but adjusts the speed profile based on the Intelligent Driver Model (IDM) to speed up/slow down based on the positions of surrounding vehicles. Users can additionally plug in their own models into simulation to control arbitrary objects in the scene.

RL Interfaces

Waymax provides adapters to commonly used interfaces for Reinforcement Learning training. This includes a stateless simulation interface similar to the Brax API, as well as a stateful interface implementing the dm-env API.