ANWM: Aerial World Model for Long-horizon Visual Generation and Navigation in 3D Space

A physics-informed, action-conditioned world model for aerial visual generation and navigation in large-scale 3D environments.

Overview

ANWM (Aerial Navigation World Model) is an action-conditioned world model that predicts future egocentric observations for UAV navigation in large-scale 3D environments. Conditioned on recent visual history and a 4-DoF UAV action (delta x, delta y, delta z, delta yaw), it autoregressively imagines observations along candidate trajectories and ranks each path by visual similarity to the goal.

Highlights:

• TB-scale pre-training on action-conditioned world models with 4-DoF UAV trajectories, including large-scale data curation and high-throughput distributed training.
• Physics-inspired Future Frame Projection (FFP) that projects historical frames to future viewpoints, injecting coarse geometric priors and stabilizing long-horizon visual generation.
• Long-horizon visual forecasting and stronger navigation success: extending the effective prediction horizon from under 10 m (indoor) to the hundreds of meters scale in outdoor open-space environments.

Dataset

Available on Hugging Face.

The simulated benchmark is built from AerialVLN, OpenFly trajectories.

Split	Trajectory segments	Frames per segment	Actions per segment
Train	350K	48	47
Test	2.2K	48	47

The test set contains 1.1K planar and 1.1K 3D trajectories. Horizontal speed is 5 m/s, vertical speed is 2 m/s, yaw speed is 15 deg/s, and the average trajectory length is approximately 80.7 m. Real-world evaluation uses the drone-view subset of Sekai with estimated depth.

Installation

Python 3.9 or newer and a CUDA-enabled PyTorch installation are required. Install the PyTorch build matching the local CUDA runtime first, then install ANWM:

conda create -n anwm python=3.9 -y
conda activate anwm
python -m pip install -e ".[metrics]"
python scripts/check_environment.py --component metrics --verify-imports

The Stable Diffusion VAE stabilityai/sd-vae-ft-ema is downloaded by Diffusers on first use.

Data and Checkpoint

Download the dataset from Hugging Face and extract it into data/airvln_16/:

huggingface-cli download EmbodiedCity/ANWM-Dataset --repo-type dataset --local-dir data/airvln_16_shards
mkdir -p data/airvln_16
for f in data/airvln_16_shards/airvln_16-*.tar; do tar -xf "$f" -C data/airvln_16; done

Download the released checkpoint from Hugging Face into the path expected by infer.py:

huggingface-cli download EmbodiedCity/ANWM 0200000.pth.tar \
  --local-dir logs/anwm_cdit_airvln/checkpoints

The repository includes the released split manifests under data/splits/. Place processed trajectories and the ANWM checkpoint at:

data/airvln_16/<trajectory>/...
logs/anwm_cdit_airvln/checkpoints/0200000.pth.tar

The original AirVLN-16 preparation code is provided at data/preprocessing/Data_preprocessing_airvln_16.ipynb. Set root_dir in the notebook to the local AirVLN directory. The notebook writes trajectory folders containing numeric image files and traj_data.pkl.

Inference

Generate ANWM rollout predictions:

torchrun --standalone --nproc_per_node=1 infer.py \
  --exp config/anwm.yaml \
  --ckp 0200000 \
  --datasets airvln_16 \
  --eval_type rollout \
  --rollout_fps_values 1,4

Generate the matching ground truth:

torchrun --standalone --nproc_per_node=1 infer.py \
  --exp config/anwm.yaml \
  --datasets airvln_16 \
  --eval_type rollout \
  --rollout_fps_values 1,4 \
  --gt 1

Evaluate LPIPS, DreamSim, and FID:

python evaluate.py \
  --datasets airvln_16 \
  --gt_dir outputs/inference/gt \
  --exp_dir outputs/inference/anwm \
  --eval_types rollout \
  --rollout_fps_values 1,4

Training

The released run used six GPUs. config/anwm.yaml contains the paper model and optimization parameters:

torchrun --standalone --nproc_per_node=6 train.py \
  --config config/anwm.yaml \
  --global-seed 0 \
  --bfloat16 1 \
  --torch-compile 1

Checkpoints are written to logs/anwm_cdit_airvln/checkpoints/. The numbered checkpoint used for the reported model is 0200000.pth.tar.

Real-world Planning

Real-world evaluation follows the same overall workflow as simulation, but uses different data, depth, and entry scripts:

	Simulation	Real-world
Data	AirVLN-16	Sekai drone-view
Depth	simulator GT	Pi-Long
Eval	infer.py / evaluate.py	planning_eval.py
Config	config/	real/config/

Key files (other scripts under real/tools/ are optional):

Role	File
Convert videos → AirVLN-16	real/tools/process_youtube_to_airvln16_format.py
Fill depth with Pi-Long	real/tools/fill_depth_pilong.py
One-shot preprocess	real/tools/prepare_dataset.sh
Dataset adapter	real/dataset.py
Planning evaluator	real/planning_eval.py
Configs	real/config/
Split + candidates	data/splits/sekai_new/

Install deps, preprocess (needs ffmpeg; keep Pi-Long under third_party/Pi-Long or elsewhere), then evaluate. Point RAW_ROOT / OUTPUT_ROOT at your own directories to replace or re-index data.

python -m pip install -e ".[real]"
python scripts/check_environment.py --component real --verify-imports

RAW_ROOT=data/sekai_raw \
OUTPUT_ROOT=data/sekai_new \
PILONG_DIR=third_party/Pi-Long \
bash real/tools/prepare_dataset.sh

torchrun --standalone --nproc_per_node=1 planning_eval.py \
  --exp config/anwm.yaml \
  --ckp 0200000 \
  --datasets sekai_new \
  --output_dir outputs/planning \
  --num_samples 5

Repository Layout

anwm/             model, diffusion, projection, dataset, and rollout code
config/           paper training and simulation evaluation configuration
data/             original preprocessing notebook and released split metadata
real/             Sekai data adapters, depth fill, and planning evaluation
train.py          distributed training entry point
infer.py          visual generation entry point
evaluate.py       LPIPS, DreamSim, and FID evaluation
planning_eval.py  real-world trajectory-ranking entry point

Released Configuration

The paper model is defined by config/anwm.yaml:

Parameter	Value
Backbone	CDiT-XL/2
Parameters	~1.1B
Image size	224 x 224
Context frames	4
Prediction steps	16
Goals per observation	4
Batch size per GPU	1
Learning rate	8e-5
Gradient clipping	10.0
Training diffusion steps	1000
Inference sampling steps	250
Released checkpoint	200,000 steps

Limitations

• Generation may drift or collapse as trajectories approach 200 m.
• Fine structures such as windows and facades can become distorted.
• FFP depends on depth quality and camera calibration.
• Real-world results are offline evaluations, not autonomous UAV deployment.

Acknowledgements

We sincerely thank the following projects for their excellent work: AerialVLN, OpenFly, OpenUAV, Sekai, NWM, Matrix-Game, and YUME.

Citation

@article{zhang2025anwm,
  title={Aerial World Model for Long-horizon Visual Generation and Navigation in 3D Space},
  author={Zhang, Weichen and Tang, Peizhi and Zeng, Xin and Man, Fanhang and Yu, Shiquan and Dai, Zichao and Zhao, Baining and Chen, Hongjin and Shang, Yu and Wu, Wei and Gao, Chen and Chen, Xinlei and Wang, Xin and Li, Yong and Zhu, Wenwu},
  journal={arXiv preprint arXiv:2512.21887},
  year={2025}
}