DOCSEARCHbyAlgolia
Submit Docs

DiffusionDriveV2

repository·master·Indexed 43 days ago

https://github.com/hustvl/diffusiondrivev2

A framework for end-to-end autonomous driving that uses reinforcement learning to constrain truncated diffusion models. It combines diffusion modeling with GRPO-based reinforcement learning to generate high-quality, diverse, and multimodal trajectories, achieving state-of-the-art results on the NAVSIM dataset.

Tokens
3.9K
Snippets
12
Records
15
Agent score
57%

What's inside DiffusionDriveV2

  1. Overview of DiffusionDriveV2

    master

    DiffusionDriveV2 is a reinforcement learning-constrained truncated diffusion modeling framework designed for end-to-end autonomous driving. It addresses the mode collapse and quality-diversity dilemma common in generative diffusion models by using reinforcement learning to constrain low-quality modes and explore superior trajectories.

    Key technical components include:

    • Scale-adaptive multiplicative noise: Promotes broad exploration for trajectory planning.
    • Intra-anchor GRPO: Manages advantage estimation among samples generated from a single anchor.
    • Inter-anchor truncated GRPO: Incorporates a global perspective across different anchors to prevent improper advantage comparisons between distinct driving intentions (e.g., turning vs. going straight).

    It achieves state-of-the-art results on the NAVSIM dataset (91.2 PDMS on v1 and 85.5 EPDMS on v2).

  2. Getting Started with DiffusionDriveV2

    master

    To use DiffusionDriveV2, you must follow a multi-step setup process involving the NAVSIM environment, the DiffusionDriveV2 specific environment, and then proceeding to training or evaluation.

    Follow these steps in order:

    1. Prepare the NAVSIM environment: Follow the official NAVSIM getting started guide.
    2. Prepare the DiffusionDriveV2 environment: Refer to docs/install.md for specific dependency and environment setup.
    3. Run Training and Evaluation: Refer to docs/train_eval.md for instructions on how to execute training scripts and evaluate models.

  4. Evaluate DiffusionDrive models

    master

    There are two ways to evaluate the models:

    1. Fast Evaluation (Recommended)

    Use run_pdm_score_fast.py for rapid evaluation (several times faster than the official script). This requires the metric and feature caches to be present.

    2. Official Evaluation

    Use run_pdm_score.py. Note that if using the official script, you may need to uncomment a specific line in navsim/agents/diffusiondrivev2/diffusiondrivev2_model_sel.py at line 1491.

    # Fast evaluation
python $NAVSIM_DEVKIT_ROOT/navsim/planning/script/run_pdm_score_fast.py \
        agent=diffusiondrivev2_sel_agent \
        experiment_name=diffusiondrivev2_agent_eval \
        train_test_split=navtest \
        agent.checkpoint_path=ckpts/diffusiondrivev2_sel.ckpt \
        +metric_cache_path="${NAVSIM_EXP_ROOT}/metric_cache/" \
        +test_cache_path="${NAVSIM_EXP_ROOT}/metric_feature_cache/"

# Official evaluation
python $NAVSIM_DEVKIT_ROOT/navsim/planning/script/run_pdm_score.py \
        train_test_split=navtest \
        agent=diffusiondrivev2_sel_agent \
        worker=ray_distributed \
        agent.checkpoint_path=ckpts/diffusiondrivev2_sel.ckpt \
        experiment_name=diffusiondrivev2_agent_eval

  5. Train DiffusionDrive with Reinforcement Learning

    master

    To train the DiffusionDrive agent using reinforcement learning, use the run_training.py script. You must specify the RL agent, the training split, and the path to the base DiffusionDrive checkpoint (diffusiondrive_navsim_88p1_PDMS).

    Key parameters:

    • agent=diffusiondrivev2_rl_agent
    • agent.checkpoint_path: Path to the base model checkpoint.
    • use_cache_without_dataset=True: Enables using the cached data.
    # Training DiffusionDrive with reinforcement learning
python $NAVSIM_DEVKIT_ROOT/navsim/planning/script/run_training.py \
        agent=diffusiondrivev2_rl_agent \
        experiment_name=training_diffusiondrivev2_rl_agent \
        train_test_split=navtrain \
        agent.checkpoint_path=ckpts/diffusiondrive_navsim_88p1_PDMS \
        split=trainval \
        trainer.params.max_epochs=10 \
        cache_path="${NAVSIM_EXP_ROOT}/training_cache/" \
        use_cache_without_dataset=True \
        force_cache_computation=False

  6. Cache datasets for DiffusionDrive training and evaluation

    master

    To speed up training and evaluation processes, you can pre-calculate and cache datasets. There are four specific caching tasks depending on your goal:

    1. Training Cache: Caches the dataset for training.
    2. Evaluation Metric Cache: Caches metrics for evaluation.
    3. PDMS Training Cache: Caches data specifically for calculating PDMS during training.
    4. Fast Evaluation Feature Cache: (Optional) Caches features to enable significantly faster evaluation.

    Ensure $NAVSIM_EXP_ROOT is defined in your environment to specify the cache destination.

    # cache dataset for training
python navsim/planning/script/run_dataset_caching.py agent=diffusiondrivev2_rl_agent experiment_name=diffusiondrivev2_cache train_test_split=navtrain

# cache dataset for evaluation
python navsim/planning/script/run_metric_caching.py train_test_split=navtest cache.cache_path=$NAVSIM_EXP_ROOT/metric_cache

# cache dataset for calculating PDMS during training.
python navsim/planning/script/run_metric_caching.py train_test_split=navtrain cache.cache_path=$NAVSIM_EXP_ROOT/train_pdm_cache

# cache dataset for fast evaluation (optional)
python navsim/planning/script/run_dataset_caching.py agent=diffusiondrivev2_rl_agent experiment_name=diffusiondrivev2_cache train_test_split=navtest cache_path=$NAVSIM_EXP_ROOT/metric_feature_cache

  7. Download GTRS trajectories for data augmentation

    master

    If you intend to use GTRS trajectories as data augmentation during mode selector training, you must download the simulated ground-truths for the GTRS vocabulary. Navigate to the gtrs_traj directory and use wget to fetch the required .pkl file.

    cd gtrs_traj
wget https://huggingface.co/Zzxxxxxxxx/gtrs/resolve/main/navtrain_16384.pkl

  8. Train the DiffusionDrive Mode Selector

    master

    The mode selector is trained on top of the RL-trained model. Use the diffusiondrivev2_sel_agent and point agent.checkpoint_path to the diffusiondrivev2_rl.ckpt checkpoint.

    # Training mode selector
python $NAVSIM_DEVKIT_ROOT/navsim/planning/script/run_training.py \
        agent=diffusiondrivev2_sel_agent \
        experiment_name=training_diffusiondrivev2_sel_agent \
        train_test_split=navtrain \
        agent.checkpoint_path=ckpts/diffusiondrivev2_rl.ckpt \
        split=trainval \
        trainer.params.max_epochs=20 \
        cache_path="${NAVSIM_EXP_ROOT}/training_cache/" \
        use_cache_without_dataset=True \
        force_cache_computation=False

  9. Download DiffusionDriveV2 assets using super_download.sh

    master

    To speed up the asset download process, you can use the super_download.sh script which utilizes tmux to parallelize downloads. Navigate to the download directory within the DiffusionDriveV2 repository and execute the script.

    cd /path/to/DiffusionDriveV2/download
bash super_download.sh

  10. Download required checkpoints

    master

    Before training or evaluation, download the following checkpoints and place them in the ckpts directory:

    • resnet34.a1_in1k (from timm)
    • diffusiondrive_navsim_88p1_PDMS (DiffusionDrive model)
    • diffusiondrivev2_rl.ckpt (The RL-trained DiffusionDrive model)
    • diffusiondrivev2_sel.ckpt (DiffusionDriveV2 RL + mode selector)

    Checkpoints can be found on Hugging Face at hustvl/DiffusionDriveV2 and hustvl/DiffusionDrive.

  11. Create custom plots using NAVSIM primitives

    master

    You can build custom matplotlib plots by combining low-level visualization functions. For example, you can manually add annotations and LiDAR to an axis and then use configure_bev_ax to set the correct Birds-Eye-View perspective.

    from navsim.visualization.plots import configure_bev_ax
from navsim.visualization.bev import add_annotations_to_bev_ax, add_lidar_to_bev_ax
import matplotlib.pyplot as plt

fig, ax = plt.subplots(1, 1, figsize=(6, 6))
ax.set_title("Custom plot")

# Add specific elements to the axis
add_annotations_to_bev_ax(ax, scene.frames[frame_idx].annotations)
add_lidar_to_bev_ax(ax, scene.frames[frame_idx].lidar)

# Apply BEV configuration (axes limits, etc.)
configure_bev_ax(ax)

plt.show()

  12. Initialize SceneLoader for visualization

    master

    To visualize driving scenes, you must first set up a SceneLoader. This requires an environment variable OPENSCENE_DATA_ROOT pointing to your data directory. You use a SceneFilter (often instantiated via Hydra) to define which scenes to load from a specific split (e.g., mini, test, or trainval).

    import os
from pathlib import Path
import hydra
from hydra.utils import instantiate
import numpy as np

from navsim.common.dataloader import SceneLoader
from navsim.common.dataclasses import SceneFilter, SensorConfig

SPLIT = "mini"  # ["mini", "test", "trainval"]
FILTER = "all_scenes"

hydra.initialize(config_path="../navsim/planning/script/config/common/scene_filter")
cfg = hydra.compose(config_name=FILTER)
scene_filter: SceneFilter = instantiate(cfg)
openscene_data_root = Path(os.getenv("OPENSCENE_DATA_ROOT"))

scene_loader = SceneLoader(
    openscene_data_root / f"navsim_logs/{SPLIT}",
    openscene_data_root / f"sensor_blobs/{SPLIT}",
    scene_filter,
    sensor_config=SensorConfig.build_all_sensors(),
)