@benbye: I trained a 750K-parameter world-model agent on my M5 Pro to play video games! Presenting NanoDreamer: a tiny Dreamer-i…
Summary
Benjamin Bai unveils NanoDreamer, a 750K-parameter Dreamer-inspired world-model and actor-critic agent trained entirely on a MacBook M5 Pro, complete with his own autograd engine, tensor library, and transformer, playing all 5 MinAtar games from imagination.
View Cached Full Text
Cached at: 10/29/26, 04:03 PM
I trained a 750K-parameter world-model agent on my M5 Pro to play video games!
Presenting NanoDreamer: a tiny Dreamer-inspired agent that learns game dynamics and trains an actor–critic in latent imagination.
Write-up: https://t.co/bXoODifBDX https://t.co/BIHsI6t0HK
NanoDreamer: Learning to Act Through Imagination
Source: https://www.benjamin-bai.com/projects/nanodreamer Benjamin BaiOctober 2026GitHub
IntroductionHow hard could it be?Uh, it did not workChanging how the model learnedActor–critic checkTaking a little inspirationWell, did it fix it?The final draftThe moment we’ve been waiting forTakeaways and next stepsAdditional notesBuild it from scratch
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 playingSpace 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 aseparate 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. Aworld modellearns how the game works, anactorchooses actions inside its predictions, and acriticestimates 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 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 anMLP autoencoderthat compressed each frame into a 64-dimensional latent embedding. Once it could reconstruct the screen, I froze its weights.
Asmall causal transformerread the latent embeddings alongside action and positional embeddings. Its three prediction heads estimated the next latent state, reward, and probability the game would continue. Theactor and criticwere MLPs that saw the current latent concatenated with its change from the previous frame.
The actor then trained through5-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 usedweighted binary cross-entropyto make missing an occupied cell more costly.
Open any equation to see what the symbols mean.
Lreconstruction=−meani[wixilogx^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 only0.03 points. Even a random policy scored4.3.
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.
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.
Step1
72%83/116
Step3
18%23/128
Step5
19%23/124
Step8
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 was10.46, while real discounted returns averaged only3.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 about18%and reward prediction error by about9%. 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 around4 pointsin roughly50k stepsand... plateaued from there.
0246mean score50k100k150k200k250konline policy stepsHover 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 olderDreamerV2 conceptsand 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 from6 of 95 to 40 of 95.
- Closer reward totals.Five-step cumulative reward error fell from0.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.
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 dida 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 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 stepsAgentRandom
Finally, something was working. At 200k steps, Space Invaders scored nearly3× randomand Breakout nearly7× random. But would the same recipe work across the rest of MinAtar?
With the algorithm working, I trained it for500,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

Space Invaders
Breakout
Asterix
Seaquest
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.
Want to build it from scratch?
Follow the foundations, from algebra and gradients to the networks behind the agent.
Similar Articles
I trained a 1.57B-parameter Dreamer 4 World Model from scratch for under $150
The author trained a 1.57B-parameter Dreamer 4 world model from scratch for under $150, demonstrating the accessibility of low-cost AI model training and using procedural data for validation.
Local AI World Model Part 2 - Deep NN to turn Images into Playable Characters, with prompt switching mid rollout
A solo researcher shares Part 2 of a local AI world model: a ~960M-parameter transformer using diffusion forcing and text cross-attention that runs at 30-60 FPS on an RTX 5090, accepts live WASD actions, and supports mid-rollout text prompt switching (e.g., 'add a pond to the desert').
@MaximeRivest: wow moondream!! 9b params and 2b active. Did I mention I love small specialized models?
Users hype Moondream 3.1, a small vision-language model with 9B total params and 2B active params, praising it for outperforming Rex-Omni and LocateAnything on image tasks while being cheaper, fast, and locally runnable.
A tiny world model running games on a Raspberry Pi Pico 2 W
tinyworldformer runs two tiny neural networks—a looped-transformer world model and an MLP actor—on a Raspberry Pi Pico 2 W to generate and play simple games like Breakout, Pong, and Flappy Bird entirely on-device, with optional on-board RL learning via stream-x.
@heyrobinai: THE ENTIRE AI INDUSTRY JUST GOT HUMILIATED a tiny model trained in just a few hours on a single graphics card is planni…
Yann LeCun's team releases LeWorldModel, a tiny 15M-parameter physics model trained on a single GPU in hours that outperforms billion-dollar foundation models in planning speed and physical plausibility, challenging the dominant scaling paradigm.