Master Waymax & JAX
A hands-on learning project to understand Waymo's high-performance autonomous driving simulator. Explore data-driven simulation, JAX optimization, and sim agent research.
Why Learn Waymax?
Waymax achieves a breakthrough by combining real-world data, GPU acceleration, and multi-agent simulation—100x faster than traditional simulators while maintaining realism.
Data-Driven Simulation
Learn how Waymax uses real Waymo Open Motion Dataset (100K+ scenarios) instead of synthetic environments for realistic simulation.
JAX-Powered Performance
Achieve 1000+ Hz simulation speeds using JAX's JIT compilation, vmap vectorization, and GPU/TPU acceleration.
Multi-Agent Simulation
Simulate all agents in a scene, not just the ego vehicle. Train realistic sim agents that match human driving behavior.
Closed-Loop Evaluation
Understand why closed-loop testing is critical for AV validation and how to properly evaluate planning systems.
Deep Dive Papers
Comprehensive guides covering theory, implementation, mental models, and interview preparation for each major topic.
Waymax Deep Dive
Core simulator architecture, metrics, and evaluation framework
WOSAC Challenge
Sim agents challenge evaluation and winning strategies
JAX Scaling RL
Distributed training, vmap, pmap, and performance optimization
V-Max Framework
Complete RL training pipeline on top of Waymax
BehaviorGPT
State-of-the-art sim agent modeling with transformers
Sim-to-Real Gap
Neural rendering, world models, and bridging virtual to physical
Long-Tail Scenarios
Adversarial generation and safety-critical testing at scale
Distributed Training
PureJaxRL, actor-learner architectures, and billion-step training
Structured Learning Path
Progress from JAX fundamentals to advanced simulation techniques with our structured curriculum.
JAX Fundamentals
- Arrays & Immutability
- Automatic Differentiation
- JIT Compilation
- PyTrees
Waymax Basics
- Data Loading
- Metrics System
- Log Playback
- Agent Interfaces
Multi-Agent Simulation
- IDM Agents
- Vectorization
- Batch Simulation
- Performance Tuning
Advanced Topics
- RL Training
- Behavior Cloning
- Scenario Generation
- WOSAC Evaluation
Key Insights
Critical lessons from studying Waymo's simulation infrastructure
Waymax achieves 1000+ Hz on GPU with real-world data—100x faster than CARLA while maintaining realism through data-driven approach.
Functional paradigm, XLA compilation, and vmap enable massive parallelization that traditional OOP simulators cannot match.
Open-loop metrics don't correlate with real-world performance. Error compounding in closed-loop evaluation reveals true system capabilities.
Success means matching the distribution of human behavior, not just avoiding collisions. WOSAC metrics capture this nuance.
Ready to Start Learning?
Dive into the deep dive papers to understand the theory, explore the hands-on code examples, and prepare for ML infrastructure engineering roles.