NanoDreamer: Learning to Act Through Imagination

Benjamin BaiOctober 2026GitHub

Recently, I've been getting into world models and embodied intelligence, and because I love video games, obviously Dreamer caught my attention. However, I had one problem. I have no GPU.

So what do I do instead? Build my own tiny Dreamer, small enough to train on my MacBook. And because I wanted to understand every piece, I started from the bottom with my own autograd engine, tensor library, neural network layers, optimizers, and transformer.

After a lot of debugging, that became NanoDreamer! A 750k-parameter world model and actor–critic that learns to play all 5 MinAtar games, with the actor trained entirely through imagination and a model small enough to run and train on my M5! It trains from scratch on its own gameplay, without expert demonstrations or game-specific strategies, using the same architecture and training setup with separate weights for each game.

Here's NanoDreamer playing all 5 games after 500k steps per game (around 3 hours on an M5 Pro).

NanoDreamer playing Space Invaders
NanoDreamer playing Space Invaders.
How to play

Move your cannon along the bottom, shoot the descending aliens, and dodge their shots. A hit from an alien or enemy shot ends the game.

On screen
  • Your cannonalong the bottom row
  • Alienscolor indicates their movement direction
  • Your shotstravel upward
  • Enemy shotstravel downward
Training reward

+1 for each alien shot. Moving, firing, and surviving a step do not give extra points.

In this write-up I'll show you the bugs, experiments, and architectural changes it took to get Dreamer working at a fraction of the cost. If you're new to ML, I also wrote a separate guide to the foundations, with my own implementation and learnings!

How hard could it be?

I started by reading DreamerV3 and DreamerV4, trying to figure out what I actually needed to build. The idea seemed simple enough. A world model learns how the game works, an actor chooses actions inside its predictions, and a critic estimates the rewards beyond them. Here's what my first attempt looked like.

First NanoDreamer architectureFrames pass through an encoder and history into a causal transformer and transition heads. The actor supplies the chosen action. The critic estimates future returns. Separate training summaries show reconstruction, prediction and behavior objectives. Green dots identify changes in the working version.01 / REPRESENTATION & PREDICTIONDecoderMLP · reconstruction lossActorreads zₜ; later [zₜ, Δzₜ]FramesEncoderMLP · frozen after pretrainingLatent historyTransformercausal · current action includedNext latentz′ = z + ΔzRewardscalar · MSEContinuationprobability · binary CEPrediction losscombined learning signalCriticscalar V(actor input)02 / BEHAVIOR TRAININGActor objective−log π(a | state) × advantageReturn targetsrewards + critic · λ-returnsCritic objectiveMSE(V, imagined return)Click a component to explore
First draft
First NanoDreamer architectureFrames pass through an encoder and history into a causal transformer and transition heads. The actor supplies the chosen action. The critic estimates future returns. Separate training summaries show reconstruction, prediction and behavior objectives. Green dots identify changes in the working version.01 / REPRESENTATION & PREDICTIONDecoderMLP · reconstruction lossActorreads zₜ; later [zₜ, Δzₜ]FramesEncoderMLP · frozen after pretrainingLatent historyTransformercausal · current action includedNext latentz′ = z + ΔzRewardscalar · MSEContinuationprobability · binary CEPrediction losscombined learning signalCriticscalar V(actor input)02 / BEHAVIOR TRAININGActor objective−log π(a | state) × advantageReturn targetsrewards + critic · λ-returnsCritic objectiveMSE(V, imagined return)

Click a component to jump to its explanation.

I started with an MLP autoencoder that compressed each frame into a 64-dimensional latent embedding. Once it could reconstruct the screen, I froze its weights.

A small causal transformer read the latent embeddings alongside action and positional embeddings. Its three prediction heads estimated the next latent state, reward, and probability the game would continue. The actor and critic were MLPs that saw the current latent concatenated with its change from the previous frame.

The actor then trained through 5-step imagined rollouts, sampling an action and feeding each predicted latent back into the model. We combined the predicted rewards with the critic's estimate of future return to give the actor feedback on its actions.

In my head, this all worked in theory. Now I just had to see whether it actually would.

Inside the first draftArchitecture, losses, and imagined returns

First, describe the game in 64 numbers

A small MLP autoencoder compressed each game frame into 64 numbers, then tried to reconstruct it.

Most cells were empty. A model could therefore get many cells right by predicting an almost empty screen. I used weighted binary cross-entropy to make missing an occupied cell more costly.

Open any equation to see what the symbols mean.

Lreconstruction=−mean⁡i[wixilog⁡x^i+(1−xi)log⁡(1−x^i)]\mathcal L_{\text{reconstruction}}=-\operatorname{mean}_i\left[w_i x_i\log\hat x_i+(1-x_i)\log(1-\hat x_i)\right]

Each xᵢ is the real occupancy of one cell in one channel, either zero or one. Its prediction x̂ᵢ is a probability. A hat means predicted, wᵢ gives occupied cells more weight based on their rarity in that channel, and mean averages the errors across the batch. If a cell is occupied, its loss is −wᵢ log x̂ᵢ. Predicting 0.9 gives about 0.105wᵢ loss, while predicting 0.1 gives about 2.303wᵢ. Confident mistakes get larger penalties. The other term handles empty cells.

I then froze the autoencoder. The transformer had to work with whatever those 64 numbers preserved.

Then, predict what an action does

The transformer read eight steps of latent and action history. It had two causal blocks, with width 64 and four attention heads.

Three heads read the output. The next-latent head predicted a change to the current latent. The reward head predicted a number, and the continuation head predicted the probability that the episode kept going.

z^t+1=zt+Δzt,Lnext latent=MSE⁡(z^t+1,E(xt+1))\hat z_{t+1}=z_t+\Delta z_t,\quad \mathcal L_{\text{next latent}}=\operatorname{MSE}(\hat z_{t+1},E(x_{t+1}))
Lreward=MSE⁡(r^t,rt),Lcontinuation=BCE⁡(c^t,ct)\mathcal L_{\text{reward}}=\operatorname{MSE}(\hat r_t,r_t),\qquad \mathcal L_{\text{continuation}}=\operatorname{BCE}(\hat c_t,c_t)

zₜ is the current 64-number latent and Δzₜ is the predicted change. E is the frozen encoder, so E(xₜ₊₁) is the encoding of the real next frame. MSE means mean squared error. Subtract each predicted value from its target, square the difference, then average across latent values and sampled transitions.

rₜ is the real reward and r̂ₜ is its prediction. MSE averages squared errors, so predicting 0.2 for a reward of 1 gives (0.2 − 1)² = 0.64. cₜ is one if the episode continues and zero if it ends. BCE means binary cross-entropy. It penalizes confident incorrect continuation probabilities, just as reconstruction BCE penalizes incorrect occupancy probabilities.

The predicted-frame loss decoded the predicted latent and compared it with the real next frame. Cells that changed received more weight, so a small moving object was harder to ignore. I combined the four losses like this.

Lworld=Lnext latent+2Lreward+0.2Lcontinuation+0.5Lpredicted frame\mathcal L_{\text{world}}=\mathcal L_{\text{next latent}}+2\mathcal L_{\text{reward}}+0.2\mathcal L_{\text{continuation}}+0.5\mathcal L_{\text{predicted frame}}

The coefficients are manual balancing weights. Relative to next-latent loss, 2 increases the reward contribution, while 0.2 and 0.5 reduce the continuation and predicted-frame contributions. These were prototype settings, not values derived from Dreamer or established as optimal by a controlled comparison. The losses have different scales, so 2 does not mean reward is exactly twice as important.

Now, let the actor practise

I started with real history from replay. The actor chose an action, the frozen model predicted the next latent, reward, and continuation, and the actor chose again. I repeated this for five imagined actions.

The critic estimated rewards beyond the five-step window. Assuming the episode continues, a simple return adds predicted rewards to that final estimate.

G0=r^0+γr^1+γ2r^2+γ3r^3+γ4r^4+γ5V(s5)G_0=\hat r_0+\gamma\hat r_1+\gamma^2\hat r_2+\gamma^3\hat r_3+\gamma^4\hat r_4+\gamma^5V(s_5)

G₀ is the return target at the starting state. Each r̂ is a reward predicted by the world model. γ discounts rewards further into the future. V(s₅) is the critic’s estimate of rewards beyond the fifth imagined action, which keeps the rollout ending from being treated as the game ending. This simplified equation assumes the episode continues throughout.

The actual script used lambda returns, mixing the next critic estimate with the longer return. I used γ = 0.99 and λ = 0.95. Here s is the actor–critic input state.

Gtλ=r^t+γc^t[(1−λ)V(st+1)+λGt+1λ],G5λ=V(s5)G_t^\lambda=\hat r_t+\gamma\hat c_t\left[(1-\lambda)V(s_{t+1})+\lambda G_{t+1}^\lambda\right],\qquad G_5^\lambda=V(s_5)

Gₜ with the λ superscript is the return target at decision t. The next predicted reward is added to a mixture of the next critic estimate and the longer return. λ controls that mixture, γ discounts the future, and ĉₜ is the continuation probability. Work backward from G₅ = V(s₅). A continuation probability of zero removes future rewards after an ending.

Each decision got its own return target. I subtracted what the critic had expected to get the advantage. A positive advantage made the sampled action more likely, while a negative one made it less likely.

At=sg⁡(Gtλ−V(st)),Lactor=−Et[log⁡π(at∣st)At]A_t=\operatorname{sg}(G_t^\lambda-V(s_t)),\qquad \mathcal L_{\text{actor}}=-\mathbb E_t[\log\pi(a_t\mid s_t)A_t]
Lcritic=Et[(V(st)−sg⁡(Gtλ))2]\mathcal L_{\text{critic}}=\mathbb E_t[(V(s_t)-\operatorname{sg}(G_t^\lambda))^2]

Aₜ is the advantage, the return minus what the critic expected. π(aₜ | sₜ) is the probability of the action the actor sampled in that state. Multiplying its negative log-probability by a positive advantage encourages that action, while a negative advantage discourages it. 𝔼 means average over sampled decisions. sg means stop-gradient, so the advantage is a fixed learning signal during this update.

V(sₜ) is the critic’s estimate and Gₜ with the λ superscript is the return target. Their squared difference measures the prediction error. 𝔼 averages it over sampled decisions. sg means stop-gradient, so the critic learns toward a fixed target rather than changing both sides of the comparison.

These are the core losses. An entropy bonus encouraged exploration. The targets were detached, and actor gradients flowed through action probabilities, not through the world model.

But fitting imagined returns does not make them right. I later checked the critic against discounted returns from real gameplay.

End of the first-draft details

Uh, it did not work

The first training run used roughly 20,000 steps. And it didn't even shoot.. a complete disaster. Across 30 games, it averaged only 0.03 points. Even a random policy scored 4.3.

The original Space Invaders actor moves right and left without firing, then loses with zero points.

My first doubt was whether the model was predicting the rewards the actor needed to learn from. I replayed the same actions in the real game and the model, then compared what actually happened with what it predicted.

Here is an example I took from the experiment.

Following the same shotReal frames above · decoder confidence below
t + 1
t + 2
t + 3
t + 4
t + 5
Real game
0
0
0
0
+1
World model
-0.02
-0.01
-0.06
1.19
-0.02

The model was quite off with its reward predictions. The model expected about 1.19 points on step four, but the actual shot earned one point on step five, when the model predicted almost nothing. It even predicted small negative rewards when nothing was earned. The actor was being given the wrong answer to learn from. D:

Across more trajectories, reward timing quickly deteriorated beyond one step into the future.

Step 1
72% 83/116
Step 3
18% 23/128
Step 5
19% 23/124
Step 8
1% 1/121

One step ahead, about 72% of the reward events were timed correctly. Eight steps ahead, that was only about 1%. The model could start from a real history, but its imagined future quickly became unreliable.

Those unreliable imagined rewards could also distort the critic's predictions. In one diagnostic, its average estimate was 10.46, while real discounted returns averaged only 3.16. So the actor could be learning from both misplaced rewards and an overly optimistic critic.

Changing how the model learned

The model could start from a real history, but struggled once it had to continue from its own predictions. I had two ideas for improving that experience: learn from the actor's actual gameplay and train on several predicted steps.

First, the model had only seen random gameplay, so it might struggle in situations our actor reached. I added the actor's real gameplay to replay so the model could learn from those mistakes.

01Play real games02Save trajectories03Update encoder + world model04Practise in imaginationactor + critic

repeat with the new policy ↻

Second, our earlier probe showed predictions falling apart further into the future. I tested whether training on several predicted steps, rather than one step from a real history, would help.

The idea was simple: feed each predicted latent back into the model to predict the next step, then compare the predictions against the real latents using the recorded actions. In this experiment, I fine-tuned the model for 600 updates on five-step rollouts.

Lrolloutlatent=error at step 1+error at step 2+⋯+error at step 5\mathcal L_{\text{rollout}}^{\text{latent}}=\text{error at step 1}+\text{error at step 2}+\cdots+\text{error at step 5}

In the separate five-step check, 600 rollout updates reduced latent prediction error by about 18% and reward prediction error by about 9%. Not bad! But better predictions didn't necessarily mean a better agent.

I combined online replay and rollout training and tried Space Invaders again. The run reached around 4 points in roughly 50k steps and... plateaued from there. 🫩

Space Invaders
AgentRandom
0246mean score50k100k150k200k250konline policy steps
Hover or tap a point to see its score

Better than 0.03, but still close to random. I needed to check whether the model was really what held it back.

Taking a little inspiration

The initial training changes hadn't been enough. I went back to older DreamerV2 concepts and used them alongside the diagnostics to rethink how the model represented the game and learned from experience.

Read the screen with shared filters. A projectile could appear in any column, so I replaced the MLP with a CNN. Instead of learning each position separately, it could reuse the same filters across the screen.

Flatten → MLP
Shared 3 × 3 filters

Practise the futures it would learn from. Our earlier diagnostics showed predictions drifting during imagination. I kept three-step generated rollouts so the predictor practised using its own outputs. I also sampled rare reward events more often, with importance weights correcting the sampling bias.

Initially one predicted step
Three steps checked against replay

Well, did it fix it?

The architecture looked more complete, but did it actually predict rewards better? I replayed identical actions through the early and rebuilt Space Invaders models to check. Using trajectories from 30 random-play games, both started with eight real frames and predicted eight steps.

Across 1,314 rollout windows, here is how often each model put a real reward event on the right step.

First draftWorking agent
Step 1
57% 53/93
87% 81/93
Step 3
22% 21/97
55% 53/97
Step 5
6% 6/95
42% 40/95
Step 8
0% 0/93
28% 26/93
  • Better timing. Five steps ahead, correctly timed events increased from 6 of 95 to 40 of 95.
  • Closer reward totals. Five-step cumulative reward error fell from 0.313 to 0.125, about 60% lower.
  • Still some imaginary points. False alarms at step five increased from 6 to 27. The rebuilt model found more real rewards, but also predicted more that never happened.

A big success! Now, you'd expect the imagined frames to look better too, right? Here is the same shot through both models, with the real game on top and each model's decoded predictions below. Hover over a heatmap to see its decoded frame.

Same shot, two checkpointsHover predictions to see decoded frames
t + 1
t + 2
t + 3
t + 4
t + 5
Real game
First draft
Working model

Interestingly enough, the first draft actually matched these particular frames more closely, even though the rebuilt model predicted rewards better across the larger test. A cleaner-looking imagined screen wasn't the same as a more useful reward prediction.

If you're interested, I also did a deeper dive into the model-training choices and their effects.

The final draft

Better predictions were only useful if the actor could learn from them. The rebuild also changed how the actor and critic used the model. These were already part of the checkpoint tested above; here's how they fit together.

  • Shared history. The actor and critic had only seen the current latent and a one-frame change. I gave them the transformer's summary of up to 16 frames and preceding actions, so both could use a longer history to choose actions and estimate future rewards.
  • An action MLP. Previously, the action went into the transformer. Now this MLP combines the history summary with the chosen action before predicting the next latent, reward, and continuation. The actor can choose from the history before the model predicts that action's consequences.
  • Real targets for the critic. I added real replay returns alongside imagined ones to help ground its optimistic estimates. It still guides the actor through return targets, shown by the dashed training path.
Working NanoDreamer architectureFrames pass through an encoder and history into a causal transformer and transition heads. The actor supplies the chosen action. The critic estimates future returns. Separate training summaries show reconstruction, prediction and behavior objectives. Green dots identify changes in the working version.01 / REPRESENTATION & PREDICTIONDecoderCNN · weighted BCEActorreads history features hₜFramesEncoderCNN · 64-D latent · jointLatent historyTransformer2 blocks · width 128 · causalAction MLPhₜ + aₜNext latentMSE + decoded-frame BCEReward255 bins · cross-entropyContinuationprobability · binary CEPrediction losscombined learning signalCriticdistributional V(hₜ)02 / BEHAVIOR TRAININGActor objective−log π(a | state) × advantageReturn targetsrewards + critic · λ-returnsCritic objectiveCE(real) + 0.3 × CE(imagined)Click a component to explore
Working NanoDreamer
Working NanoDreamer architectureFrames pass through an encoder and history into a causal transformer and transition heads. The actor supplies the chosen action. The critic estimates future returns. Separate training summaries show reconstruction, prediction and behavior objectives. Green dots identify changes in the working version.01 / REPRESENTATION & PREDICTIONDecoderCNN · weighted BCEActorreads history features hₜFramesEncoderCNN · 64-D latent · jointLatent historyTransformer2 blocks · width 128 · causalAction MLPhₜ + aₜNext latentMSE + decoded-frame BCEReward255 bins · cross-entropyContinuationprobability · binary CEPrediction losscombined learning signalCriticdistributional V(hₜ)02 / BEHAVIOR TRAININGActor objective−log π(a | state) × advantageReturn targetsrewards + critic · λ-returnsCritic objectiveCE(real) + 0.3 × CE(imagined)

Green dots mark changes from the first draft.

Inside the working versionArchitecture, losses, and imagined returns

First, describe the game in 64 numbers

The CNN encoder compresses each frame into 64 numbers. Shared 3 × 3 filters look for local patterns across the grid. The decoder reconstructs the observation channels using weighted binary cross-entropy.

Unlike the first draft, prediction losses also update the encoder. A slowly moving target encoder provides the real latent targets, so the predictor isn't chasing every encoder update immediately.

Open any equation to see what the symbols mean.

θˉE←0.99θˉE+0.01θE\bar\theta_E\leftarrow0.99\bar\theta_E+0.01\theta_E

θᴇ denotes the trainable encoder's weights. The bar marks the target encoder. Each update keeps 99% of its old weights and takes 1% of the current encoder's weights. Gradients do not update this target network.

Reconstruction checks what the latent preserves visually. Prediction losses give it feedback about the consequences the world model needs to learn.

Then, predict what an action does

The transformer reads up to 16 frames and preceding actions. It has two causal blocks, width 128 and four attention heads. A small MLP combines its summary hₜ with the current action before the three prediction heads.

The heads predict the next latent, a reward distribution over 255 bins, and continuation probability. Each generated step checks all three, plus the frame decoded from the predicted latent.

ℓk=Lnext latent+Lpredicted frame+2Lreward+0.2Lcontinuation\begin{aligned}\ell_k={}&\mathcal L_{\text{next latent}}+\mathcal L_{\text{predicted frame}}\\&+2\mathcal L_{\text{reward}}+0.2\mathcal L_{\text{continuation}}\end{aligned}

Next-latent loss is mean squared error against the target encoder's output. Predicted-frame loss is weighted binary cross-entropy. Reward loss is cross-entropy against a two-hot target, which splits a symlog-transformed reward between neighboring bins. Continuation uses binary cross-entropy. The coefficients balance contributions with different numerical scales.

Three generated steps feed predictions back into the model. A quarter of sampling draws come from reward events, with importance weights correcting that mixture. Padding is excluded from losses.

Lworld=Lreconstruction+Lone-step+13∑k=13ℓk\begin{aligned}\mathcal L_{\text{world}}={}&\mathcal L_{\text{reconstruction}}+\mathcal L_{\text{one-step}}\\&+\frac13\sum_{k=1}^{3}\ell_k\end{aligned}

Reconstruction checks real frames. The one-step term checks latent and decoded-frame predictions from real histories. The averaged three-step term checks futures generated through the model's own outputs. Target latents are detached.

Now, let the actor practise

The actor and critic read the transformer's history summary. Starting from replay, the actor chooses five imagined actions. The model predicts their rewards and continuation probabilities, and the target critic estimates what comes afterward.

Gtλ=r^t+γc^t[(1−λ)Vˉ(ht+1)+λGt+1λ]G5λ=Vˉ(h5)\begin{aligned}G_t^\lambda&=\hat r_t+\gamma\hat c_t[(1-\lambda)\bar V(h_{t+1})+\lambda G_{t+1}^\lambda]\\G_5^\lambda&=\bar V(h_5)\end{aligned}

G is the return target. Each step adds predicted reward r̂ to discounted future value. γ = 0.99 discounts the future; λ = 0.95 mixes the next critic estimate with the longer return. Continuation ĉ removes future value when the model expects an ending. Work backward from the final target-critic estimate V̄.

Actions with returns above the critic's expectation become more likely. A running return scale keeps the advantage comparable as rewards change, and a small entropy bonus encourages exploration.

At=sg⁡ ⁣(Gtλ−V(ht)max⁡(1,S))Lactor=−Et[wt(Atlog⁡π(at∣ht)+0.001H)]\begin{aligned}A_t&=\operatorname{sg}\!\left(\frac{G_t^\lambda-V(h_t)}{\max(1,S)}\right)\\\mathcal L_{\text{actor}}&=-\mathbb E_t[w_t(A_t\log\pi(a_t\mid h_t)+0.001\mathcal H)]\end{aligned}

A is the scaled advantage and S is the running return range. π is the sampled action's probability, ℋ is policy entropy, and w discounts later states using predicted continuation. sg stops gradients through the learning target. Actor gradients pass through action probabilities, not through the world model.

This recipe trains the critic on three-step real replay returns, bootstrapped with a target critic, plus a smaller imagined-return loss. The actor still learns entirely in imagination.

Lcritic=Lreal replay+0.3Limagined\mathcal L_{\text{critic}}=\mathcal L_{\text{real replay}}+0.3\mathcal L_{\text{imagined}}

Both terms use cross-entropy to fit a distribution over value bins to detached two-hot return targets. Real examples receive replay importance weights; imagined examples receive continuation weights. The target critic moves 2% toward the current critic per update.

End of the working-version details

After 4,000 initial random transitions, the actor keeps collecting real gameplay. Every eight policy steps, I update the world model, actor, and critic. Exploration falls from 40% random actions to a 5% floor, so the agent can still try something outside its current preference and learn from the result.

The moment we've been waiting for

The predictions were better. Now for the test that mattered—could an actor learning inside them play the real game? Here are the training runs for Space Invaders and Breakout.

Space Invaders
051015mean score50,000 steps · agent 5.81 · random 4.7750k100,000 steps · agent 10.52 · random 3.94100k150,000 steps · agent 11.14 · random 4.47150k200,000 steps · agent 12.87 · random 4.51200k12.874.51online policy steps
Breakout
01234mean score50,000 steps · agent 2.96 · random 0.3750k100,000 steps · agent 2.97 · random 0.29100k150,000 steps · agent 2.94 · random 0.24150k200,000 steps · agent 3.04 · random 0.44200k3.040.44online policy steps
AgentRandom

Finally, something was working. At 200k steps, Space Invaders scored nearly 3× random and Breakout nearly 7× random. But would the same recipe work across the rest of MinAtar?

With the algorithm working, I trained it for 500,000 policy steps on each MinAtar game, using separate weights but the same architecture and hyperparameters. Here are the resulting scores, averaged over 100 fresh games per agent.

Average game score · higher is better

GameRandomNanoDreamer
Space Invaders4.1422.61
Breakout0.375.03
Asterix0.401.50
Seaquest0.283.93
Freeway0.1511.23

And there it was. The same roughly 750k-parameter architecture, trained separately on each MinAtar game, beat the random baseline across the board. After the failed runs, misplaced rewards, and nights spent debugging, NanoDreamer was finally learning to play through imagination. :D

NanoDreamer playing Space Invaders
Space Invaders
NanoDreamer playing Breakout
Breakout
NanoDreamer playing Asterix
Asterix
NanoDreamer playing Seaquest
Seaquest
NanoDreamer playing Freeway
Freeway

Takeaways and next steps

NanoDreamer changed what I pay attention to when reading a paper. I started by looking at architecture diagrams. Now I'm much more interested in the assumptions, training data, and experiments behind the results.

It also taught me to be more careful about what an experiment actually tells me. Getting an agent to work and understanding why it works are different achievements. I'm happy we reached the first, but there's still work to do on the second.

For my next project, I want to spend less time recreating a complete system and more time investigating a specific question. What information does a latent representation lose? When does an imagined rollout become unreliable, and can the agent recognize that?

But I'm also tempted to just build more working algorithms and architectures through my favourite video games, like Smash Ultimate or Deadlock. We'll see where that takes me!

And lastly, thank you for reading! I really appreciate you making it this far. I spent way too many sleepless nights debugging this thing and pulled an all-nighter to finish the write-up, so it means a lot if you enjoyed it!

Additional notes

Extra experiments, hyperparameters, and my training and evaluation setup.

Game rewards

I used the rewards supplied by MinAtar, without custom shaping. There are no extra points for firing, moving, or surviving a step.

GameWhat earns rewardWhat ends the game
Space Invaders+1 for each alien shot.An alien reaches or hits the cannon, or an enemy shot hits it.
Breakout+1 for each brick broken.The paddle misses the ball.
Asterix+1 for each treasure collected.The player touches an enemy.
Seaquest+1 for each enemy shot. Surfacing with six divers adds a bonus based on remaining oxygen.An enemy or bullet hits the submarine, oxygen runs out, or it surfaces without a diver.
Freeway+1 for reaching the top of the road.The environment timer expires. A car hit resets the chicken to the bottom; it doesn’t end the episode.

Scores are sums of these rewards. The games have different scoring opportunities, so a score of 5 means something different in each game. Sampling reward events more often changes the training data mix; reward bins change how the model predicts a number. Neither adds rewards to the environment.

Our rollout limits can also end a recorded episode before the environment itself ends. Freeway’s own timer is longer than our 1,000-action final evaluation limit.

Working-version hyperparameters

These are the settings for the full rebuilt recipe behind the five-game results. Each game uses the same settings and architecture, with separate weights trained from scratch.

SettingValue
Latent size64
CNN channels16 → 32 → 64, with 3 × 3 filters
TransformerWidth 128 · 2 blocks · 4 attention heads
History context16 steps
Reward / value distributions255 bins, symlog range −5 to 5
World-model / real-return horizon3 steps
Actor imagination horizon5 steps
World-model batch32 replay contexts
Imagination batch64 starting contexts
Initial random experience4,000 environment actions
Update frequency1 world, actor, and critic update per 8 policy steps
Replay capacity200,000 transitions
Reward-event sampling mixture25% event starts + 75% uniform starts, importance-corrected
OptimizerAdam
World-model learning rate0.0002
Actor learning rate0.00004
Critic learning rate0.0001
Gradient norm clipping5 for each optimizer
Discount γ / return mixing λ0.99 / 0.95
Actor entropy coefficient0.001
Policy probability mixing99% learned probabilities + 1% uniform
Random-action exploration40% initially, falling to a 5% floor by 26,250 policy steps
Target encoder / critic update1% / 2% toward the current network per update
Generated-step loss weightsLatent 1 · frame 1 · reward 2 · continuation 0.2
Critic loss weightsReal replay 1 · imagined 0.3

The controlled comparisons change the named encoder, critic, or sampling choice. Their settings should not be confused with the full recipe above.

Training and evaluation details

A policy step means one action taken in the real environment. The 500k-step runs contain 500,000 policy actions, plus 4,000 initial random actions. Imagined transitions and gradient updates do not count toward that number. Each training episode allows up to 1,000 actions.

ResultEvaluation
Longer-run learning curves100 games per checkpoint · 200-action limit · different evaluation seeds at each checkpoint
Final five-game table500k training steps · 100 fresh games per game · 1,000-action limit
Training-choice comparison chartsMatched 150k training steps · 100 games · evaluation seed 920000 · 200-action limit
Training-choice endpoint checks300 games · evaluation seed 940000 · 200-action limit; one run ended at 184,712 training steps

The agent samples its learned action probabilities during evaluation, without the additional training-time random-action exploration. Random baselines use the same game, evaluation seeds, and action limits as the corresponding agent check.

The five-game results have one training seed per game. The additional comparisons include two seeds for some recipes and one for others. More evaluation games reduce uncertainty about a saved policy; they don’t establish reliability across training seeds. The 32-point clip is a selected episode, not the average.

Additional training comparisons

After the longer runs, I went back to check which training signals helped the policy. I compared versions with one signal removed, then followed up on the promising alternatives. The actor learned entirely in imagination throughout.

All comparisons below use 150k real policy steps and the same 100 evaluation games, with a 200-action limit. That keeps the training budget equal, including the run that later hit its time cap.

Did prediction errors need to train the encoder?

My original encoder learned to reconstruct frames. Reconstructing a projectile didn't necessarily make its next collision easy to predict, so the full rebuild also let prediction losses train the encoder.

To test that choice, I trained another version where the CNN learned only reconstruction. The architecture stayed the same, and the encoder still learned online. It just stopped receiving gradients from the prediction losses.

GameJoint encoderReconstruction only
Space Invaders9.3211.26
Breakout2.922.98

In this first seed, reconstruction-only training scored higher in Space Invaders and almost the same in Breakout. So I couldn't credit joint encoder training for the improvement. A representation learned through reconstruction was still competitive in this setup.

Did the critic benefit from real rewards?

Remember the critic expecting more reward than the real game gave? In the full rebuild, I tried grounding it with real replay returns alongside imagined targets. That sounded useful, but the actor was practising inside the model. I wanted to check whether the mixed targets actually helped.

I compared that critic with one trained entirely on imagined returns, then repeated both versions with a fresh training seed. Each dot below is a separately trained agent.

150k policy steps Seed 0 Seed 1

Space Invaders

01020Full recipeFull recipe · seed 0 · 9.32 pointsFull recipe · seed 1 · 11.16 pointsImagination-only criticImagination-only critic · seed 0 · 15.82 pointsImagination-only critic · seed 1 · 15.14 points

Breakout

036Full recipeFull recipe · seed 0 · 2.92 pointsFull recipe · seed 1 · 1.48 pointsImagination-only criticImagination-only critic · seed 0 · 4.82 pointsImagination-only critic · seed 1 · 1.40 points
Hover or tap a dot to see its score

In Space Invaders, imagination-only critic training improved the score in both seeds. In Breakout, it helped in the first seed but not the second. Giving the critic real targets wasn't a consistently better choice.

What if I combined the simpler choices?

That left a follow-up worth trying. What happened if the encoder learned only reconstruction and the critic learned only in imagination? I tested that combination in the second seed.

GameFull recipeImagined critic+ Reconstruction encoder
Space Invaders11.1615.1412.59
Breakout1.481.405.08

In Breakout, the combination reached 5.08, while the other two versions stayed around 1.4–1.5. In Space Invaders, it beat the full version but stayed behind the imagination-only critic with a jointly trained encoder.

So the changes weren't simply additive. I had promising simpler alternatives, but no single recipe had established itself as the most reliable across games. The combined version had only one seed; these comparisons measured the effects of the training choices, not the mechanism behind them.

The architecture detours

I tried more capacity because the predictor might not have been able to track objects and delayed rewards. Increasing transformer width from 64 to 128 and blocks from 2 to 3 grew the world model from roughly 109,000 to 613,000 parameters. Some reward metrics improved, but gameplay didn't improve consistently.

I also tried a small RSSM for recurrent memory. Spatial-token experiments kept four regions separate instead of compressing a frame into one token, and a separate denoising experiment refined a noisy next-state guess over several steps. Those were attempts to preserve local detail and model more complicated futures, inspired by DreamerV4.

These trials didn't isolate which architecture was best. They used different checkpoints and training conditions, and some runs were short. The detailed log preserves them, but they don't establish a single architectural cause for the early failures.

Inside the small RSSM trial

RSSM stands for recurrent state-space model. Our trial used a 64-number recurrent memory and 8 groups of 8 categorical choices. A GRU-like transition updated the memory from the previous stochastic state and action. A prior predicted the next stochastic state from memory, while a posterior could also use the real observation to correct that prediction.

The decoder, reward head, and continuation head read the combined memory and stochastic state. Training used image reconstruction, reward and continuation losses, and balanced KL losses to bring the prior and posterior together. In imagination there was no new observation, so only the prior was available. This was a compact DreamerV3-inspired trial, not the original DreamerV3 model size.

There was also a replay bug in one RSSM trial. Replay required eight-step windows and discarded shorter episodes, including the actor's six-step failures. In one random dataset, 231 of 300 episodes were too short. The rebuilt replay keeps short terminal windows and masks padding, so those failures remain training examples.

After 2,000 updates on 1,000 random episodes, the Space Invaders prior missed all 73 one-step reward windows in the held-out probe, even though the observation-conditioned loss had improved. That was a failure of this implementation and training setup, not evidence that small RSSMs cannot work.

References