Training a Tiny Doom Defender

Some time ago VAGO Solutions released SauerkrautLM-Doom-MultiVec-1.3M: a small ModernBERT encoder model that plays Doom Defend the Center scenario in real time on a CPU. It was trained with Supervised Fine-Tuning on 31k human gameplay examples.

My first thought: cool! I love both Doom and Small Language Models.

Then another idea: I bet I can improve the model. :-)

This sounded like a fun project. Also, since I had just released a little course on Reinforcement Learning Environments for LLMs, this could be an opportunity to see if RL can once again help go beyond human imitation.

If I've made you curious, keep reading to see how I changed the architecture, added PPO on top of SFT, and trained a stronger Doom player.

Scenario: Doom - Defend the Center

Defend the Center is a classic scenario from VizDoom, a popular library for developing AI bots that play Doom using visual information.

Here, the player is at the center of a large circular map. Monsters spawn along the wall and advance towards the player to attack him. Enemies can be killed with a single shot. 26 ammo is provided at the beginning and never refilled.

The best outcome achievable in this environment is killing 26 monsters.

The inspiration: SauerkrautLM-Doom-MultiVec-1.3M

Let's briefly recap the original model I took inspiration from. For more details, check out the paper and the GitHub repository.

In short, they model Doom gameplay as a classification problem: game image frame + additional info -> model -> action.

The model consists of three parts: input handling, encoder, classification head.

• First input: a 1000-character ASCII representation of the 640×480 RGB frame. To produce this, the image is converted to grayscale, downscaled to 40x25 and finally mapped to 10 chars indicating brightness. To support this representation, the model uses a character-level tokenizer.

• Second input: a depth buffer coming from VizDoom game engine, downscaled to 40×25 resolution, normalized to [0,1], and quantized into 16 discrete distance bins. Each depth bin has a learned 128-dim embedding that is added to the corresponding token embedding.

• Encoder: a 5-layer ModernBERT encoder modified with two-stage hash embeddings to reduce embedding parameters.

• Classification head:

  • Per-token embeddings are aggregated via learned attention pooling into a single 128-dim vector.
  • A linear layer classifies this pooled vector into 4 discrete actions: shoot, move forward, turn left, turn right.

Since Doom allows multiple actions at the same time, and combining them is beneficial, at inference time they adopt a composite strategy: a Python function that uses those action probabilities to combine them based on some thresholds.

SauerkrautLM-Doom-MultiVec-1.3M was trained on 31K human demonstration frames at frame skip = 4 (ViZDoom standard, an action every 4 game tics). They used soft action scores (0.85 for active keys, 0.05 baseline) and adopted KL-divergence loss between the model's softmax and the soft teacher scores.

They report that the model kills 17.8 enemies on average, beating all tested LLMs and running at 31 ms per move on CPU.

Independent evaluation

The reported score is the average of 10 episodes. I wanted a better estimate of the number to beat, so I decided to evaluate the original model on 1000 episodes. Seeds in this scenario determine where monsters appear and how they move; I used seeds in the range 10000-10999.

What I found: SauerkrautLM-Doom-MultiVec-1.3M kills 20.38 enemies on average; the standard deviation is 5.35.

A better estimate but a harder opponent to beat. Let's see what we can achieve.

Building the Tiny Doom Defender

Rethinking the action prediction

My initial idea was straightforward: they trained the model to imitate human examples; I can improve it by making it actually play the game with Reinforcement Learning (PPO).

Then I encountered a wall.

SauerkrautLM-Doom-MultiVec-1.3M outputs probabilities for 4 different actions, but at inference time the pipeline decides which action to actually apply using a Python function:

• Select the top action (e.g., move_forward)
• Add shoot if P(shoot)>0.75*P(top)
• Combine compatible actions (movement + rotation) if the second action's probability exceeds 0.15.

The action executed is not a sample from the distribution the model predicts.

While this is fine for inference, it's a big issue for policy-gradient Reinforcement Learning: these methods learn by sampling from the policy and alter the probability distribution based on rewards. Here the Python function is a middle-man between the model and the environment. If you include the Python function in the environment, you optimize a policy that never plays directly; if you include the function in the policy, it's frozen, non-differentiable code. Either way, the decisive part of the action lives outside the model, and PPO only trains the model. This might seem philosophical but I tried both and the model collapsed. :-)

Maybe there's a more natural way to design the model: multiple prediction heads, one for each axis (shoot, turn, forward). This brings the model closer to the actual game action space. (And we drop the annoying Python function)

This means that now we have a different model to train from scratch: tiny-doom-defender. Might be fun.

Why not deviate more from the original model?

The standard Doom Defend the Center scenario is designed to be solvable with just turning and shooting. Let's also drop moving forward.

In summary, our model now produces actions across 2 axes:

• turn: left, no, right
• shoot: no, yes.

Input handling

SauerkrautLM-Doom-MultiVec-1.3M receives 2 inputs: an ASCII representation derived from the image frame brightness and depth buffer embeddings obtained from VizDoom game engine.

While thinking about how to create a better model, I considered a few things.

First, the depth buffer is an extra image that tells the agent how far away each visible pixel is. This can be considered privileged info: normal Doom players can only see pixels, not the distance of elements in the game. So, for my model, I'd like to exclude this information, playing from pixels alone.

Let's focus for a moment on the image frame input instead. The RGB image is converted to grayscale and downscaled. Its brightness values are then used to produce an ASCII representation, which is tokenized and fed into the model.

Intuitively, there might be some limitations:

• overall, this process is lossy and can make our model "see" poorly
• converting the image to text which is then tokenized is not necessary per se
• there are explicit feature extraction choices that can probably be learned by the neural network instead.

Based on these considerations, I decided to radically simplify the input handling layer for tiny-doom-defender, drawing inspiration from Convolutional Neural Networks.

My model receives a 160×100 RGB frame, which is processed by two 2D convolutional layers. Each layer reduces the spatial dimensions by a factor of two, producing a final 40×25 feature map: this allows us to re-use the ModernBERT encoder without changing the sequence dimensions used by the existing architecture. With this approach, I hope the convolutional layers can learn to extract the most significant image features themselves.

As you may have noticed, the models discussed so far do not have access to temporal information: each action is selected only based on the current game status. So, for example, these models have no way to understand if a monster is approaching or standing still. To fight this limitation without significantly affecting inference time, in tiny-doom-defender I decided to stack 3 frames (at times t, t-1 and t-2). Since the recent changes in the image are also determined by the player actions, we also pass the actions at times t-1 and t-2.

In short, tiny-doom-defender sees the game directly: 3 stacked RGB frames go through a few small convolutional layers, while the last two actions are embedded separately and added in. This lets the model learn its own features and a sense of motion without touching the ModernBERT encoder.

Architecture summary

In addition to the changes discussed above, I also decided to try reducing the ModernBERT layers in the Encoder from 5 to 4; I hope they will be sufficient to make the model reason about the scenario. Now our model is 1.1M parameters instead of 1.3M.

Below is a simple diagram showing the final architecture of tiny-doom-defender. For more details, you can read the code here.

Training our model

It's now time to train our tiny model and see if it lives up to its promise.

We'll first make it learn from examples via Supervised Fine-Tuning and then refine its skills by playing, using Reinforcement Learning.

Recording gameplay data for Supervised Fine-Tuning

During this experiment, I found out that in the Reinforcement Learning world, Supervised Fine-Tuning is often referred to as Behavior Cloning. I like this term because it captures the idea that the model statistically imitates the training data and is limited by its quality.

Leaving aside terminology, we now need a certain amount of SFT data with this format: last 3 RGB frames, last 2 actions -> action.

For SauerkrautLM-Doom-MultiVec-1.3M, the authors played 2 hours and produced 31K examples. Armed with a bit of patience, I thought of doing the same: after all, it'd be nice to play Doom again, like in my childhood. But after a few games, I changed my mind: getting good scores in this scenario is not easy; recording hours of gameplay requires focus and can be stressful.

How else can I produce SFT data?

Maybe using a scripted Oracle. Let me explain... This could be a script with a simple heuristic that uses privileged information from the game engine (e.g. the enemy positions) to play well. Is this cheating? If the Oracle is only used to produce training data in the specified format, without exposing any of the privileged information, I'd say definitely not. To make this work, we just need to be a bit smart: we want this data to actually be learnable by our model that only sees RGB frames, so every oracle decision should be a function of what the model can see. After some iterations, I came up with an Oracle script based on this idea: spin one way, stop and shoot when a monster crosses the crosshair, resume spinning.

The oracle can be easily invoked to produce and save data:

record-oracle \
--episodes 400 --max-frames 100000 --save-to-dataset anakin87/doom-defend-the-center-100k-oracle

Using this command, in a few minutes, I produced a dataset with 100K examples. Nice!

The data were produced using seeds in the range 0-399.

Using 10000-10999 seeds, the oracle scores 23.05 mean kills; standard deviation is 2.85. One caveat: the score is useful as a reference but comparing it with a vision-only model is not apples-to-apples.

Supervised Fine-Tuning

Now that we have data, we can train.

I put together a simple SFT training script:

• builds the model
• loads the data
• as loss, uses the sum of the shoot and turn cross-entropy losses
• uses 90%/10% training/evaluation split
• measures joint accuracy: the fraction of samples where both actions are correct
• stops early if the evaluation metric does not improve for N evaluation
• saves the best performing checkpoint.

While the script is device agnostic, I recommend running it in a machine with a GPU (6GB VRAM should be enough) and a few CPU cores to make data processing fast and keep the GPU busy. In my case, I ran SFT on an NVIDIA L40 48GB with 12 CPU cores, which I used for less than an hour at $0.86/hr; you should be able to use smaller GPUs. I also tried running this on a Macbook M4 PRO (48GB RAM) but it was taking too long (hours) and I stopped; it would be nice to know exactly why, but I won't pretend to be an expert by pasting a Claude explanation here. :-)

On the L40 machine, I ran the script with:

train-sft \
  --data anakin87/doom-defend-the-center-100k-oracle \
  --bf16 --batch-size 256 --num-workers 12 --epochs 40

You can find the SFT model on Hugging Face.

Evaluating our model with the usual 10000-10999 test seeds, we get 22.19 mean kills, with a 3.76 standard deviation. Not bad at all! As expected, this is below the oracle score but beats SauerkrautLM-Doom-MultiVec-1.3M by +1.81 and has more consistent performance.

Reinforcement Learning

Time to let our model play!

In many experiments online, Reinforcement Learning is used to make models learn Doom from scratch. Here, the SFT model is already competent, so we use RL to refine its gameplay and improve its action distribution.

RL/PPO Recap

In Reinforcement Learning, the agent lives in an environment, observes its state and takes actions that alter it. The agent receives rewards based on its actions. Its goal is maximizing cumulative reward over time.

In our case, the agent is our tiny-doom-defender model; the environment is VizDOOM Defend the center scenario; we use a reward of +1 for each enemy killed and −1 for dying.

Proximal Policy Optimization is a popular RL algorithm.

PPO is typically implemented using an actor–critic architecture. The actor is the model that produces an action based on state. It is the part we ultimately want to improve, since it is what we use at inference time. The critic is another model (or a different prediction head in the same model) that estimates the expected return (value) based on the current state. The critic is used to train the actor more effectively: the value provides a baseline based on the current state; if the chosen action does better than the baseline (positive advantage), we want to reinforce it and vice-versa. During training, the critic is also updated over time to produce better estimates.

PPO improves the training stability of the policy (=our model) by preventing excessively large policy updates. Published in 2017 by OpenAI, it quickly gained popularity for several reasons: simplicity, computational efficiency and sample efficiency.

If you come from the generative Language Models world (like me), you might ask: why not use GRPO?

This is a simpler variant of PPO. The idea is to sample the model multiple times, creating a group; for computing the advantage, the baseline is the group mean and we no longer need a separate critic.

I considered this idea but then discarded it.

• Sample efficiency. Group-based setup is suitable for generative LMs: one prompt, multiple sampled responses to compare. This is less natural in an evolving Doom environment: to create a comparable group, I'd need to run several trajectories from the same initial state.
• Credit assignment. In standard GRPO, you get a reward at the end of the episode and reinforce all the actions of the same trajectory; in this Doom scenario, you get a reward for killing an enemy at a specific timestep, so with PPO it's easier to assign different advantages to actions at different timesteps.

PPO training

I can finally train the model.

I use a self-contained PPO implementation.

Here's the process I followed:

• for each hyperparameter configuration, train for 30 iterations using 12 parallel environments seeded sequentially from 42
• evaluate all checkpoints on 40 episodes using 50000-50039 validation seeds and pick the best one (the best policy is often not the final checkpoint)
• only at the end, get an unbiased score on 1000 episodes using 10000-10999 test seeds

While potentially expensive, this process ensures that we don't pick the best model based on noisy training returns, and that the final test score is not contaminated by training or model selection.

Programmatically, this can be done using the scripts:

train-ppo \
  --sft-checkpoint my-sft-model \
  --output my-ppo-model \
  ... # specify hyperparameters

select-ppo-snapshots \
  --output-dir my-ppo-model 

As part of the training and validation process, I conducted a quick hyperparameter search.

Chosen hyperparameters: γ 0.99, GAE λ 0.95, value coef 0.5, entropy coefficient 0.03, clip 0.15, grad clip 0.5, target-KL 0.05, 256 steps per env per rollout, 2 update epochs over 16 minibatches.

Based on these experiments, I also decided to train only the heads and the top encoder block and keep the main visual backbone of the model untouched. As is common, the heads are trained with a higher learning rate (5e-4) than the encoder (5e-5).

Finally, the best model can be evaluated on the test seeds:

eval-model \
  --ckpt my-ppo-model/policy_best \
  --episodes 1000

We get 23.12 mean kills, with a 2.81 standard deviation. Once again, we improved both the mean kill count and consistency.

I used the Macbook M4 PRO for PPO. Each round of training, selection and evaluation took less than two hours.

Conclusions

A table is worth a thousand words :-)

Evaluation on 1000 test episodes

Speaking of inference performance, it's hard to tell what hardware the authors used to measure the 31 ms per-decision inference time. It's only described as "CPU". So I ran a quick benchmark on my M4 Pro Mac. Compared to SauerkrautLM-Doom-MultiVec-1.3M, tiny-doom-defender is at least 40% faster when running on a single core; depending on how many cores you use, the difference can be even larger. Not a very serious estimate but I just wanted to be sure that my model is not slower.

Compared to SauerkrautLM-Doom-MultiVec-1.3M, I trained a smaller, simpler and faster model that can play without accessing privileged information. Despite this, it manages to kill more enemies with more consistent performance, matching the Oracle that has access to enemy positions!

The theoretical ceiling is 26 kills so there could still be room for improvement. In addition to hyperparameter search, I tried unfreezing more layers and introducing reward shaping; no success.

My impression is that the most promising lever is adding memory to the architecture (instead of just three stacked frames), and then letting PPO run longer.

Would you like to try? Let me know about your experiments!

Bonus 💾: I tried quantizing the model to 8 bits to make it fit into a floppy disk and I got a 1.1 MB quantized version that plays indistinguishably. On the same 1000 episodes, it gets 23.13 with 2.78 standard deviation.

I hope you liked this experiment as much as I did. Feel free to follow me or connect with me online.