RohanNaga/doomdit-nexttic

DoomShift world models

The checkpoints behind DoomShift: How Much Adaptation Do World Models Need Under Domain Shift? (Nagabhirava, Chirumamilla, Liu; CoRL 2026 Workshops PhysWM and TAG).

Project page · Paper · Code · Frames · Latents

Three pretrained diffusion image models, each turned into a next-tic Doom world model under one recipe: predict the next frame's latent from 32 context latents and the newest action (29 discrete actions). Each model is trained on 4 Doom maps and tested on 13 maps it never saw. This repo also holds the fine-tuned VAE decoders the reported numbers are decoded through, and every adaptation run (LoRA and full fine-tuning) on the 13 unseen maps.

Contents

Path	What it is	Size
040-unet-nexttic/	SD 1.4 U-Net (860M), 200k updates; snapshots every 10k	20 × 3.4 GB
041-pixart-nexttic/	PixArt-α (628M), 200k updates; snapshots every 50k	4 × 2.5 GB
042-sd35-nexttic/	SD 3.5 Medium (2.27B), 200k updates; snapshots every 10k	20 × 9.5 GB
decoders/vae_decoder_sd1x_mse_lpips/	SD 1.x VAE decoder tuned on training-map frames, MSE + 0.1 LPIPS, encoder frozen. Used for the U-Net and PixArt-α rows
decoders/vae_decoder_sd35_mse_lpips/	The same recipe in the 16-channel SD 3.5 latent space
decoders/vae_decoder_sd1x_mse/	MSE-only tune (the GameNGen decoder recipe), kept for comparison
adapters/adapt/runs/unet200k_arenas13_mapNN_r16_k8_s0/	U-Net rank-16 LoRA, 8 adaptation episodes, one run per unseen map	13 runs
adapters/adapt/pixart200k_arenas13_mapNN_r16_k8_s0/	PixArt-α rank-16 LoRA, same protocol	13 runs
adapters/adapt_sd35/mapNN/	SD 3.5 rank-16 LoRA, same protocol	13 runs
adapters/adapt_fullft_g8k/mapNN/	U-Net full fine-tuning, the comparison row	13 runs
adapters/adapt/splits/	The per-map adaptation splits (which 8 episodes each run trained on)
adapters/adapt_*_full*, adapters/adapt_fullft/	Additional and preliminary runs, kept for completeness

Each world-model directory holds config.json (the exact training arguments), log.jsonl (the training log), train_episodes.json, and the snapshots snap_<step>.pt. The final snapshot has a .sha256 file. Each snapshot carries the live weights (model) and the fp32 EMA at decay 0.9999 (ema). The paper's zero-shot rows and every adaptation start from the 200k EMA weights.

Each adaptation run directory holds adapters at updates 0, 250, 500, 1k, 2k and 4k (adapter_<step>.pt, with sha256 files), plus config.json, log.jsonl and scores.jsonl. The scores file records every evaluated checkpoint: scene PSNR, LPIPS, the decoder used, the source snapshot's sha256 and the split's sha256. An adapter file holds only the LoRA weights and the few parts trained in full (action embedding, input projection, noise embedding). It names its frozen source snapshot, which you download separately.

Results

Medians over maps, one tic ahead, scored on the scene (HUD excluded) against the raw frame:

Backbone	Scene PSNR (dB) ↑ train → unseen	LPIPS ↓ train → unseen	Turn-response ↑ train → unseen
SD 1.4 U-Net	25.20 → 22.30	0.158 → 0.303	0.86 → 0.81
PixArt-α	25.18 → 22.38	0.159 → 0.285	0.84 → 0.79
SD 3.5 Medium	25.38 → 22.09	0.126 → 0.262	0.84 → 0.81

Rank-16 LoRA adapters trained on 8 episodes of the new map for 4k updates (2% of the pretraining updates, about 1% of its compute) recover 72–76% of the LPIPS rise and about half of the lost PSNR. Full fine-tuning of the U-Net recovers more (88% vs. 72%), but it trains 205× more parameters and forgets 2.7× more on the training maps (1.08 vs. 0.40 dB).

The turn-response score flips the last turn action and counts the windows whose predicted horizontal shift reverses direction. A model that ignores actions scores 0. A perfect model scores 0.89, because 11% of recorded turns move the view too little to measure in one tic.

Usage

Clone the code, install requirements.txt, then download what you need:

hf download RohanNaga/doomdit-nexttic --include "040-unet-nexttic/snap_0200000.pt*" "040-unet-nexttic/config.json" \
    "decoders/vae_decoder_sd1x_mse_lpips/*" --local-dir weights
hf download RohanNaga/doom-dense-arnold-latents --repo-type dataset --include "sd15/*" --local-dir data/latents

Score a world model or an adapter with eval_tf.py. Give it the backbone, the checkpoint, the tuned decoder and the evaluation latents, and pass --use-ema to score the EMA weights. Given an adapter, eval_tf.py rebuilds the frozen source snapshot the adapter names and applies the adapter on top. python eval_tf.py --help lists every flag. Sampling is 10-step DDIM with v-prediction.

Training

All three backbones share one recipe:

• v-prediction with context-noise augmentation
• 200k updates at global batch 32, learning rate 5e-5 with 2k warmup
• bf16 autocast with fused AdamW, and the EMA kept in fp32
• 32 context tics and the newest action

The U-Net and PixArt-α share the SD 1.x KL-f8 latent space (4 channels, stabilityai/sd-vae-ft-mse), and SD 3.5 uses its own 16-channel VAE. Frames are 320×240, padded to 320×256, giving 32×40 latents. Training data is the 4 training maps of RohanNaga/doom-dense-arnold. config.json in each directory has the exact arguments and the git commit. The U-Net and SD 3.5 configs show steps: 400000. Both runs were stopped at 200k, the checkpoint the paper reports. The learning rate is constant after warmup, so the stop matches a planned 200k run.

Limitations

The models are trained on one game, Doom, played by one agent, Arnold. The paper evaluates them one tic ahead and measures only basic control, so long-horizon rollout quality is not characterized. The weights are research artifacts for studying domain shift, not general-purpose video models.

Licensing

Each world model is a fine-tune of its pretrained backbone and is released under that backbone's license:

• 040-unet-nexttic/ and the U-Net adapters: CreativeML OpenRAIL-M (Stable Diffusion 1.4)
• 041-pixart-nexttic/ and the PixArt-α adapters: CreativeML OpenRAIL++-M (PixArt-α)
• 042-sd35-nexttic/, its adapters and decoders/vae_decoder_sd35_mse_lpips/: Stability AI Community License (Stable Diffusion 3.5 Medium)
• decoders/vae_decoder_sd1x_*: MIT, the license of stabilityai/sd-vae-ft-mse

The use restrictions of those licenses carry over to these weights. The training code is MIT apart from the DiT-derived files; see the code repository's NOTICE.md.

Citation

@inproceedings{nagabhirava2026doomshift,
  title     = {DoomShift: How Much Adaptation Do World Models Need Under Domain Shift?},
  author    = {Nagabhirava, Rohan and Chirumamilla, Keerthana and Liu, Changliu},
  booktitle = {CoRL 2026 Workshop on Bringing Physics Simulation and World Models Together for Robotic Manipulation},
  year      = {2026}
}