DiffAV

A physics-informed, RL-aligned evaluation engine for autonomous driving

Diffusion models for autonomous vehicles

Early Development — API Unstable

DiffAV is in early development and undergoing rapid iteration. Breaking changes are expected. Pin to specific commits if stability is required.

Overview

DiffAV generates and evaluates adversarial driving scenarios to stress-test autonomous vehicle (AV) stacks. It combines generative trajectory prediction with physics constraints and reinforcement learning alignment to produce realistic yet safety-critical counterfactual scenarios, evaluated against the Waymo Open Dataset Sim Agents Challenge (WOSAC) metrics.

Why DiffAV?

• Physics-Informed: Bicycle model kinematics constraints ensure generated trajectories are physically plausible — no teleporting vehicles or impossible accelerations
• RL-Aligned: Reinforcement learning fine-tuning steers generation toward safety-critical scenarios that expose AV stack weaknesses
• JAX-Native: Built entirely on JAX with Flax NNX modules for JIT compilation, automatic differentiation, and hardware acceleration
• WOSAC-Compatible: Evaluates against standard Waymo challenge metrics (ADE, FDE, collision rate, miss rate), computed as JAX-native Euclidean proxies rather than the official leaderboard implementation
• Modular Architecture: Clean protocol-based design with frozen dataclass configuration — easy to swap generative models, physics validators, and evaluation metrics

Status & Results

DiffAV is a working research scaffold: every pillar (diffusion world model, DPO/steering alignment, physics feasibility, occupancy flow, evaluation metrics) is built, and its core has been exercised on real Waymo Open Dataset data. It is not a leaderboard-tuned system — results are reported honestly.

• Trajectory prediction. A map-conditioned diffusion baseline trained on real WOD reaches minADE₆ ≈ 5.6 m on held-out validation (tracks_to_predict, WOMD 2 Hz). That is roughly 9× the ~0.6 m of full-scale WOSAC leaders: the model is over-dispersed (best-of-64 ≈ 2.5 m) and under-fit, trained on ~1.6% of WOMD. The map is load-bearing — zeroing the scene tokens degrades minADE ~8.6× (4.4 → 38 m).
• Adversarial steering (the differentiator). Test-time reward guidance on the frozen baseline steers a real scene's adversary from 7.30 m → 2.88 m from its victim while improving off-road feasibility (0.21 → 0.08) and holding WOSAC realism (0.673 → 0.653). A map-conditioned Diffusion-DPO fine-tune reproduces this at training time.
• Roadmap. Close the trajectory-quality gap (scale up, fix over-dispersion), then a persistent-WOSAC leaderboard submission and GRPO/R1-style RL fine-tuning (now standard in the winning recipe).

The map-conditioned model (MapConditionedTrajectoryModel) and its trainer are driven end-to-end by scripts/train_wod.py — warmup-cosine schedule, EMA, and held-out validation.

Design

Generative Trajectory Prediction

DiffAV models multi-agent future trajectories conditioned on scene context (ego state, surrounding agents, HD map features). The generative backbone is a diffusion model:

• Diffusion Models — iterative denoising for high-quality multi-modal trajectory distributions

The protocol-based design leaves room for other families (normalizing flows, score-based models); those are not yet implemented.

All models predict in the WOSAC state space: 11 history steps (1.1s) conditioning 80 future steps (8.0s) at 10Hz, with state dimension (x, y, heading, velocity).

Physics Constraints

Generated trajectories pass through a differentiable physics validation layer based on the bicycle kinematic model. Constraints include maximum acceleration, curvature bounds, and collision detection. Physics violations feed back as a loss term during training, ensuring the generator learns to produce plausible trajectories without post-hoc rejection sampling.

Adversarial Alignment

An RL fine-tuning stage (DPO) optimizes the generator to produce scenarios that are simultaneously realistic (high WOSAC scores) and challenging (expose planning failures). This closes the loop between generation quality and safety-critical scenario discovery.

Architecture

src/diffav/
  core/         # Domain types, protocols, and configuration
  data/         # WOD TFRecord loading and scene tokenization
  models/       # Trajectory prediction model wrappers (diffusion)
  physics/      # Bicycle model kinematics and constraint validation
  alignment/    # RL fine-tuning for adversarial scenario generation
  evaluation/   # WOSAC metrics (ADE, FDE, collision rate) and orchestration
  occupancy/    # Fourier Neural Operator occupancy flow prediction
  sensor/       # NeRF/LiDAR sensor simulation for perception testing
  api/          # Public SDK for scenario mining and evaluation

Sister Repositories

DiffAV builds on four companion libraries in the JAX ecosystem:

Installation

# Clone and set up the environment (auto-detects CUDA/Metal/CPU)
git clone https://github.com/avitai/DiffAV.git
cd DiffAV
./setup.sh

# Or manually with uv (dev tooling is installed by default)
uv sync                 # CPU
uv sync --extra gpu     # Linux with NVIDIA GPU (CUDA 12)

Requirements

• Python 3.11, 3.12, or 3.13
• JAX >= 0.6.1, < 0.12 (tested on 0.9.1)
• Flax >= 0.12.0
• TensorFlow >= 2.20.0 (CPU-only, for WOD proto parsing)
• Waymo Open Dataset access (requires license agreement)

Quick Start

from diffav.api import MinerConfig, create_scenario_miner

# Configure and create the scenario miner
config = MinerConfig(model_path="checkpoints/wod-mini")
miner = create_scenario_miner(config)

# Generate scenarios and evaluate a planner
scenarios = miner.generate("unprotected_left_turn", density="high", count=100)
report = miner.evaluate_planner(my_planner_fn, scenarios)
print(report.metric_values)  # {"ade": ..., "fde": ...}

Development

DiffAV uses uv as its package manager and enforces code quality via pre-commit hooks.

Running Tests

# Core unit tests
uv run pytest tests/core/ -v

# Full test suite with coverage
uv run pytest -v --cov=src/diffav --cov-report=term-missing

# Skip tests requiring WOD data
uv run pytest -v -m "not wod"

Code Quality

# All pre-commit hooks
uv run pre-commit run --all-files

# Individual checks
uv run ruff check src/          # Lint
uv run ruff format src/         # Format
uv run pyright src/             # Type check

License

DiffAV is licensed under the MIT License.