Waymax: An accelerated simulator for autonomous driving research.

Overview

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.

Installation

Waymax can be installed via pip using the following command:

conda create -n waymax python=3.10 -c conda-forge
conda activate waymax
pip install --upgrade pip
pip install waymo-open-dataset-tf-2-11-0==1.6.0
pip install -e .

Please refer to JAX for specific instructions on how to setup JAX with GPU/CUDA support if needed.

Configure access to Waymo Open Motion Dataset

Waymax is designed to work with the Waymo Open Motion dataset out of the box.

A simple way to configure access 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.

Please reference TF Datasets for alternative methods to authentication.

Components

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.

Dataloading

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)

Metrics

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:

Agents

The waymax.agents module defines intelligent simulated agents for realistic simulation. Waymax currently supports:

Log-playback
IDM

Environments and dynamics

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)

Tutorials

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.

Citing Waymax

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}}

SafeRoboticsLab/waymax