Dream to Catch
A from-scratch rebuild of IRIS: a tokenizer turns each frame into 16 tokens, a transformer dreams the next frame and its reward, and the policy trains only on imagined games. It touches the real environment for the first time at test — and catches every ball. Including the honest failure that almost sank it.
The build, step by step
The exact path we followed — and the one you'll follow in the code kit.
Three small parts
A VQ-VAE tokenizer (0.8M params), a transformer world model (3.3M), and an actor-critic (0.8M). Five million parameters, end to end.
Train the dream
Every 64×64 frame becomes 16 discrete tokens. The transformer learns to predict the next tokens, the reward, and whether the episode ends — a language model for gameplay. The policy then trains purely inside its imagination.
The failure
Every metric looked healthy — reconstruction near-perfect, next-frame accuracy 98%. Yet the policy stayed exactly random. The tokenizer had collapsed to 3 of its 256 codes and quietly deleted the ball from every frame. A policy cannot learn to catch a ball that doesn't exist in its dream.
The fix
EMA codebook updates with dead-code revival, and one decisive change: weight the reconstruction loss toward the ball. Ball recall went to 1.00 on every seed.
Awake
On 500 fresh real games the imagination-trained policy caught the ball every single time — 1.00, matching a hand-coded oracle, against a 0.16 random baseline.
From the run
