kicgit/occworld

3D World Model for Autonomous Driving

OccWorld

Paper | Project Page

OccWorld: Learning a 3D Occupancy World Model for Autonomous Driving

Wenzhao Zheng*, Weiliang Chen*, Yuanhui Huang, Borui Zhang, Yueqi Duan, Jiwen Lu

* Equal contribution

OccWorld models the joint evolutions of 3D scenes and ego movements.

Combined with self-supervised (SelfOcc), LiDAR-collected (TPVFormer), or machine-annotated (SurroundOcc) 3D occupancy, OccWorld has the potential to scale up to large-scale training, paving the way for interpretable end-to-end large driving models.

Demo

Overview

Given past 3D occupancy observations, our self-supervised OccWorld trained can forecast future scene evolutions and ego movements jointly. This task requires a spatial understanding of the 3D scene and temporal modeling of how driving scenarios develop. We observe that OccWorld can successfully forecast the movements of surrounding agents and future map elements such as drivable areas. OccWorld even generates more reasonable drivable areas than the ground truth, demonstrating its ability to understand the scene rather than memorizing training data. Still, it fails to forecast new vehicles entering the sight, which is difficult given their absence in the inputs.

Installation

1. Create conda environment with python version 3.8.0

2. Install all the packages in environment.yaml

3. Anything about the installation of mmdetection3d, please refer to mmdetection3d

Preparing

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:

OccWorld/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

Getting Started

Training

Train the OccWorld on RTX 4090 with 24G GPU memory (a VQVAE should be trained using similar command before training).

python train.py --py-config config/occworld.py --work-dir out/occworld

Evaluation

Eval the model on RTX 4090 with 24G GPU memory.

python eval_metric_stp3.py --py-config config/occworld.py --work-dir out/occworld

Related Projects

Our code is based on TPVFormer, SelfOcc, and PointOcc.

Also thanks to these excellent open-sourced repos: SurroundOcc OccFormer BEVFormer

Citation

If you find this project helpful, please consider citing the following paper:

@article{zheng2023occworld,
    title={OccWorld: Learning a 3D Occupancy World Model for Autonomous Driving},
    author={Zheng, Wenzhao and Chen, Weiliang and Huang, Yuanhui and Zhang, Borui and Duan, Yueqi and Lu, Jiwen },
    journal={arXiv preprint arXiv: 2311.16038},
    year={2023}
}