DOME: Taming Diffusion Model into High-Fidelity Controllable Occupancy World Model

Project Page | Paper

output_video.mp4

output_video.mp4

output_video.mp4

Our Occupancy World Model can generate long-duration occupancy forecasts and can be effectively controlled by trajectory conditions.

📖 Overview

Our method consists of two components: (a) Occ-VAE Pipeline encodes occupancy frames into a continuous latent space, enabling efficient data compression. (b)DOME Pipeline learns to predict 4D occupancy based on historical occupancy observations.

🗓️ News

• [2025.1.1] We release the code and checkpoints.
• [2024.11.18] Project page is online!

🗓️ TODO

• Code release.
• Checkpoint release.

🚀 Setup

clone the repo

git clone https://github.com/gusongen/DOME.git
cd DOME

environment setup

conda env create --file environment.yml

data preparation

1. Create soft link from data/nuscenes to your_nuscenes_path

2. Prepare the gts semantic occupancy introduced in Occ3d

3. Download our generated train/val pickle files and put them in data/

   nuscenes_infos_train_temporal_v3_scene.pkl

   nuscenes_infos_val_temporal_v3_scene.pkl

The dataset should be organized as follows:

.
└── data/
    ├── nuscenes            # downloaded from www.nuscenes.org/
    │   ├── lidarseg
    │   ├── maps
    │   ├── samples
    │   ├── sweeps
    │   ├── v1.0-trainval
    │   └── gts             # download from Occ3d
    ├── nuscenes_infos_train_temporal_v3_scene.pkl
    └── nuscenes_infos_val_temporal_v3_scene.pkl

ckpt preparation

Download the pretrained weights from here and put them in ckpts folder.

🏃 Run the code

(optional) Preprocess resampled data

cd resample

python launch.py \
    --dst ../data/resampled_occ \
    --imageset ../data/nuscenes_infos_train_temporal_v3_scene.pkl \
    --data_path ../data/nuscenes

OCC-VAE

# train 
sh tools/train_vae.sh

# eval
sh tools/eval_vae.sh

# visualize
sh tools/vis_vae.sh

DOME

# train 
sh tools/train_diffusion.sh 

# eval
sh tools/eval.sh 

# visualize
sh tools/vis_diffusion.sh

🎫 Acknowledgment

This code draws inspiration from their work. We sincerely appreciate their excellent contribution.

• OccWorld
• Latte
• Vista
• PyTorch-VAE
• A* serach

🖊️ Citation

@article{gu2024dome,
  title={Dome: Taming diffusion model into high-fidelity controllable occupancy world model},
  author={Gu, Songen and Yin, Wei and Jin, Bu and Guo, Xiaoyang and Wang, Junming and Li, Haodong and Zhang, Qian and Long, Xiaoxiao},
  journal={arXiv preprint arXiv:2410.10429},
  year={2024}
}