Waymax: An accelerated simulator for autonomous driving research

Waymax is a lightweight, multi-agent, JAX-based simulator for autonomous driving research based on the Waymo Open Motion Dataset.

Getting Started

API Documentation

Installation

pip install git+https://github.com/waymo-research/waymax.git@main#egg=waymo-waymax

For additional installation options, see the Install Guide in the project README.