Dreams That Last — the RSSM on a Real Robot
Last time our world model could see one step ahead — and then its dream fell apart. Today we rebuild it so the dream holds, sixty steps deep, on a real robot arm.
The previous pod ended on an honest failure. Our MiniPong world model predicted the next frame almost perfectly — but the whole point of a world model is to run it in a closed loop, its own output becoming its next input, and there it degraded fast. Within a handful of imagined steps the paddle still obeyed, but the ball dissolved into a smear and then into nothing.
Today we rebuild it so the dream holds — and we do it on a real robot arm, with real recorded data. The idea we need is the RSSM, the Recurrent State-Space Model from the PlaNet paper (Learning Latent Dynamics for Planning from Pixels, Hafner et al., 2019) — the state design behind every Dreamer agent since.
The Dream That Fell Apart
Three things went wrong last time. Errors compound: a tiny mistake at step 5 is eaten at step 6 and amplified at step 7, and nothing pushes the model back toward reality. The encoder never knew its job: trained first, alone, to redraw frames, it built a code space that decodes beautifully and predicts badly. And we patched it by hand: the "smoothness" fix, forcing consecutive frames to nearby codes, was us guessing at a principle. Today we meet the real one.
The uncomfortable diagnosis ties all three together: every one of our fixes was a training fix applied to a design problem. A loss can only push around the numbers a model already carries. If those numbers are the wrong shape — if there is nowhere to put "which way is the ball going" or "the cube is behind the hand" — no amount of training will invent that place. So most of this lecture is about one question: what should the state be made of?
And there is a deeper failure hiding inside the smear. When MiniPong's ball approached the paddle, its future genuinely forked: bounce left, or bounce right. A model forced to output one state, trained under squared error, has exactly one safe move: predict the average of the possible futures — a frame that never occurs in the real world. A single vector has no way to express width; certain and uncertain look identical, so the model cannot say "I don't know yet." The fix, stated now and earned over the rest of this article: a world model should not track a state. It should track a belief — a distribution over states — wide when it is guessing, narrow when it knows.
A Real Robot, One Row at a Time
Our subject is the SO-101 robot arm doing a pick-and-place, from the openly available lerobot/svla_so101_pickplace dataset: 50 tele-operated episodes, 11,939 frames. Every timestep stores exactly three things — a camera frame, six joint angles, and the six joint commands that were sent.
Look at the sizes. The state of the arm is six numbers — shoulder pan, shoulder lift, elbow flex, wrist flex, wrist roll, gripper: the arm's entire configuration. The observation is a 64×64 image: 12,288 numbers pointed at those six, plus a cube, a table, and lighting the robot does not control. A world model's job is to recover the small thing from the huge one.
But some of the state lives in neither column. Is the cube gripped, or just touched? No joint angle says so, and at the wrong camera angle no pixel does either. At the grasp moment the gripper closes over the cube, and for roughly half a second the cube is behind the hand — in neither the frame nor the joint angles. The model has to hold it. And "held" is not one number: was the cube nudged as the fingers closed? Did it settle a millimetre left? The honest answer is not a position but a region — narrow if the grasp looked clean, wider if it did not. That is a belief, on real data.
There is a catch: the dataset contains pixels and joint commands, not beliefs. Nobody ever wrote down "at frame 40 the model should be 70% sure the cube is held." So where does the belief come from? To answer that, we build the two designs anyone would try first — and watch precisely how each one fails on this robot.
Design A: A Memory That Cannot Choose
Start with the design you already own — the MiniPong model drawn with today's shapes, now trained jointly so the comparison is fair to it. An encoder squeezes each frame down to 1,024 numbers. A GRU carries a single state vector forward, updating it from (that vector, the frame, the action). A decoder paints the next frame from the state. Every arrow is a computation; nothing is ever sampled. Hand it the same history twice and it produces the identical future twice.
Before we break it, let us be honest about what this design is good at — with a measurement, not an adjective. People say a recurrent state "carries information forward" and move on. We can do better: open the vector and ask it questions with a linear probe, about ten lines of code. Run the trained model over 25 real episodes, recording the state and the six true joint angles at every timestep. Fit a single linear map — no hidden layers — from state to joint angles, and as a control, fit the same readout against shuffled labels.
The result: 99.9% of every joint angle is recovered by a plain linear readout, against a shuffled control that scores essentially zero. The model was never given a joint angle as a target; it built that representation because predicting frames required it. And crucially, those readouts keep tracking the real robot 40 steps into a camera-less rollout. Nothing is refreshing them; they are being carried. This is the property worth protecting: lose the arm during a one-second occlusion and the model is useless on a robot.
Now ask Design A a question with two answers. Freeze the recording at the instant the gripper closes on the cube: what happens next? The grasp catches, or the cube slips. Design A must pick one future, so the loss makes it split the difference — half-gripped, half-slipped, exactly MiniPong's smear, now on a robot. Dreamed 60 steps ahead, the hedging compounds: the arm softens, the scene drifts, and the joint error reaches 0.409, forty-five times worse than the design we finish with. Averaging is not caution; it is a confident wrong answer.
Design B: Dice That Cannot Remember
So make the state honest about forks: have the network output a distribution — a centre and a width — and draw the state from it, fresh, every step. This is the paper's own stochastic baseline, and it really does fix the fork. We froze the trained model at the grasp moment and asked for six futures from one identical starting state and identical actions. The six answers differ — the gripper values genuinely spread across the draws — and each future is sharp, not mushy. An average of two futures is a blur; a sample of one future is a clean, physically plausible frame you could act on.
But Design B bought honesty at a fork and lost something Design A had for free. If the state is redrawn at every timestep, every fact must be re-transmitted through a fresh sample, forever. Fifteen frames of occlusion is fifteen consecutive dice rolls the cube's position has to survive. It won't. The noise you added for honesty erodes exactly what you were trying to remember.
Measurement makes it worse than the argument suggests: trained on the same data and budget, Design B's joint error rose during training (0.26 → 0.39) while its KL cost climbed past 90 — paying more and more to re-transmit what it should simply have kept. Dreamed 60 steps, its joint error reaches 2.598, the worst of the three by a wide margin.
Two designs, two ceilings, opposite directions: a perfect memory that cannot choose, and honest dice that cannot remember. Neither is a world model you would trust for sixty steps.
The RSSM: One State, Two Halves
The RSSM's whole idea in one line: stop choosing. Carry both, and give each the job the other cannot do.
At every timestep the model carries one state with two parts. h — 256 numbers updated by a GRU, with no noise ever applied: the deterministic memory, the belt that carried the arm through 40 camera-less steps. It is where facts live. s — 32 numbers drawn fresh from a distribution the model computes from h: where the six futures came from, where doubt lives, and the only place randomness enters the whole model. The frame is painted from neither half alone: the decoder literally receives the concatenation [h, s] — 256 + 32 = 288 numbers, one vector — so every prediction reflects both the facts and the branch the model currently believes it is on.
One timestep runs in three stages. Stage 1: the GRU updates the memory from (previous h, previous s, the action just taken) — no randomness and, notice, no camera frame. This stage alone is Design A. Stage 2: the updated memory feeds a small network that outputs a distribution — a centre and a width — and the new s is a draw from it. Stage 3: the decoder paints the frame from h and s together. Then the loop repeats.
And here is the crucial wiring detail, the mechanism that makes the pair more than the sum: that sampled s flows straight back into stage 1 of the next step. Suppose at the grasp moment the dice draw "the cube is held." That sample enters the GRU on the very next step — it stops being a guess and becomes part of the memory. Every later step is computed from a memory that contains "the cube is held": the arm lifts, and the cube goes with it. The model never contradicts a choice it already made. Call it commit, then remember. Design B could never do this — with no memory to write into, a sampled choice evaporates at the next step. Design A could never make the choice at all — it would have averaged "held" and "slipped", then faithfully carried the smear forward for sixty steps.
Is the pairing really load-bearing, or just two useful features side by side? We sharpened the test: take one trained RSSM and run it three ways — freeze the sample, or starve the memory of it — same weights, same episode, only the wiring changed. It degrades exactly the way the three separate models did. Cut either link and the dream falls apart.
One more structural payoff, the direct answer to MiniPong: our old model decoded from only half its split state, so dreaming one step meant a round trip out to a code and back — sixty round trips for a sixty-step dream, each one leaking error. Here the dream never leaves the state: state → state → state, internally. We decode a picture only when a human wants to watch.
Training Without Labels
We have the machine; nobody has told us what it should contain. Nowhere in those 11,939 rows is there a column saying "at this moment the state should be these 288 numbers," and nobody could write one — the state is the model's own private summary. So the training signal has to be built from the only two things we actually have: the frames we recorded, and the model's own second opinion. Each becomes one loss term.
Loss 1 — repaint the frame. Bolt a decoder onto the state, paint a 64×64 image, compare it pixel by pixel with what the camera saw. This term forces the state to contain the world — the only way to repaint the arm, the cube and the table is to have kept them. But it is not enough: MiniPong had this term and still failed, because a state can be perfectly repaintable and impossible to predict forward. Loss 1 fills the state; it does nothing to make it predictable.
Loss 2 — ask the same question twice. There is no label for "predictable," so we manufacture one. Ask the model once with the frame: "given the memory, the action, and the camera image you just received — what is the state now?" That is the posterior; it has all the evidence. Ask again without the frame — that is the prior; it must guess blind, exactly the situation it will face when it dreams. Loss 2 is the distance between those two answers. Minimise it, and you are training the blind guess to match the informed one — and there is no way to win that game except to actually carry forward whatever determines the next frame. Which is the definition of predicting. The model grades itself.
One subtlety is where the real magic sits: loss 2 measures a gap between two networks, so the optimiser can close it by moving either one. It improves the prior — the predictor learns to carry the cube, the contact, the velocity. But it also disciplines the posterior: perception is punished for encoding anything the prior could never anticipate, and pushed toward representations that are predictable in the first place. That is precisely what our hand-rolled "smoothness" penalty was groping toward — here it is not a patch, it falls out of asking one question twice.
Two rules make this work. The camera frame enters the model on exactly one path — the posterior — and the prior never touches it; break that asymmetry and both guessers just copy the frame and learn nothing about time. And both losses fire at every timestep — 32 times per training sequence, for 16,000 steps. The model is asked "could you have predicted this?" millions of times, and every answer becomes a gradient. MiniPong, which trained a predictor afterwards on frozen codes, was never once asked that question.
The Verdict: Three Designs, Sixty Steps
Now the honest question: does it work? Take five pick-and-place episodes the model has never seen, give it five frames to warm up, then switch the camera off, hand it only the joint commands that were actually sent, and ask it to imagine the next sixty steps: two full seconds of robot motion. The model is 7.7 million parameters — convolutional encoder, the 256-number memory, the 32-number stochastic state, a decoder, and a small head that reads the six joint angles back out of the state — trained on 45 tele-operated episodes for roughly 50 minutes on a single GPU.
Because arguments are cheap, we trained all three designs ourselves — same 45 episodes, same 16,000 steps, same seed, matched parameter counts — and ran the identical 60-step test:
| design | parameters | pixel error @60 | joint error @60 |
|---|---|---|---|
| deterministic only | 7,718,313 | 0.0106 | 0.409 |
| stochastic only | 7,543,881 | 0.0155 | 2.598 |
| both — the RSSM | 7,665,993 | 0.0050 | 0.009 |
Each design fails exactly the way the argument predicted. And the combination is not a small improvement: 0.009 is 45× better than deterministic-only and 288× better than stochastic-only, at matched parameter counts. This is not tuning. It is the state design.
Note the pixel column, too: it separates the three far less than joint error, because pixels are dominated by the static table and background while the robot is where the designs differ. Choose the metric that measures what you care about.
The dreamed trajectories are worth staring at. All six joints track the real robot — including the sharp, discrete event: the gripper signal drops as the hand closes on the cube, and the dreamed curve drops at the same timestep, without ever seeing a frame of it. The joint error runs about 0.002–0.02 from step 1 through step 60 and does not climb; on one episode the error at step 60 is lower than at step 30 — the dream re-converges toward reality instead of running away. That flat curve is the entire promise: a model you can run forward, on its own, long enough to plan with.
What We Did Not See — and What Comes Next
One claim we expected to make, and could not. We did see the belief story measured on a real robot: at the very first timestep, with no history, the model's uncertainty is at its maximum — and one frame collapses it. But we expected uncertainty to spike at contact, the moment the future genuinely forks, and averaged over the state's dimensions it does not. Either the forks live in a few dimensions the average washes out, or fifty episodes of a reliably successful grasp simply contain little to be unsure about. We are telling you because the alternative is a plot that flatters the story — and "the experiment did not confirm this" is itself a finding.
Four things to keep. A world model tracks a belief, not a state — a model that cannot express doubt invents confident nonsense. Prediction is taught by a tug-of-war — pulling the guess-without-looking toward the guess-with-looking is the single term that turns an autoencoder into a world model. Keep the two paths separate — a belt for facts, dice for forks; mix them into one channel and you lose both. And train the pieces together — a representation learned without the prediction task will not survive it.
To build this yourself — the two-half state on real SO-101 data, the linear probe into the memory, and the 60-step camera-off dream — the companion build lives in the Production Lab as Dreams That Last.
You now have a model that can imagine two seconds of a real robot's future from actions alone. The next question is the one this series has been walking toward: what if the robot practised in there? Next time — Dreamer: freeze this world model, put an agent inside it, and train it entirely on imagined rollouts. Does skill learned in a dream transfer to the real world? We will test it, and report what we find.