Ranked by Relevance

Project Ideas

Hands-on projects to demonstrate your skills with Waymax and JAX. Organized by relevance to ML infrastructure and sim agents work.

Tier 1

High Relevance

Do these first. These projects directly align with sim agents and ML infrastructure work.

Tier 1 ProjectAdvanced

Sim Agent Policy with WOSAC Metrics

What: Train a behavior model that controls all non-ego agents
Why: This is exactly what the sim agents team works on
Deliverable: Model that achieves competitive WOSAC metrics
JAXBehavior ModelingMetrics Computation
Tier 1 ProjectIntermediate

Batch Simulation Profiler

What: Build tooling to profile and optimize Waymax simulation speed
Why: ML infra engineers care about throughput and efficiency
Deliverable: Dashboard showing steps/second, memory usage, bottlenecks
Performance EngineeringJAX Profiling
Tier 1 ProjectAdvanced

Distributed Training Pipeline

What: Scale training across multiple GPUs with gradient synchronization
Why: This is core ML infrastructure work
Deliverable: Training script using JAX pmap that scales near-linearly
Distributed SystemsJAX pmapCollective Operations
Tier 2

Strong Relevance

These projects demonstrate valuable skills and understanding of simulation challenges.

Tier 2 ProjectIntermediate

Custom Metrics Implementation

What: Implement additional realism metrics beyond built-in ones
Why: Metrics design is crucial for simulation validity
Ideas: Speed limit compliance, Lane-following quality, Interaction complexity
Metrics DesignJAX
Tier 2 ProjectIntermediate

Adversarial Scenario Mining

What: Find scenarios that cause policy failures
Why: Long-tail testing is a key challenge
Deliverable: Pipeline that identifies high-failure-rate scenarios
TestingScenario AnalysisML Evaluation
Tier 2 ProjectIntermediate

Model Serving Benchmark

What: Benchmark inference latency of different policy architectures
Why: Serving ML models in sim needs to be fast
Deliverable: Comparison of MLP vs Transformer inference times
BenchmarkingModel OptimizationJAX
Tier 3

Good Learning Exercises

Great for building foundational skills and understanding the ecosystem.

Tier 3 ProjectBeginner

Visualization Dashboard

What: Interactive visualization of simulation rollouts
Why: Debugging and understanding agent behavior
Tools: Matplotlib, Plotly, Web-based
Data VisualizationFrontend Development
Tier 3 ProjectIntermediate

Dataset Converter

What: Convert other datasets (nuPlan, Argoverse) to Waymax format
Why: Multi-dataset support (like ScenarioMax in V-Max)
Data EngineeringDataset Formats

Tips for Success

Make the most of these projects to build your portfolio

Start Small

Begin with Tier 3 projects to build familiarity before tackling advanced ones.

Measure Everything

Include benchmarks and metrics in your deliverables to demonstrate rigor.

Write Clean Code

Follow JAX best practices: pure functions, explicit state, type annotations.

Document Well

Write clear READMEs explaining your approach, results, and learnings.

Ready to Build?

Start with the deep dive papers to build foundational knowledge, then pick a project that interests you.

Read Waymax Deep DiveBrowse All Papers