DEV Community

Devanshu Biswas
Devanshu Biswas

Posted on

World models: let an agent dream rollouts inside a learned simulator — until compounding prediction error makes the dream drift

An agent that only learns by acting for real is expensive: every trial burns time, money, wear, and sometimes safety, and reinforcement learning is famously sample-hungry — millions of steps. A world model is the escape hatch: a learned function that captures the environment's dynamics, so given a state and an action it predicts the next state. Once the agent carries a fast, private copy of the world in its head, it can plan and "dream" thousands of rollouts inside the model — cheap, fast, safe — instead of paying to act in reality. I built a demo with real state-transition math and a real drift metric to show both the power and the catch. Here it is.

A world is just (state, action) → next state

The environment in the demo is a ball bouncing in a box under gravity, nudged by actions. Its true dynamics are a real hand-coded physics step — apply the action, gravity, drag, integrate, reflect off the walls with restitution. This is the slow, costly real world.

function physStep(s, a, p){                  // ONE parameterised physics step
  let {x,y,vx,vy} = s; const {ax,ay} = ACT(a);
  vx += ax*p.imp;  vy += ay*p.imp;           // action = velocity impulse
  vy -= p.g;  vx += p.vbias;                 // gravity (+ systematic bias)
  vx *= p.drag; vy *= p.drag;                // drag
  x += vx; y += vy;                          // integrate
  if (x<R){x=R;vx=-vx*p.e;} if (x>1-R){x=1-R;vx=-vx*p.e;}   // wall bounce
  if (y<R){y=R;vy=-vy*p.e;} if (y>1-R){y=1-R;vy=-vy*p.e;}   // (restitution e)
  return {x,y,vx,vy};
}
Enter fullscreen mode Exit fullscreen mode

The learned model is the same shape, imperfect parameters

A trained world model never recovers the true physics exactly — it has a slightly wrong idea of gravity, bounciness, drag, and how much an action does. I model that honestly: the same physStep, but with parameters offset by an err knob. At err = 0 the model equals the true env; crank it up and it's confidently wrong.

function modelParams(err){                    // err in [0,1]
  return {
    g:    TRUE.g    * (1 + 0.45*err),         // thinks gravity is stronger
    e:    TRUE.e    * (1 - 0.18*err),         // thinks walls are less bouncy
    drag: TRUE.drag - 0.004*err,              // thinks there's more drag
    imp:  TRUE.imp  * (1 - 0.35*err),         // underestimates its own actions
    vbias: 0.0018*err                         // a small systematic sideways bias
  };
}
Enter fullscreen mode Exit fullscreen mode

Dreaming is an autoregressive rollout

To dream, give both the same start and the same action sequence. Step the true env with trueStep. Step the model with modelStep — and crucially feed the model its own previous prediction (si), never the true state. That autoregressive feedback is exactly how an agent imagines a future it hasn't lived.

function rollout(s0, actions, err, N){
  const trueT=[s0], imagT=[s0], errs=[0];
  let st=s0, si=s0;                           // st = reality, si = imagination
  for (let k=0;k<N;k++){
    const a = actions[k];
    st = trueStep(st, a);                     // what really happens
    si = modelStep(si, a, err);               // dream: prediction feeds itself
    trueT.push(st); imagT.push(si);
    errs.push(Math.hypot(st.x-si.x, st.y-si.y)); // drift at this step
  }
  return { trueT, imagT, errs };
}
Enter fullscreen mode Exit fullscreen mode

The catch: error compounds, and bounces amplify it

Because each prediction becomes the input to the next, a tiny error at step 1 shifts step 2's input, whose error shifts step 3, and errors don't stay small — they accumulate, roughly geometrically. Worse, near a wall a small position error flips whether the ball has bounced yet, after which the two trajectories head opposite ways. In the demo you slide model error and rollout length and watch the dashed (dreamed) path peel away from the solid (true) one; a bounce blows the gap wide open.

The honest metric is the drift per step, from which I read the reliable horizon — the last step before the dream crosses a tolerance. That horizon is the real answer to "how many steps can I plan on?"

function metrics(errs, tol=0.06){
  let horizon = errs.length - 1;
  for (let k=1;k<errs.length;k++) if (errs[k] > tol){ horizon = k-1; break; }
  return { finalDrift: errs[errs.length-1], horizon };
}
Enter fullscreen mode Exit fullscreen mode

Why this matters now

One learned dynamics function, two uses: plan by dreaming, and train a controller on imagined experience instead of costly real steps (Ha & Schmidhuber's Vision-Memory-Controller; Dreamer; MuZero-style latent planning). It's wildly sample-efficient — but only within the horizon. Real systems fight the drift with uncertainty and ensembles, periodic re-grounding on real data, and rollouts in a compact latent space rather than raw pixels. And it's why large video-generation models like Sora are increasingly read as implicit world models — they learned intuitive physics and object permanence just to predict the next frame. The craft of world models is buying as many trustworthy imagined steps as you can before reality and imagination part ways.

Dream a rollout and watch it drift:
https://dev48v.infy.uk/ai/days/day42-world-models.html

Top comments (0)