Action-Conditioned Diffusion Transformers for DOOM

I trained a diffusion-style world model for DOOM using a Conditional DiT, with actions injected through adaLN and past visual context handled through cross-attention.

The live video interaction program and the model weights are here, check it out! https://huggingface.co/signaldetectionweary

My original inspiration came from comma.ai’s A Drive in the Office. I had already been thinking a lot about world models, and seeing their VQGAN+GPT2 style system made me want to try something out [1].

I did try a VQVAE+GPT implementation here without action conditioning.

I wanted to try action conditioning to make a diffusion-style World Model. I liked Navigation World Models (NWM), so I followed their style, using cross-attention for visual context, and AdaLN for action conditioning [2].

Instead of getting real-world footage or buying a robot, I built it around DOOM so I could train a RL agent to play DOOM and generate data synthetically.

I'd point out here that GameNGen and the DOOM implementation as a consequence of NeuralOS are good, modern points of comparison, although they are very different underlying structures [3,15].

1. Background

The why of all this came as inspiration from comma.ai’s A Drive in the Office [1]. I really liked the problem shape: learn from paired image and action streams, make the model act along a latent internalized memory.

Navigation World Models pushed that interest a bit further for me [2]. I liked the way it handled conditioning in a Conditional DiT: visual context comes in through cross-attention, while action and other global signals modulate the network more directly. Considering the issues with managing generative drift, learned latent denoisers are a stronger choice over token focused autoregressive transformers.

So the actual lineage of this project is a little different from what the final demo might suggest. The thing that got me excited was comma.ai. The architectural direction I wanted to try was closer to NWM. DOOM was the practical environment I could actually use. And GameNGen later became the natural comparison point once I realized how close the task framing was [3].

GameNGen is the right thing to compare against later, because it is the clearest diffusion-for-DOOM reference point I know of [3]. But the architecture I actually wanted to try was not really a GameNGen reproduction. It was closer to taking the CDiT conditioning style I liked from NWM and seeing what happened if I used it to run a game.

More broadly, this sits inside a wave of world-model and controllable video work that has been getting a lot more interesting lately — DIAMOND, Genie, Oasis, OpenAI’s world-simulator writeup, Vista, iVideoGPT, and Genie 2 all come to mind [7,8,9,10,11,12,13]. I am not claiming this project is on the same scale as those, but they were definitely part of the air around this build.

2. Preparation for Pretraining

I needed a lot of continuous DOOM video, and I needed it fast enough that repeated VAE and diffusion training would actually be practical. So before the model itself, a big part of this project was just getting data collection and the storage/loading pipeline into a shape that made pretraining feasible.

Using HASARD on top of Sample Factory and VizDOOM, plus a bunch of extra scenario boilerplate, I trained agents and generated a large corpus of gameplay clips across several setups [4,5,6].

3. KL-Regularized Autoencoder

I wrote an autoencoder in the Stable Diffusion style — four upsamples or downsamples, targeting an 8x8 latent grid — and followed a training scheme similar in spirit to the VAE recipes that have become common in recent latent video work [14]. In practice that meant: first warm up with L1 and a tiny KL term, then lean on LPIPS once the reconstructions stabilize, and finally use a short adversarial finetune to tighten detail.

Pretrain warmup

Because LPIPS can introduce ugly artifacts early, I like to let the autoencoder stabilize first. A few thousand steps is usually enough to get the basic reconstruction behavior into a sane regime.

VAE LPIPS pretrain

It looks better pretty quickly once LPIPS comes in, but seeing is not really the whole story.

I usually train until the obvious metrics flatten out in a stable way: L1, LPIPS, KL, gradnorm, and the latent mean statistics. Having a well-behaved $\mu$ and latent scale matters later once diffusion starts trying to work inside that representation.

A strong VAE solves a surprising number of downstream problems. That has been true in a lot of my past work, and it was definitely true here too.

VAE adversarial loss finetune

Once I had a reasonably strong autoencoder, I used a short adversarial finetune to tighten up details. Sometimes this matters less if the latent grid is large enough, but with an 8x8 latent grid it still helps.

At that point the reconstructions were good enough to move on, even if the grass still had a bit of a fruity-pebbles texture to it.

4. Using a Diffusion Transformer

The denoiser here is a latent-space Diffusion Transformer. The autoencoder is Stable-Diffusion-style in spirit [14], but the denoiser itself is not a U-Net path adapted from SD1.4. I wanted to try a DiT directly.

Adaptive LayerNorm (adaLN) and Cross-Attention

One straightforward way to think about conditioning in a diffusion transformer is that not every conditioning signal has to enter the model in the same way.

For this project I ended up liking the NWM / Conditional DiT split more than the GameNGen split [2]. Past visual context is rich, structured, and high-dimensional, so cross-attention makes a lot of sense there. Actions are much smaller and more global, so feeding them through adaLN-style modulation felt more natural to me than turning them into the main attended token stream.

So in this implementation, cross-attention is used to attend over patchified latent context frames, while the action sequence is encoded by a small transformer, pooled, and injected as part of the global conditioning vector used by adaLN.

This was not some grand miraculous design decision so much as me following the kind of architecture I wanted to use later for broader video work. I liked the NWM-style split better, and I wanted to see it in a DOOM setting.

Action-Conditioned AdaLN + Cross-Attention DiT

My denoiser is a latent-space DiT that predicts the reverse diffusion direction for a noisy target latent $z_t$ while conditioning on both past latent context and an action history in the style of Amir Bar's Navigation World Models.

Rather than shoving everything into one token stream, we separate the input frame, context frames, and action conditioning. Context latents are exposed as memory through cross-attention. Action information is encoded with a separate TransformerEncoder and used as conditioning for AdaLN-Zero. And finally, diffusion timestep, prediction horizon, action history, and context-noise level are fused into a single conditioning vector that drives adaLN modulation throughout the block.

NWM's Action-conditioned DiT block with adaLN modulation, self-attention, cross-attention into context memory, and gated residual branches. The embedding transformer for AdaLN-Zero is a design divergence.

The noisy target latent $z_t$ is patchified into a 2D token sequence and given spatial positional embeddings. Context frames are patchified with the same embedder and receive both spatial and frame-time embeddings before being flattened into a memory bank.

The global conditioner combines diffusion time, prediction horizon, action sequence, and optional context-noise metadata into a single vector. The action history itself is encoded with a small masked transformer and pooled before being added into that conditioning stream.

Each DiT block applies three pre-norm residual updates: self-attention over the noisy target tokens, cross-attention into context memory, and an MLP. AdaLN predicts a shift, scale, and residual gate for each branch, so conditioning controls both the normalized features and the strength of each residual path.

After the transformer stack, I apply one final adaLN modulation before projecting tokens back to latent patches and unpatchifying them into the original spatial layout.

So the short version is: GameNGen gave me a very good point of comparison, but the DiT architecture I actually wanted to use was much closer to NWM/CDiT.

I also still think GameNGen’s context-noising idea is just very good, so I keep that practical lesson from their setup even while changing the conditioning split [3].

Training DiT to visualize a world model

Diffusion-style world models are very much a real thing now, even if there are lots of different ways to structure them [3,7,8,9,10,11,12,13]. In my case, the main training intuitions that mattered were using noised context frames, progressively longer context run-ups, and just accepting that the DiT part is a bit of a slog. Even with only a context capacity of four preceding frames, single-frame prediction converges.

The DOOM diffusion world

A peek at live autoregressive inference at different stages of training.

Inference is done autoregressively, using a seed sequence taken from the data as the initial four-frame context, then dequeuing and caching future generated frames for the next context cycle, on and on.

Coming soon: the DiT and VAE checkpoints will live on Hugging Face.

5. Another perspective

After the initial version of this project, I found two really nice public GameNGen-style reproductions: Arnaud’s work here and Masao Taketani’s implementation here. I followed their idea and constrained my models to train on VizDOOM data trained on the Deathmatch map.

I still think GameNGen is the right comparison point for this page, because the question “can you get a coherent DOOM diffusion world model?” is the obvious one. GameNGen is the benchmark I want to be legible against. NWM is the architectural direction I liked more.

Remaking the data, new RL agent, new DiT pretrain

Still, constraining the training data to a singular DOOM map, "Deathmatch," should be a simplification that will help remove the previous issue I have with player teleporting into new maps and their gun appearing and disappearing. So I train a new RL agent, make a new dataset, and train a new DiT.

Using sample factory, I train a new RL agent on VizDOOM's simple deathmatch. I used a IMPALA style ConvNet with LSTM. 20M steps later and the agent appears to win, I've stuck it here.

And here it shows I should really retrain the VAE to include these weapons, as the shotgun doesn't look so convincing. But I've been at this for a while and I figure I'll just give it a whirl.

And finally, here we are.

Thanks for reading.