Waymax
Project record
Waymax is a multi-agent driving simulator written in JAX that builds scenarios from the Waymo Open Motion Dataset and represents road users as bounding boxes rather than raw sensor data. It includes data loading, visualization, metrics, simulated agents, and adapters to reinforcement learning interfaces such as dm-env and Brax.16
- Repository: waymax repository (external site: github.com)
- Documentation: Waymax documentation (external site: waymo-research.github.io)
- License: Waymax License Agreement for Non-Commercial Use (external site: github.com)
- Paper: Waymax paper (arXiv 2310.08710) (external site: arxiv.org)
Availability and license
Overall availability
Source code is public on GitHub and installable with pip from the main branch, free of charge for non-commercial purposes only. Using it with the Waymo Open Motion Dataset requires registering for Waymo Open Dataset access.12
Availability is separate from permission: read the license before using or redistributing.
Custom non-commercial license (headed October 17, 2023). It bars use for real-world vehicle operation or its development, testing real vehicles, commercial driving simulation, production systems, or other primarily commercial purposes; unmodified copies may not be redistributed, though research results and derivative work may be shared with a required notice that passes the terms on. The license also forbids using the Waymax code or documentation to train or improve AI foundation models. The license field in setup.py's package metadata reads Apache-2.0, but the header of setup.py itself and the LICENSE file apply the non-commercial agreement.23
Public materials checklist
| Item | Status | Notes and evidence |
|---|---|---|
| Source codeIs the source code publicly readable? | Public | Source is published on GitHub under the non-commercial license.12 |
| DocumentationIs user documentation published? | Public | A documentation site with an API reference and getting-started material, plus Colab tutorial notebooks.51 |
| InstallationAre installation instructions or packages publicly available? | Public | The README documents installation with pip from the GitHub main branch; setup.py requires Python 3.10 or later.13 |
| Supported platformsAre supported operating systems or hardware documented? | Public | Built on JAX; the README says it can run on hardware accelerators such as GPUs and TPUs and refers to JAX's instructions for GPU/CUDA setup.1 |
| Release statusAre versioned releases published? | Not public | The GitHub releases page lists no releases; setup.py declares version 0.1.0 and the README installs from the main branch.431 |
What it is useful for
Documented for autonomous driving behavior research, including closed-loop simulation for planning and simulated-agent research and open-loop behavior prediction, on GPUs or TPUs; a tutorial covers preparing Waymo Open Sim Agents Challenge submissions.1
Run and use notes
- setup.py lists Python 3.10 or later and dependencies including JAX 0.4.6+, TensorFlow 2.11+, Flax, chex, and dm_env.3
Organization context
U.S. eligibility
Sources
This listing is not an endorsement, a safety assessment, or a federal approval.