COME: Adding Scene-Centric Forecasting Control to Occupancy World Model

Demo Videos

The comparison of ground-truth, DOME generation with official checkpoint and COME. The task setting is to use 4-frame 3D-Occ sequences as input and predict the next 6-frame (3-s prediction) sequences.

gt_dome_come_3s_results.mp4

The comparison of ground-truth, DOME generation with reproduced checkpoint and COME. The task setting is to use 4-frame 3D-Occ sequences as input and predict the next 16-frame (8-s prediction) sequences.

gt_dome_come_8s_results.mp4

The COME generation with BEV layouts. The task setting is to use 2-frame 3D-Occ sequences and 8-frame BEV layouts as input and predict the next 6-frame (3-s) sequences.

come_with_bev_layout_3s_results.mp4

Overview

COME = Forecasting Guided Generation

Results

🚀 Setup

environment setup

conda env create --file environment.yml
pip install einops tabulate 
cd occforecasting 
python setup.py develop
cd ..

data preparation

1. Create soft link from data/nuscenes to your_nuscenes_path

2. Prepare the gts semantic occupancy introduced in Occ3d

3. Download generated train/val pickle files from OccWorld or DOME.

4. Prepare the train/val pickle files for scene-centric forecasting.

python -m occforecasting.datasets.nusc_occ3d_dataset

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
    ├── nuscenes_train_occ3d_infos.pkl
    └── nuscenes_val_occ3d_infos.pkl

The four pickle files can also be downloaded in infos.

optinal inputs

For testing under different conditions, more inputs are needed.

1. motion planning results with yaw angles from BEVPlanner. Please put the json file on the project root directory. We simply add a yaw regression branch on BEV-Planner Project, Thanks for their great work.

2. BEV layouts for training and validation sets at 2Hz labels. Please unzip files and put them in './data/step2'. The pre-processing script is from UniScene, Thanks for their great work.

3. 3D occupancy prediction results from BEVDet and EFFOcc. Please unzip files and put them in './data/occpreds'. Thanks for their open-source checkpoints.

4. AE evaluation protocol from UniScene, please download AE_checkpoint for request and put in './ckpts/'.

Model Zoos

We recommend to download checkpoints with folders under './work_dir'.

Task Setting	Inputs	Method	Config	Checkpoint
Input-4frame-Output-6frame	3DOcc + GT Traj	Stage1-COME-World Model	Config	CKPT
Input-4frame-Output-6frame	3DOcc + GT Traj	Stage2-COME-Scene-Centric-Forecasting	Config	CKPT
Input-4frame-Output-6frame	3DOcc + GT Traj	Stage3-COME-ControlNet	Config	CKPT
Input-4frame-Output-6frame	3DOcc + Pred Traj	Stage3-COME-ControlNet	Config	Same As Above
Input-4frame-Output-6frame	BEVDet + Pred Traj	Stage3-COME-ControlNet	Config	Same As Above
Input-4frame-Output-6frame	BEVDet + GT Traj	Stage3-COME-ControlNet	Config	Same As Above
Input-4frame-Output-6frame	EFFOcc + Pred Traj	Stage3-COME-ControlNet	Config	Same As Above
Input-4frame-Output-6frame	EFFOcc + GT Traj	Stage3-COME-ControlNet	Config	Same As Above
Input-4frame-Output-16frame	3DOcc + GT Traj	Stage1-COME-World Model	Config	CKPT
Input-4frame-Output-16frame	3DOcc + GT Traj	Stage2-COME-Scene-Centric-Forecasting	Config	CKPT
Input-4frame-Output-16frame	3DOcc + GT Traj	Stage3-COME-ControlNet	Config	CKPT
Input-2frame-Output-6frame	3DOcc + GT Traj + BEV Layouts	Stage1-COME-World Model	Config	CKPT
Input-2frame-Output-6frame	3DOcc + GT Traj + BEV Layouts	Stage2-COME-Scene-Centric-Forecasting	Config	CKPT
Input-2frame-Output-6frame	3DOcc + GT Traj + BEV Layouts	Stage3-COME-ControlNet	Config	CKPT
Input-4frame-Output-6frame	3DOcc + GT Traj	Stage1-COME-Small-World Model	Config	CKPT
Input-4frame-Output-6frame	3DOcc + GT Traj	Stage2-COME-Scene-Centric-Forecasting	Config	Same As Above
Input-4frame-Output-6frame	3DOcc + GT Traj	Stage3-COME-Small-ControlNet	Config	CKPT

🏃 Run the code

OCC-VAE

By default, we use the VAE checkpoint provided by DOME, thanks for their greak work.

# train 
python tools/train_vae.py --py-config ./configs/train_occvae.py --work-dir ./work_dir/occ_vae 

# eval
python tools/train_vae.py --py-config ./configs/train_occvae.py --work-dir ./work_dir/occ_vae --resume-from ckpts/occvae_latest.pth

# visualize
python tools/visualize_demo_vae.py \
    --py-config ./configs/train_occvae.py \
    --work-dir ./work_dir/occ_vae \
    --resume-from ckpts/occvae_latest.pth \
    --export_pcd \
    --skip_gt

Scene-Centric Forecasting

cd occforecasting
# train 
bash train.sh occforecasting/configs/unet/unet_aligned_past2s_future_3s.py

# eval
bash test.sh occforecasting/configs/unet/unet_aligned_past2s_future_3s.py

COME World Model

# train 
python tools/train_diffusion.sh --py-config ./configs/train_dome_v2.py --work-dir ./work_dir/dome_v2

# eval
python tools/eval_metric.py --py-config ./configs/train_dome_v2.py --work-dir ./work_dir/dome_v2 --resume-from ./work_dir/dome_v2/best_miou.pth --vae-resume-from ckpts/occvae_latest.pth 


# visualize
python tools/visualize_demo.py --py-config ./configs/train_dome_v2.py --work-dir ./work_dir/dome_v2 --resume-from ./work_dir/dome_v2/best_miou.pth --vae-resume-from ckpts/occvae_latest.pth 

COME ControlNet

# train 
python tools/train_diffusion_control_ddp.py --py-config configs/train_dome_controlnet_mask_invisible_v2.py --work-dir work_dir/train_dome_controlnet_mask_invisible_v2 

# eval
python tools/test_diffusion_control.py --py-config configs/train_dome_controlnet_mask_invisible_v2.py --work-dir work_dir/train_dome_controlnet_mask_invisible_v2 


# visualize
python tools/visualize_demo_control_mask_invisible.py --py-config configs/train_dome_controlnet_mask_invisible_v2.py  --work-dir work_dir/train_dome_controlnet_mask_invisible_v2  --vae-resume-from ckpts/occvae_latest.pth  --skip_gt 

Acknoweldgements

This project is built on top of DOME and OccWorld. Thanks for the excellent open-source works!

DOME

OccWorld

ControlNet

HuyuanDiT