Docs edits (#178)

* Simplify data downloading.

* Add links.

* Update WOSAC eval section.

* Minor.

* Rephrase.

* Fixes.

* Naming edits.

* Naming edits.

Author: daphne-cornelisse

docs/benchmark-readme.md

@@ -1,10 +1,10 @@
 # Data
 
-PufferDrive consumes map binaries generated from the Waymo Open Motion Dataset (WOMD) JSON files. This page covers where to obtain data and how to convert it into the binary format expected by the simulator.
+PufferDrive consumes map binaries generated from multiple data sources, including the [Waymo Open Motion Dataset (WOMD)](https://github.com/waymo-research/waymo-open-dataset) JSON files, [ScenarioMax](https://github.com/valeoai/V-Max), and [CARLA](https://carla.org/). This page covers how to obtain data and convert it into the binary format expected by the simulator.
 
 ## Download options
-- **Mini dataset**: [GPUDrive_mini](https://huggingface.co/datasets/EMERGE-lab/GPUDrive_mini) with ~1k training files for quick experiments.
-- **Full dataset**: [GPUDrive](https://huggingface.co/datasets/EMERGE-lab/GPUDrive) with ~100k scenes for large-scale training.
+- [`pufferdrive_womd_train`](https://huggingface.co/datasets/daphne-cornelisse/pufferdrive_womd_train): **10k scenarios** from the Waymo Open Motion _training_ dataset
+- [`pufferdrive_womd_val`](https://huggingface.co/datasets/daphne-cornelisse/pufferdrive_womd_val): **10k scenarios** from the Waymo Open Motion _validation_ dataset
 - Additional compatible sources: [ScenarioMax](https://github.com/valeoai/ScenarioMax) exports JSON in the same format.
 - Included CARLA maps: CARLA exports live in `data_utils/carla/carla`.
 
@@ -13,29 +13,19 @@ Install the CLI once:
 
 ```bash
 uv pip install -U "huggingface_hub[cli]"
-```
-
-Download the mini dataset:
+```https://github.com/valeoai/V-Max
 
+Download:
 ```bash
-huggingface-cli download EMERGE-lab/GPUDrive_mini \
+huggingface-cli download daphne-cornelisse/pufferdrive_womd_train \
   --repo-type dataset \
-  --local-dir data/processed/training \
-  --include "training/*"
-```
-
-Download the full dataset:
-
-```bash
-huggingface-cli download EMERGE-lab/GPUDrive \
-  --repo-type dataset \
-  --local-dir data/processed/training \
-  --include "training/*"
+  --local-dir data/processed/training
 ```
 
 Place raw JSON files under `data/processed/training` (default location read by the conversion script).
 
 ## Convert JSON to map binaries
+
 The conversion script writes compact `.bin` maps to `resources/drive/binaries`:
 
 ```bash
@@ -64,5 +54,6 @@ Refer to [Simulator](simulator.md) for how the binaries are consumed during rese
 - With binaries in place, run `puffer train puffer_drive` from [Getting Started](getting-started.md) as a smoke test that the build, data, and bindings are wired together.
 - To inspect the binary output, convert a single JSON file with `load_map(<json>, <id>, <output_path>)` inside `drive.py`.
 
-## Waymo map editor
-See [WOMD Editor](womd-editor.md) for a browser-based workflow to inspect, edit, and export Waymo/ScenarioMax JSON into the `.bin` format consumed by the simulator.
+## Interactive scenario editor
+
+See [Interactive scenario editor](scene-editor.md) for a browser-based workflow to inspect, edit, and export Waymo/ScenarioMax JSON into the `.bin` format consumed by the simulator.

docs/data.md

@@ -1,8 +1,9 @@
-# Evaluation
+# Evaluations and benchmarks
 
-Benchmarks are provided to measure how closely agents match real-world driving behavior and how well they pair with human trajectories.
+Driving is a safety-critical multi-agent application, making careful evaluation and risk assessment essential. Mistakes in the real world are costly, so simulations are used to catch errors before deployment. To support rapid iteration, evaluations should ideally run efficiently. This is why we also paid attention to optimizing the speed of the evaluations. This page contains an overview of the available benchmarks and evals.
+
+## Sanity maps 🐛
 
-## Sanity maps
 Quickly test the training on curated, lightweight scenarios without downloading the full dataset. Each sanity map tests a specific behavior.
 
 ```bash
@@ -15,9 +16,10 @@ Or run them all at once:
 puffer sanity puffer_drive --wandb --wandb-name sanity-all
 ```
 
-- Tip:Turn learning-rate annealing off for these short runs (`--train.anneal_lr False`) to keep the sanity checks from decaying the optimizer mid-run.
+- Tip: turn learning-rate annealing off for these short runs (`--train.anneal_lr False`) to keep the sanity checks from decaying the optimizer mid-run.
 
 Available maps:
+
 - `forward_goal_in_front`: Straight approach to a goal in view.
 - `reverse_goal_behind`: Backward start with a behind-the-ego goal.
 - `two_agent_forward_goal_in_front`: Two agents advancing to forward goals.
@@ -31,21 +33,24 @@ Available maps:
 
 ![Sanity map gallery placeholder](images/maps_screenshot.png)
 
-## WOSAC distributional realism
-Evaluate how realistic your policy behaves compared to the Waymo Open Sim Agents Challenge (WOSAC):
+## Distributional realism benchmark 📊
+
+We provide a PufferDrive implementation of the Waymo Open Sim Agents Challenge (WOSAC) for fast, easy evaluation of how well your trained agent matches distributional properties of human behavior.
 
 ```bash
 puffer eval puffer_drive --eval.wosac-realism-eval True
 ```
 
-Add `--load-model-path <path_to_checkpoint>.pt` to score a trained policy instead of a random baseline.
-See [WOSAC Benchmark](benchmark-readme.md) for the metric pipeline and links to the official configs.
+Add `--load-model-path <path_to_checkpoint>.pt` to score a trained policy, instead of a random baseline.
+
+See [the WOSAC benchmark page](wosac.md) for the metric pipeline and all the details.
+
+## Human-compatibility benchmark 🤝
 
-## Human-compatibility
-Test how a policy coexists with human-controlled agents:
+You may be interested in how compatible your agent is with human partners. For this purpose, we support an eval where your policy only controls the self-driving car (SDC). The rest of the agents in the scene are stepped using the logs. While it is not a perfect eval since the human partners here are static, it will still give you a sense of how closely aligned your agent's behavior is to how people drive. You can run it like this:
 
 ```bash
 puffer eval puffer_drive --eval.human-replay-eval True --load-model-path <path_to_checkpoint>.pt
 ```
 
-During this evaluation the self-driving car (SDC) is controlled by your policy while other agents replay log data.
+During this evaluation the self-driving car (SDC) is controlled by your policy while other agents replay log trajectories.

docs/evaluation.md

@@ -12,12 +12,14 @@
 </div>
 
 ## Highlights
-- Multi-agent drive environment with batched stepping for speed.
+
+- Multi-agent drive environment that trains agents at 300K steps per second.
 - Scripts to convert Waymo Open Motion Dataset JSON into lightweight binaries (CARLA sample map included).
 - Benchmarks for distributional realism and human compatibility.
 - Raylib-based visualizer for local or headless render/export.
 
 ## Quick start
+
 - Follow [Getting Started](getting-started.md) to install, build the C extensions, and run `puffer train puffer_drive`.
 - Consult [Simulator](simulator.md) for how actions/observations, rewards, and `.ini` settings map to the underlying C environment and Torch policy.
 - Prepare drive map binaries with the steps in [Data](data.md).
@@ -41,7 +43,7 @@
   <div class="step-card">
     <div class="badge">Step 3</div>
     <h3>Train & Evaluate</h3>
-    <p>Run `puffer train puffer_drive` and score policies with WOSAC and human-replay benchmarks.</p>
+    <p>Train agents and evaluate them with WOSAC and human-replay benchmarks.</p>
     <a href="evaluation/">Open guide</a>
   </div>
 </div>

docs/index.md

@@ -1,4 +1,4 @@
-# WOMD Editor
+# Interactive scenario editor
 
 A browser-based playground for inspecting and editing Waymo Open Motion Dataset (WOMD) scenes. The tool runs fully client-side at <https://womd-editor.vercel.app/> and works directly with the JSON format produced by Waymo/ScenarioMax exports and PufferDrive conversions.
 

docs/womd-editor.md docs/scene-editor.mddocs/womd-editor.md renamed to docs/scene-editor.md

@@ -0,0 +1,71 @@
+# Waymo Open Sim Agent Challenge (WOSAC) Benchmark
+
+We provide a re-implementation of the [Waymo Open Sim Agent Challenge (WOSAC)](https://waymo.com/research/the-waymo-open-sim-agents-challenge/), which measures _distributional realism_ of simulated trajectories compared to logged human trajectories. Our version preserves the original logic and metric weighting but uses PyTorch on GPU for the metrics computation, unlike the original TensorFlow CPU implementation. The code is also simplified for clarity, making it easier to understand, adapt, and extend.
+
+**Note:** In PufferDrive, agents are conditioned on a "goal" represented as a single (x, y) position, reflecting that drivers typically have a high-level destination in mind. Evaluating whether an agent matches human distributional properties can be decomposed into: (1) inferring a person's intended direction from context (1 second in WOSAC) and (2) navigating toward that goal in a human-like manner. We focus on the second component, though the evaluation could be adapted to include behavior prediction as in the original WOSAC.
+
+[TODO: ADD bar graphs]
+
+## Usage
+
+### Running a single evaluation from a checkpoint
+
+The `[eval]` section in `drive.ini` contains all relevant configurations. To run the WOSAC eval once:
+
+```bash
+puffer eval puffer_drive --eval.wosac-realism-eval True --load-model-path <your-trained-policy>.pt
+```
+
+The default configs aim to emulate the WOSAC settings as closely as possible, but you can adjust them:
+
+```ini
+wosac_num_rollouts = 32      # Number of policy rollouts per scene
+wosac_init_steps = 10        # When to start the simulation
+wosac_num_agents = 256       # Total number of WOSAC agents to evaluate
+wosac_control_mode = "control_wosac"  # Control the tracks to predict
+wosac_init_mode = "create_all_valid"  # Initialize from the tracks to predict
+wosac_goal_behavior = 2      # Stop when reaching the goal
+wosac_goal_radius = 2.0      # Can shrink goal radius for WOSAC evaluation
+```
+
+### Log evals to W&B during training
+
+During experimentation, logging key metrics directly to W&B avoids a post-training step. Evaluations can be enabled during training, with results logged under a separate `eval/` section. The main configuration options:
+
+```ini
+[train]
+checkpoint_interval = 500    # Set equal to eval_interval to use the latest checkpoint
+
+[eval]
+eval_interval = 500          # Run eval every N epochs
+map_dir = "resources/drive/binaries/training"  # Dataset to use
+```
+
+## Baselines
+
+We provide baselines on a small curated dataset from the WOMD validation set with perfect ground-truth (no collisions or off-road events from labeling mistakes).
+
+| Method | Realism meta-score | Kinematic metrics | Interactive metrics | Map-based metrics | minADE | ADE |
+|--------|-------------------|-------------------|---------------------|-------------------|--------|------|
+| Ground-truth (UB) | 0.832 | 0.606 | 0.846 | 0.961 | 0 | 0 |
+| π_Base self-play RL | - | - | - | - | - | - |
+| [SMART](https://arxiv.org/abs/2405.15677) | 0.805 | 0.534 | 0.830 | 0.949 | 1.124 | 3.123 |
+| π_Random | 0.485 | 0.214 | 0.657 | 0.408 | 6.477 | 18.286 |
+
+*Table: WOSAC baselines in PufferDrive on 229 selected clean held-out validation scenarios.*
+
+---
+
+> ✏️ Download the dataset from [Hugging Face](https://huggingface.co/datasets/daphne-cornelisse/pufferdrive_wosac_val_clean) to reproduce these results or benchmark your policy.
+
+---
+
+## Useful links
+
+- [WOSAC challenge and leaderboard](https://waymo.com/open/challenges/2025/sim-agents/)
+- [Sim agent challenge tutorial](https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_sim_agents.ipynb)
+- [Reference paper introducing WOSAC](https://arxiv.org/pdf/2305.12032)
+- [Metrics entry point](https://github.com/waymo-research/waymo-open-dataset/blob/master/src/waymo_open_dataset/wdl_limited/sim_agents_metrics/metrics.py)
+- [Log-likelihood estimators](https://github.com/waymo-research/waymo-open-dataset/blob/master/src/waymo_open_dataset/wdl_limited/sim_agents_metrics/estimators.py)
+- [Configurations proto file](https://github.com/waymo-research/waymo-open-dataset/blob/99a4cb3ff07e2fe06c2ce73da001f850f628e45a/src/waymo_open_dataset/protos/sim_agents_metrics.proto#L51)
+- [Default sim agent challenge configs](https://github.com/waymo-research/waymo-open-dataset/blob/master/src/waymo_open_dataset/wdl_limited/sim_agents_metrics/challenge_2025_sim_agents_config.textproto)

docs/wosac.md

@@ -28,12 +28,12 @@ nav:
   - Docs:
       - Getting started: getting-started.md
       - Simulator: simulator.md
-      - Interactive scenario editor: womd-editor.md
+      - Interactive scenario editor: scene-editor.md
       - Visualizer: visualizer.md
   - Data: data.md
   - Benchmarks:
       - Overview: evaluation.md
-      - WOSAC: benchmark-readme.md
+      - WOSAC: wosac.md
   - Blog:
     - pufferdrive 2.0: pufferdrive-2.0.md