docs/CHECKPOINTS.md

# Checkpoints and Configs

Released checkpoints are hosted in the
[OmniV2X Hugging Face model repo](https://huggingface.co/AndyPJT/OmniV2X).

```bash
mkdir -p checkpoints
hf download AndyPJT/OmniV2X \
  omniv2x_dairv2x_no_map.ckpt \
  omniv2x_dairv2x_map.ckpt \
  --local-dir checkpoints
```

## Recommended Files

Use these names for local checkpoints:

| Setting | Local filename | Agent config | Inference caps |
| --- | --- | --- | --- |
| OmniV2X no-map | [`checkpoints/omniv2x_dairv2x_no_map.ckpt`](https://huggingface.co/AndyPJT/OmniV2X/resolve/main/omniv2x_dairv2x_no_map.ckpt) | `omniv2x_flow` | `MAP_TOKENS=0 INFRA_BBOX_TOKENS=16 NUM_INFERENCE_STEPS=20` |
| OmniV2X map | [`checkpoints/omniv2x_dairv2x_map.ckpt`](https://huggingface.co/AndyPJT/OmniV2X/resolve/main/omniv2x_dairv2x_map.ckpt) | `omniv2x_map_flow` | `MAP_TOKENS=128 INFRA_BBOX_TOKENS=16 NUM_INFERENCE_STEPS=20` |

The released checkpoints are sanitized Lightning checkpoints for public
inference. They keep model weights and inference config, but omit optimizer,
scheduler, callback, and training loop state. EMA checkpoints can be archived
for ablations, but the reproducibility baseline uses raw `traj_ade=*.ckpt`
weights.

Each public checkpoint should include an `agent_config` entry in the Lightning
checkpoint. If you publish an older checkpoint without embedded config, keep the
matching Hydra file next to the run at `code/hydra/config.yaml`; otherwise
inference falls back to the public no-map defaults.

## Config Contract

Both DAIR-V2X checkpoints use:

- DINOv3 ViT-B/16 visual encoder with 448x448 inputs
- two front-camera frames per sample
- ray positional encoding with camera parameters
- flow matching with 100 training timesteps and 20 inference steps
- history trajectory ego conditioning with five ego tokens
- infrastructure object-token conditioning with model capacity for 32 objects;
  the public reproduction command caps this at 16 bbox tokens

The map checkpoint additionally enables map conditioning with 128 map tokens,
implemented as 32 lanes and 4 tokens per lane.

## Released Checkpoint Verification

The released checkpoints were verified on the DAIR-V2X validation split with
2,390 samples, seed 42, 20 flow-matching inference steps, and the inference caps
listed above. PDMS values are reported on a 0-100 scale.

| Setting | Samples | Avg. L2 (m) | 3s Avg. Collision (%) | PDMS |
| --- | ---: | ---: | ---: | ---: |
| OmniV2X no-map | 2,390 | 0.89 | 0.07 | 88.62 |
| OmniV2X map | 2,390 | 0.90 | 0.01 | 90.08 |

## Reproducing DAIR-V2X Inference

No-map:

```bash
export CHECKPOINT_PATH=checkpoints/omniv2x_dairv2x_no_map.ckpt
export DAIRV2X_DATA_ROOT=/path/to/dair-v2x-seq
MAP_TOKENS=0 INFRA_BBOX_TOKENS=16 NUM_INFERENCE_STEPS=20 \
  RUN_NAME=omniv2x_no_map_val \
  bash scripts/infer_dairv2x.sh
```

Map:

```bash
export CHECKPOINT_PATH=checkpoints/omniv2x_dairv2x_map.ckpt
export DAIRV2X_DATA_ROOT=/path/to/dair-v2x-seq
export DAIRV2X_MAP_ROOT=/path/to/dair-v2x-seq/maps
MAP_TOKENS=128 INFRA_BBOX_TOKENS=16 NUM_INFERENCE_STEPS=20 \
  RUN_NAME=omniv2x_map_val \
  bash scripts/infer_dairv2x.sh
```

The inference script defaults to `--device cuda`, matching the paper
reproduction setting. For CPU-only installation checks, append `--device cpu`
and use `--max_scenarios` for a small smoke test.

Fine-tuning defaults to the no-map agent. For the map-conditioned variant, set
`AGENT_CONFIG=omniv2x_map_flow` when launching `scripts/finetune_dairv2x.sh`.

Evaluate a saved trajectory pkl:

```bash
export TRAJECTORY_RESULTS=outputs/dairv2x_inference/omniv2x_no_map_val_seed42.pkl
export DAIRV2X_DATA_ROOT=/path/to/dair-v2x-seq
export DAIRV2X_MAP_ROOT=/path/to/dair-v2x-seq/maps
bash scripts/eval_dairv2x_pdms.sh
```