Waymax is a lightweight, multi-agent, JAX-based simulator for autonomous driving research based on the Waymo Open Motion Dataset. Waymax is designed to support research for all aspects of behavior research in autonomous driving - from closed-loop simulation for planning and sim agent research to open-loop behavior prediction. Objects (e.g. vehicles, pedestrians) are represented as bounding boxes, rather than raw sensor outputs, in order to distill behavior research into its simplest form.
As all components are entirely written in JAX, Waymax is easily distributed and deployed on hardware accelerators, such as GPUs and TPUs. Waymax is provided free of charge under the terms of the Waymax License Agreement for Non-Commercial Use.
Waymax can be installed via pip using the following command:
pip install --upgrade pip
pip install git+https://github.com/waymo-research/waymax.git@main#egg=waymo-waymax
Please refer to JAX for specific instructions on how to setup JAX with GPU/CUDA support if needed.
Waymax is designed to work with the Waymo Open Motion dataset out of the box.
A simple way to configure access via command line is the following:
-
Apply for Waymo Open Dataset access.
-
Install gcloud CLI
-
Run
gcloud auth login <your_email>with the same email used for step 1. -
Run
gcloud auth application-default login.
If you are using colab, run the following inside of the colab after registering in step 1:
from google.colab import auth
auth.authenticate_user()Please reference TF Datasets for alternative methods to authentication.
Structurally, Waymax is comprised of a collection of libraries for loading Open Motion data, visualization, computing common metrics, intelligent sim agents, and adapters to common RL interfaces such as dm-env. These libraries can be used as standalone modules, or used together in full closed-loop simulation.
The waymax.dataloader module contains utilities for loading data from the
Waymo Open Motion Dataset.
from waymax import config
from waymax import dataloader
scenarios = dataloader.simulator_state_generator(config.WOD_1_1_0_TRAINING)
scenario = next(scenarios)The waymax.metrics module defines commonly used metrics for evaluating agents.
These metrics can be used to evaluate simulated rollouts, or open-loop
predictions from behavior models. Supported metrics include:
The waymax.agents module defines intelligent simulated agents for realistic
simulation. Waymax currently supports:
- Log-playback
- IDM
The waymax.env module defines a stateless, closed-loop simulator interface as
well as adapters to common RL interfaces such as
dm-env and
brax.
A multi-agent simulation with rewards computed for all agents can be run as follows:
from waymax import env, config, dynamics, datatypes
# Initialization
dynamics_model = dynamics.InvertibleBicycleModel()
env_config = config.EnvironmentConfig()
scenarios = dataloader.simulator_state_generator(config.WOD_1_1_0_TRAINING)
waymax_env = env.MultiAgentEnvironment(dynamics_model, env_config)
# Rollout
state = waymax_env.reset(next(scenarios))
total_returns = 0
while not state.is_done:
action = datatypes.Action(data=..., valid=...) # Compute action here
total_returns += waymax_env.reward(state, action)
state = waymax_env.step(state, action)We provide a few colab tutorials for getting started:
- data_demo.ipynb shows how to load the data and use the top-down view visualization.
- multi_actors_demo.ipynb shows how to instantiate multiple agents and run a simple closed-loop simulation.
- wosac_submission_via_waymax.ipynb shows how to create a Waymo Open Sim Agents Challenge submission file.
If you use Waymax for your own research, please cite Waymax in accordance with the requirements of the Waymax License Agreement for Non-Commercial Use, including using the following bibtex entry:
@inproceedings{waymax,
title={Waymax: An Accelerated, Data-Driven
Simulator for Large-Scale Autonomous Driving Research},
author={Cole Gulino and Justin Fu and Wenjie
Luo and George Tucker and Eli Bronstein and Yiren Lu and Jean Harb and Xinlei Pan and Yan Wang and Xiangyu Chen and John
D. Co-Reyes and Rishabh Agarwal and Rebecca Roelofs and Yao Lu and Nico Montali and Paul Mougin and Zoey Yang and
Brandyn White and Aleksandra Faust, and Rowan McAllister and Dragomir Anguelov and Benjamin Sapp},
booktitle={Proceedings of the Neural Information Processing Systems Track on Datasets and
Benchmarks},year={2023}}
Please email any questions to waymax@google.com, or raise an issue on Github.