导航世界模型，CVPR 2025（口头报告）
_{官方 PyTorch 实现}

论文 | 项目页面 | Notebook 演示 | 模型

本仓库包含导航世界模型——条件扩散变换器（CDiT）模型训练代码的官方 PyTorch 实现。更多结果请参阅项目页面。

导航世界模型
Amir Bar，Gaoyue "Kathy" Zhou，Danny Tran，Trevor Darrell，Yann LeCun
Meta AI、加州大学伯克利分校、纽约大学

设置

首先，克隆并设置仓库：

git clone https://github.com/facebookresearch/nwm
cd nwm

数据

要下载和预处理数据，请按照 NoMaD 中的步骤操作，具体如下：

• 下载数据集
• 将预处理分辨率从 (160, 120) 改为 (320, 240)，以获得更高分辨率
• 运行 process_bags.py 和 process_recon.py，将每个处理后的数据集保存到 path/to/nwm_repo/data/<dataset_name>。

对于 SACSon/HuRoN，我们使用的是包含更高分辨率图像的私有版本。请与数据集作者联系以获取访问权限（我们无法分发）。

最终，您的目录结构应如下所示：

nwm/data
├── <dataset_name>
│   ├── <name_of_traj1>
│   │   ├── 0.jpg
│   │   ├── 1.jpg
│   │   ├── ...
│   │   ├── T_1.jpg
│   │   └── traj_data.pkl
│   ├── <name_of_traj2>
│   │   ├── 0.jpg
│   │   ├── 1.jpg
│   │   ├── ...
│   │   ├── T_2.jpg
│   │   └── traj_data.pkl
│   ...
└── └── <name_of_trajN>
    	├── 0.jpg
    	├── 1.jpg
    	├── ...
        ├── T_N.jpg
        └── traj_data.pkl

依赖项：

mamba create -n nwm python=3.10
mamba activate nwm
pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu126
mamba install ffmpeg
pip3 install decord einops evo transformers diffusers tqdm timm notebook dreamsim torcheval lpips ipywidgets

训练

使用 torchrun：

export NUM_NODES=8
export HOST_NODE_ADDR=<HOST_ADDR>
export CURR_NODE_RANK=<NODE_RANK>

torchrun \
  --nnodes=${NUM_NODES} \
  --nproc-per-node=8 \
  --node-rank=${CURR_NODE_RANK} \
  --rdzv-backend=c10d \
  --rdzv-endpoint=${HOST_NODE_ADDR}:29500 \
  train.py --config config/nwm_cdit_xl.yaml --ckpt-every 2000 --eval-every 10000 --bfloat16 1 --epochs 300 --torch-compile 0

或者使用 submitit 和 slurm（8 台机器，每台 8 张 GPU）：

python submitit_train_cw.py --nodes 8 --partition <partition_name> --qos <qos> --config config/nwm_cdit_xl.yaml --ckpt-every 2000 --eval-every 10000 --bfloat16 1 --epochs 300  --torch-compile 0

或者在本地单 GPU 上进行调试：

python train.py --config config/nwm_cdit_xl.yaml --ckpt-every 2000 --eval-every 10000 --bfloat16 1 --epochs 300  --torch-compile 0

注意：torch compile 可以使训练速度提升约 40%。然而，它可能会导致不稳定以及不同 PyTorch 版本之间的行为不一致。请谨慎使用。

预训练模型

要使用预训练的 CDiT/XL 模型：

• 从 Hugging Face 下载预训练模型
• 将检查点放置在 ./logs/nwm_cdit_xl/checkpoints 目录下

评估

用于保存评估结果的目录： export RESULTS_FOLDER=/path/to/res_folder/

在单步时间预测上进行评估

1. 准备用于评估的真值帧（一次性）

python isolated_nwm_infer.py \
    --exp config/nwm_cdit_xl.yaml \
    --datasets recon,scand,sacson,tartan_drive \
    --batch_size 96 \
    --num_workers 12 \
    --eval_type time \
    --output_dir ${RESULTS_FOLDER} \
    --gt 1

2. 根据动作预测未来状态

python isolated_nwm_infer.py \
    --exp config/nwm_cdit_xl.yaml \
    --ckp 0100000 \
    --datasets <dataset_name> \
    --batch_size 64 \
    --num_workers 12 \
    --eval_type time \
    --output_dir ${RESULTS_FOLDER}

3. 报告与 GT 的指标（LPIPS、DreamSim、FID）

python isolated_nwm_eval.py \
    --datasets <dataset_name> \
    --gt_dir ${RESULTS_FOLDER}/gt \
    --exp_dir ${RESULTS_FOLDER}/nwm_cdit_xl \
    --eval_types time

结果保存在 ${RESULTS_FOLDER}/nwm_cdit_xl/ 中。

在跟随真值轨迹上进行评估

1. 准备用于评估的真值帧（一次性）

python isolated_nwm_infer.py \
    --exp config/nwm_cdit_xl.yaml \
    --datasets recon,scand,sacson,tartan_drive \
    --batch_size 96 \
    --num_workers 12 \
    --eval_type rollout \
    --output_dir ${RESULTS_FOLDER} \
    --gt 1 \
    --rollout_fps_values 1,4

2. 使用 NWM 模拟真值轨迹

python isolated_nwm_infer.py \
    --exp config/nwm_cdit_xl.yaml \
    --ckp 0100000 \
    --datasets <dataset_name> \
    --batch_size 64 \
    --num_workers 12 \
    --eval_type rollout \
    --output_dir ${RESULTS_FOLDER} \
    --rollout_fps_values 1,4

3. 报告与真值轨迹的指标（LPIPS、DreamSim、FID）

    python isolated_nwm_eval.py \
        --datasets recon \
        --gt_dir ${RESULTS_FOLDER}/gt \
        --exp_dir ${RESULTS_FOLDER}/nwm_cdit_xl \
        --eval_types rollout

结果保存在 ${RESULTS_FOLDER}/nwm_cdit_xl/ 中。

轨迹评估 - 规划

使用 1 步交叉熵方法规划，在 8 张 GPU 上运行（采样 120 条轨迹）：

torchrun --nproc-per-node=8 planning_eval.py \
    --exp config/nwm_cdit_xl.yaml   \
    --datasets recon   \
    --rollout_stride 1   \
    --batch_size 1   \
    --num_samples 120   \
    --topk 5   \
    --num_workers 12   \
    --output_dir ${RESULTS_FOLDER}   \
    --save_preds   \
    --ckp 0100000   \
    --opt_steps 1   \
    --num_repeat_eval 3

结果保存在 ${RESULTS_FOLDER}/nwm_cdit_xl/ 中。

BibTeX

@article{bar2024navigation,
  title={Navigation world models},
  author={Bar, Amir and Zhou, Gaoyue and Tran, Danny and Darrell, Trevor and LeCun, Yann},
  journal={arXiv preprint arXiv:2412.03572},
  year={2024}
}

致谢

我们感谢 Hirose Noriaki 在 HuRoN 数据集方面提供的帮助及分享的见解，同时也感谢 Manan Tomar、David Fan、Sonia Joseph、Angjoo Kanazawa、Ethan Weber、Nicolas Ballas 以及匿名审稿人提供的有益讨论和反馈。

许可证

代码和模型权重采用知识共享署名-非商业性使用4.0国际许可协议授权。详情请参阅LICENSE.txt。

NWM (Navigation World Models) 快速上手指南

本指南基于 CVPR 2025 Oral 论文《Navigation World Models》的官方 PyTorch 实现，帮助开发者快速部署和运行条件扩散 Transformer (CDiT) 模型。

环境准备

系统要求

• 操作系统: Linux (推荐 Ubuntu 20.04+)
• GPU: 支持 CUDA 12.6 的 NVIDIA 显卡（训练建议多卡环境，推理可单卡）
• Python: 3.10
• 包管理器: 推荐使用 mamba (比 conda 更快) 或 conda

前置依赖

确保系统已安装以下基础工具：

• Git
• FFmpeg (用于视频处理)
• CUDA Toolkit 12.6 (若使用预编译 PyTorch 包可省略，但需驱动匹配)

安装步骤

1. 克隆仓库

git clone https://github.com/facebookresearch/nwm
cd nwm

2. 创建并激活虚拟环境

使用 mamba 创建名为 nwm 的环境：

mamba create -n nwm python=3.10
mamba activate nwm

3. 安装深度学习框架

安装 PyTorch 夜间构建版（适配 CUDA 12.6）：

pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu126

提示：国内用户若下载缓慢，可尝试使用清华或阿里镜像源替换上述 index-url，但需确保版本兼容性。

4. 安装系统级依赖

mamba install ffmpeg

5. 安装 Python 依赖库

pip3 install decord einops evo transformers diffusers tqdm timm notebook dreamsim torcheval lpips ipywidgets

基本使用

1. 数据准备

NWM 依赖于特定的数据结构。请参考 NoMaD 进行数据下载和预处理。 关键修改：将预处理分辨率从 (160, 120) 改为 (320, 240) 以匹配高分辨率需求。

处理后的目录结构应如下：

nwm/data
├── <dataset_name>
│   ├── <name_of_traj1>
│   │   ├── 0.jpg
│   │   ├── 1.jpg
│   │   ├── ...
│   │   ├── T_1.jpg
│   │   └── traj_data.pkl
│   └── ...

2. 下载预训练模型

从 Hugging Face 下载预训练的 CDiT/XL 模型检查点，并将其放置在以下目录：

./logs/nwm_cdit_xl/checkpoints

3. 单步预测推理 (最简单的使用示例)

假设已准备好数据集 <dataset_name> 并设置了结果输出目录，执行以下命令进行未来状态预测：

export RESULTS_FOLDER=/path/to/res_folder/

python isolated_nwm_infer.py \
    --exp config/nwm_cdit_xl.yaml \
    --ckp 0100000 \
    --datasets <dataset_name> \
    --batch_size 64 \
    --num_workers 12 \
    --eval_type time \
    --output_dir ${RESULTS_FOLDER}

4. 评估指标计算

预测完成后，运行以下脚本计算与真实值 (GT) 的差异指标 (LPIPS, DreamSim, FID)：

python isolated_nwm_eval.py \
    --datasets <dataset_name> \
    --gt_dir ${RESULTS_FOLDER}/gt \
    --exp_dir ${RESULTS_FOLDER}/nwm_cdit_xl \
    --eval_types time

结果将保存至 ${RESULTS_FOLDER}/nwm_cdit_xl/<dataset_name>。

5. (可选) 本地调试训练

若需在单张 GPU 上进行调试训练：

python train.py --config config/nwm_cdit_xl.yaml --ckpt-every 2000 --eval-every 10000 --bfloat16 1 --epochs 300 --torch-compile 0

注意：设置 --torch-compile 1 可提升约 40% 训练速度，但在不同 PyTorch 版本间可能存在不稳定性，调试阶段建议设为 0。