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 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 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.
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:
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:
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}}
此处可能存在不合适展示的内容,页面不予展示。您可通过相关编辑功能自查并修改。
如您确认内容无涉及 不当用语 / 纯广告导流 / 暴力 / 低俗色情 / 侵权 / 盗版 / 虚假 / 无价值内容或违法国家有关法律法规的内容,可点击提交进行申诉,我们将尽快为您处理。