๐ŸŒ Nano World Model

License: MIT
A minimalist repository for training video world models based on diffusion-forcing.
![3ร—4 rollout grid](assets/grid_video.gif)
## News and Updates - ๐Ÿšฉ **2026-05:** We release [๐Ÿ“• Technical Report](https://arxiv.org/abs/2605.23993) on arxiv, check it out! - ๐Ÿšฉ **2026-05:** We release all model checkpoints on [๐Ÿค— huggingface](https://huggingface.co/collections/knightnemo/nano-world-model). - ๐Ÿšฉ **2026-05:** We release code for training, inference and evaluation. ## Key Features - ๐Ÿš€ **Instant Start** โ€” Minimal dependencies, easy data loading. From clone to first rollout in minutes. - ๐Ÿ› ๏ธ **Unified Pipeline** โ€” Training, Validation, Evaluation; All managed with clean hydra-based configuration systems. - ๐Ÿ”ฌ **Scientific Transparency** โ€” Clean codebase with head-to-head ablations across prediction target, action injection, and model scale; Fully open-source, including model checkpoints. - ๐Ÿค– **Diverse Applications** โ€” Long-horizon rollouts, rollout to 3d point clouds, planning (MPC) out of the box. ## ๐Ÿš€ Quick Start Try the [Colab quickstart notebook](colab_quickstart.ipynb) to download the official Point Maze checkpoint and preview predictions beside recorded validation frames. Select a GPU runtime, then run all cells; the notebook uses the same lock below. Install [uv](https://docs.astral.sh/uv/getting-started/installation/) first. The supported Linux x86_64 environment uses one committed lock for Python, PyTorch, CUDA wheels, and all dependencies. See [environment setup and verification](docs/environment.md). ```bash git clone https://github.com/simchowitzlabpublic/nano-world-model.git cd nano-world-model ./nanowm sync ./nanowm doctor --require-cuda ``` For the locked RT-1/LeRobot dependencies, run `./nanowm sync --extra rt1`. Do not install `lerobot-datasets`; it is a dataset format version, not a PyPI package. Set data + results paths (or use the gitignored `src/configs/local/paths.yaml` template โ€” see [docs/config_system.md](docs/config_system.md#path-configuration)): ```bash export DATASET_DIR=/path/to/dino_wm_data # DINO-WM envs (point_maze, pusht, ...) export CSGO_DATA_DIR=/path/to/csgo # CSGO HDF5 files export RT1_DATA_ROOT=/path/to/rt1_fractal # RT-1 LeRobot mirror (optional) export RESULTS_DIR=/path/to/results # checkpoints + logs land here ``` Download the i3d torchscript used by FVD evaluation (FID uses InceptionV3, which downloads itself on first use): ```bash mkdir -p pretrained_models/i3d && curl -L \ "https://www.dropbox.com/scl/fi/c5nfs6c422nlpj880jbmh/i3d_torchscript.pt?rlkey=x5xcjsrz0818i4qxyoglp5bb8&dl=1" \ -o pretrained_models/i3d/i3d_torchscript.pt ``` For dataset downloads (DINO-WM, RT-1, CSGO), see [docs/datasets/README.md](docs/datasets/README.md). ## ๐Ÿฅท Train your first model DINO-WM PushT, NanoWM-B/2, default settings (pred-v ยท additive injection ยท cosine + ZTSNR): ```bash ./nanowm python src/main.py experiment=dino_wm_pusht dataset=dino_wm/pusht model=nanowm_b2 ``` CSGO with the L/2 model: ```bash ./nanowm python src/main.py experiment=csgo dataset=game/csgo model=nanowm_l2_csgo ``` RT-1 (fractal) main run: ```bash ./nanowm python src/main.py experiment=rt1 dataset=rt1/rt1 model=nanowm_b2 ``` For reproducibility, we provide example scripts in `src/scripts/`. See [docs/training.md](docs/training.md) for the full training guide, design choices, and ablation tables. ## ๐Ÿ“ฆ Pretrained Checkpoints Best-config runs (pred-v ยท additive ยท cosine + ZTSNR ยท NanoWM-B/2 unless noted):
| Domain | Checkpoint | Steps | |:-------|:-----------|:------| | DINO-WM Point Maze | ๐Ÿค— [nanowm-b2-dino-wm-point-maze-30k](https://huggingface.co/knightnemo/nanowm-b2-dino-wm-point-maze-30k) | 30k | | DINO-WM Wall | ๐Ÿค— [nanowm-b2-dino-wm-wall-15k](https://huggingface.co/knightnemo/nanowm-b2-dino-wm-wall-15k) | 15k | | DINO-WM Rope | ๐Ÿค— [nanowm-b2-dino-wm-rope-15k](https://huggingface.co/knightnemo/nanowm-b2-dino-wm-rope-15k) | 15k | | DINO-WM Granular | ๐Ÿค— [nanowm-b2-dino-wm-granular-15k](https://huggingface.co/knightnemo/nanowm-b2-dino-wm-granular-15k) | 15k | | DINO-WM PushT | ๐Ÿค— [nanowm-b2-dino-wm-pusht-100k](https://huggingface.co/knightnemo/nanowm-b2-dino-wm-pusht-100k) | 100k | | RT-1 (fractal) | ๐Ÿค— [nanowm-b2-rt1-300k](https://huggingface.co/knightnemo/nanowm-b2-rt1-300k) | 300k | | CSGO | ๐Ÿค— [nanowm-l2-csgo-100k](https://huggingface.co/knightnemo/nanowm-l2-csgo-100k) (NanoWM-L/2) | 100k |
We also provide RT-1 ablation tables with HF checkpoint paths. See [docs/training.md#design-choices](docs/training.md#design-choices) for the full table and ablation numbers. ## ๐ŸŽฌ Sample Predictions CSGO 50-frame auto-regressive long-rollouts (NanoWM-L/2, 100k):
![CSGO 50-frame autoregressive long rollout](assets/csgo_100k_long_rollout.gif)
Quantitative Metrics Evaluated on 256 fixed samples (seed=42), 250 DDIM steps, sequential scheduling (frame-by-frame autoregressive denoising).
| Dataset | Steps | PSNR โ†‘ | SSIM โ†‘ | LPIPS โ†“ | FID โ†“ | |:--------|:------|:-------|:-------|:--------|:------| | Point Maze | 30k | 36.74 | 0.984 | 0.019 | 9.66 | | Wall | 15k | 34.05 | 0.994 | 0.010 | 2.64 | | PushT | 100k | 33.19 | 0.982 | 0.016 | 13.63 | | Rope | 15k | 31.63 | 0.953 | 0.056 | 35.20 | | Granular | 15k | 26.08 | 0.917 | 0.073 | 40.05 | | RT-1 | 300k | 24.36 | 0.787 | 0.180 | 35.08 |
Full per-domain numbers and methodology in [docs/evaluation.md](docs/evaluation.md). ## ๐Ÿงญ Applications NanoWM rollouts can be used directly for downstream applications, including long-horizon generation, video-to-3D reconstruction, and MPC-style planning.
![Video-to-3D point cloud demo](assets/video_to_3d.gif)
- **[Long-horizon rollout](docs/applications/long_rollout.md)** โ€” autoregressive rollout from trained checkpoints - **[Video โ†’ 3D map](docs/applications/video_to_3d.md)** โ€” Depth Anything 3 point cloud reconstruction from rollout videos - **[MPC-style planning](docs/applications/planning.md)** โ€” CEM planning over world model rollouts ## ๐Ÿ“š Documentation - **[docs/config_system.md](docs/config_system.md)** โ€” Hydra config layout, overrides, environment variables - **[docs/training.md](docs/training.md)** โ€” training workflow, design choices, ablation tables, all checkpoints - **[docs/evaluation.md](docs/evaluation.md)** โ€” evaluation workflow, metric definitions, full result tables - **[docs/datasets/README.md](docs/datasets/README.md)** โ€” DINO-WM / RT-1 / CSGO formats, downloads, splits - **[docs/applications/planning.md](docs/applications/planning.md)** โ€” MPC + CEM model-predictive control - **[docs/applications/long_rollout.md](docs/applications/long_rollout.md)** โ€” long-horizon autoregressive rollout - **[docs/applications/video_to_3d.md](docs/applications/video_to_3d.md)** โ€” Depth Anything 3 point cloud pipeline ## ๐Ÿ™ Acknowledgements We build upon a number of existing codebases: [Latte](https://github.com/Vchitect/Latte), [Vid2World](https://github.com/thuml/Vid2World), [DFoT](https://github.com/kwsong0113/diffusion-forcing-transformer), and [DINO-WM](https://github.com/gaoyuezhou/dino_wm). More broadly, this repository draws inspirations and design principles from [NanoGPT](https://github.com/karpathy/nanoGPT), [NanoChat](https://github.com/karpathy/nanochat), and [Boyuan Chen's Research Template](https://github.com/buoyancy99/research-template). We sincerely thank the codebases above for open-sourcing their works. ## ๐Ÿ“ Citation If you find this repository useful in your research, please consider citing: ```bibtex @misc{huang2026nanoworldmodels, title={Nano World Models: A Minimalist Implementation of Future Video Prediction}, author={Siqiao Huang and Partha Kaushik and Michael Chen and Hengkai Pan and Omar Chehab and Fernando Moreno-Pino and Max Simchowitz}, year={2026}, eprint={2605.23993}, archivePrefix={arXiv}, primaryClass={cs.CV}, url={https://arxiv.org/abs/2605.23993}, } ```