Ma-Zhuang/OmniNWM

📚 What is OmniNWM?

OmniNWM (Omniscient Driving Navigation World Models) is a research‑grade codebase that builds a panoramic world model for autonomous‑driving simulation. The model learns to generate, from a vehicle’s planned trajectory, a full set of sensor‑level outputs – RGB images, semantic maps, depth maps, and a 3‑D occupancy grid – for all six surrounding cameras. Because the generated world is dense and multi‑modal, it can be used both to render realistic driving videos and to provide a closed‑loop environment for training or evaluating driving policies.

🎯 Core ideas

Idea	What it means
Multi‑modal generation	A single neural network predicts RGB, semantics, depth and 3‑D occupancy for a full 360° view.
Normalized Plücker ray‑maps	A geometric representation that lets the model translate pixel‑level predictions into precise vehicle actions.
Auto‑regressive stability	A “forcing” strategy lets the model keep generating beyond the length of ground‑truth videos without drifting.
Occupancy‑based dense rewards	The predicted occupancy grid can be turned into a reward signal, enabling realistic closed‑loop policy evaluation.
Zero‑shot transfer	The same checkpoints work on other driving datasets (e.g., nuPlan) and on different camera rigs without extra fine‑tuning.

🛠️ Getting started (quick‑start)

Clone & set up

git clone https://github.com/Ma-Zhuang/OmniNWM.git && cd OmniNWM
mkdir -p pretrained data
pip install -e .
pip install "huggingface_hub[cli]"

Patch transformers – edit modeling_utils.py as instructed (change the torch version check from 2.3 to 2.5).

Download checkpoints

huggingface-cli download Arlolo0/OmniNWM --local-dir ./pretrained
huggingface-cli download hpcai-tech/Open-Sora-v2 --local-dir ./pretrained

Prepare data – download the nuScenes v1.0 train/val splits and the 12 Hz depth/segmentation annotations from the links in the README. Follow the directory layout shown in the repo (e.g., data/nuscenes/CAM_FRONT, data/nuscenes_12hz_depth_unzip, etc.).

🚀 Typical workflows

Task	Command	Notes
Inference – trajectory‑to‑video	`torchrun --nproc-per-node 8 tools/inference.py configs/inference/infer.py`	Generates a 33‑frame video for all six cameras (448×800).
Out‑of‑distribution (nuPlan) inference	`torchrun --nproc-per-node 8 tools/inference.py configs/inference/infer_nuplan.py`	Uses a manually supplied trajectory on the nuPlan dataset.
Closed‑loop VLA test	`torchrun --nproc-per-node 8 tools/inference.py configs/inference/infer_with_occ_vla.py`	Runs a loop where the model’s occupancy prediction feeds back as a reward (321 frames).
Training – staged	`bash dist_train_mlp.sh configs/train/stage_1.py` (then stage 2, stage 3)	Stages increase resolution and video length to keep training stable.

📦 What’s inside the repo?

omninwm/models/OmniNWM‑VLA – implementation of the Tri‑MMI tri‑modal fusion and the VLA (Vision‑Language‑Action) pipeline.
configs/ – ready‑to‑run YAML/py configs for inference and the three training stages.
tools/ – scripts for distributed inference (inference.py).
pretrained/ – placeholder for the downloaded checkpoints (VAE, occupancy model, Open‑Sora‑v2 weights).
data/ – expected layout for nuScenes images, depth maps, and segmentation masks.

📖 Citation

If you use OmniNWM in research, cite the arXiv paper:

@article{li2025omninwm,
  title={OmniNWM: Omniscient Driving Navigation World Models},
  author={Li, Bohan and Ma, Zhuang and Du, Dalong and Peng, Baorui and Liang, Zhujin and Liu, Zhenqiang and Ma, Chao and Jin, Yueming and Zhao, Hao and Zeng, Wenjun and others},
  journal={arXiv preprint arXiv:2510.18313},
  year={2025}
}

⚖️ License

Apache License 2.0 (see LICENSE).

In short: OmniNWM is a cutting‑edge world‑model that turns a driving plan into a full‑fidelity, multi‑sensor simulation, enabling both realistic video synthesis and closed‑loop policy testing for autonomous‑driving research.