PAWN: Progress-Aware World Models for UAV Vision-and-Language Navigation

PAWN is a vision-language navigation policy for UAVs built on Qwen2.5-VL, augmented with two auxiliary signals on top of standard action-token supervision:

• World model — future-view latent supervision (WM loss). For each predicted action step, PAWN predicts the visual latent of the next top-down view (in the backbone's vision-embedding space) via cross-attention from the current-view tokens to the generated action tokens. This is a cosine-distance loss against the ground-truth next-view embeddings produced by the frozen vision tower — no pixel decoder or VAE is needed.
• Progress-aware head. A small cross-attention + MLP head predicts how much of the instruction is completed after executing a candidate action sequence (a scalar in [0, 1], trained with BCE). At inference time the policy samples several candidate action sequences and keeps the one with the highest predicted progress.

This repository contains everything needed to reproduce SFT training and inference with both signals enabled, plus a small amount of sample data. Model weights are not included.

Repository layout

PAWN/
├── train.py                     # SFT entry point (dataset, collator, model build, trainer)
├── infer.py                     # Open-loop inference on a JSON of samples (self-contained)
├── eval.py                      # Closed-loop evaluation in the NavGym simulator
├── pawn/
│   ├── trainer.py               # PawnTrainer (loss weighting, LoRA + world_model export)
│   └── model/language_model/
│       └── pawn_qwen25vl.py     # PawnForConditionalGeneration + WorldModel (WM + progress head)
├── navgym/                      # UAV top-down navigation simulator (closed-loop stepping)
├── gsamllavanav/                # navigation env utilities (episodes, maps, geometry)
├── configs/zero2.json           # DeepSpeed ZeRO-2 config
├── scripts/
│   ├── train_sft.sh             # Launch SFT (torchrun + DeepSpeed)
│   ├── infer.sh                 # Launch open-loop inference
│   └── eval.sh                  # Launch closed-loop evaluation
├── data/
│   ├── sample_train.json        # A few training samples (relative image paths)
│   ├── sample_infer.json        # A few inference samples
│   └── images/                  # Bundled sample images
├── requirements.txt
└── README.md

PAWN provides two inference paths:

• infer.py (open-loop, self-contained): runs single-step prediction on a JSON of observations. Works out of the box on the bundled sample images — no simulator or benchmark data required. Use this to sanity-check a checkpoint.
• eval.py (closed-loop): steps the policy through the NavGym simulator and reports navigation metrics (NE / SR / OSR / SPL). This requires the public AirNav benchmark data (see below), which is not bundled.

Installation

python -m venv .venv && source .venv/bin/activate
pip install -r requirements.txt
# Optional but recommended for multi-image prompts:
pip install flash-attn --no-build-isolation

You also need a base Qwen2.5-VL checkpoint (e.g. Qwen/Qwen2.5-VL-7B-Instruct from the Hugging Face Hub, or a local path). Set it via MODEL_NAME_OR_PATH / BASE_MODEL_PATH in the scripts below.

Data format

Each training sample is a JSON object. The important fields:

Field	Type	Description
instruction	str	Natural-language navigation instruction
cur_view	str	Path to the current top-down view image
history_views	list[str]	Past view images (subsampled at indices [-7,-4,-2,-1])
history_actions	list[str]	Past discrete actions
future_actions	list[str]	Ground-truth next ≤8 actions (the LM target)
next_view	list[str]	Per-step future-view images aligned with future_actions (WM targets)
total_actions	list[str]	Full episode action sequence (used for the progress label)
cur_position	list[float]	UAV pose [x, y, z, heading]

The discrete action space is MOVE_FORWARD, TURN_LEFT, TURN_RIGHT, STOP.

Image paths in the bundled samples are relative to the repository root (e.g. data/images/img_0000.jpg); run training and inference from the repo root. For your own data, either use relative paths (from the repo root) or absolute paths.

The progress label is computed automatically in the dataset as min(1, (len(history_actions) + len(future_window)) / len(total_actions)).

Training

# from the repo root
MODEL_NAME_OR_PATH=Qwen/Qwen2.5-VL-7B-Instruct \
TRAIN_DATA_PATH=data/sample_train.json \
OUTPUT_DIR=checkpoints/sft \
bash scripts/train_sft.sh

The defaults assume 8 GPUs. For a single-GPU smoke test on the bundled sample data, override the device count:

NPROC_PER_NODE=1 CUDA_VISIBLE_DEVICES=0 bash scripts/train_sft.sh

Key options (all overridable as environment variables — see scripts/train_sft.sh):

Option	Default	Meaning
--world_model_enable	True	Enable the WorldModel (WM-loss branch)
--progress_enable	True	Train the task-progress head
lm_loss_weight	1.0	Weight of the action-token LM loss
wm_loss_weight	0.01	Weight of the future-view latent (world-model) loss
progress_loss_weight	0.2	Weight of the progress loss
LoRA	r=16, α=32	LoRA on the backbone; world-model/progress modules train full-rank

Checkpoints are written to OUTPUT_DIR/checkpoint-XXXX/. Each checkpoint also gets a hf_model/ subdirectory containing the LoRA adapter and world_model.bin (the cross-attention + progress-head weights) for inference.

The bundled data/sample_train.json has only a handful of samples and is meant for a smoke test / to illustrate the data format, not to reproduce a full model.

Inference

Inference samples several candidate action sequences per step and selects the one with the highest predicted task progress.

# from the repo root
BASE_MODEL_PATH=Qwen/Qwen2.5-VL-7B-Instruct \
bash scripts/infer.sh checkpoints/sft/checkpoint-5000/hf_model

or directly:

python infer.py \
  --base_model_path Qwen/Qwen2.5-VL-7B-Instruct \
  --adapter_dir checkpoints/sft/checkpoint-5000/hf_model \
  --input data/sample_infer.json \
  --output predictions.json \
  --num_rollouts 4

infer.py reads a JSON list of samples and, for each, writes the selected action sequence and its predicted progress (plus all scored candidates) to --output.

• --adapter_dir: an SFT hf_model directory (LoRA adapter + world_model.bin).
• Alternatively, --backbone_dir (a full merged Qwen2.5-VL directory) plus --world_model_path.

Without any weights, the progress head is randomly initialized and the selection is meaningless — always pass a trained checkpoint.

Closed-loop evaluation (NavGym + AirNav)

eval.py steps the policy through the NavGym simulator for each episode and computes planning metrics. It reuses the same ProgressPolicy as infer.py and adds the closed-loop environment loop.

Required benchmark data (not bundled)

Closed-loop evaluation uses the AirNav benchmark. Obtain it separately and place it under data/ with this layout:

data/
└── airnav/
    ├── maps/                             # ortho RGB-D maps (~GBs; one pair per map)
    │   ├── <map_name>.tif                #   GeoTIFF (channel 1 = height/elevation)
    │   └── <map_name>.png                #   RGB ortho image
    ├── objects.json                      # object annotations
    ├── processed_descriptions.json
    └── trajectories/
        ├── airnav_<split>[_<difficulty>].json   # trajectories / goal targets
        └── instruction_<split>.json             # {id: {episode_id, instruction}}

<map_name> stems must match across the .tif/.png files and the episode metadata. The default paths live in gsamllavanav/defaultpaths.py and navgym/models/CityNavData.py and can be overridden from the CLI.

The rendered drone views produced during evaluation are written to --render_dir (default nav_render_output/); this directory is an output and is safe to delete.

Run

# from the repo root
BASE_MODEL_PATH=Qwen/Qwen2.5-VL-7B-Instruct \
EVAL_DATA=data/airnav/trajectories/instruction_val_seen.json \
AIRNAV_DATA=data/airnav/trajectories/airnav_val_seen.json \
IMAGE_DIR=data/airnav/maps \
bash scripts/eval.sh checkpoints/sft/checkpoint-5000/hf_model

or directly:

python eval.py \
  --base_model_path Qwen/Qwen2.5-VL-7B-Instruct \
  --adapter_dir checkpoints/sft/checkpoint-5000/hf_model \
  --eval_data data/airnav/trajectories/instruction_val_seen.json \
  --airnav_data data/airnav/trajectories/airnav_val_seen.json \
  --image_dir data/airnav/maps \
  --num_gpus -1 --workers_per_gpu 1

Metrics are written to --output (default eval_results.json) and per-episode action sequences to --actions_output.

gsamllavanav also contains baseline / SoM / GroundingDINO / AirSim modules that PAWN evaluation does not use; they carry heavy optional dependencies (see requirements.txt) and can be ignored for reproducing PAWN.

Model overview

PawnForConditionalGeneration wraps a Qwen2.5-VL backbone and an optional WorldModel:

• LM loss: causal cross-entropy on the assistant action tokens.
• WM loss (WorldModel.compute_wm_loss): for each future step k, cross-attend from current-view tokens to the first k action tokens, and match the result to the k-th next-view embedding (cosine distance).
• Progress loss (WorldModel.compute_progress_loss): cross-attend from all input tokens to the action tokens, pool, and regress the progress scalar (BCE).

Loss weights are applied in the trainer (pawn/trainer.py), not in the model forward, so you can tune them without touching the model code.