OmniV2X / docs /CHECKPOINTS.md
AndyPJT's picture
Correct DAIR-V2X release token cap
73ea6fa verified
|
Raw
History Blame Contribute Delete
3.73 kB

Checkpoints and Configs

Released checkpoints are hosted in the OmniV2X Hugging Face model repo.

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 omniv2x_flow MAP_TOKENS=0 INFRA_BBOX_TOKENS=16 NUM_INFERENCE_STEPS=20
OmniV2X map checkpoints/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:

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:

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:

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