The promise, and the trap inside it
In 3.1 you did something that should still feel slightly illicit: you learned a
simulator. An encoder, a GRU, a prior, a decoder — a little network that, given a
state and an action, dreams the next state, no mj_step inside. And you measured,
honestly, that it learned the easy half (the pusher, whose next position is
basically the integral of the velocity you commanded) and not the hard half (the
T-block's own pushed-around dynamics, where copy-last still beat it).
Here is the move this chapter is named for. If you have a simulator you can roll forward cheaply and differentiably, you don't need the real environment to learn a policy — you can learn one inside the dream. Freeze the world model, imagine trajectories, score their reward, and push a gradient back into an actor. The real PushT sim is never touched during policy learning. This is Dreamer, and it is one of the most seductive ideas in the field: learn to imagine, then learn to act in imagination, and pay for almost no real experience.
The trap is right there in the setup, and this chapter is built to make you measure it. The reward for PushT depends on the block — get the T to the target. But your world model got the block wrong. So the actor is about to optimize a reward computed inside a dream whose hard half is a hallucination. It may learn to look like a champion in imagination and then fail completely in reality. That gap — imagined success minus real success — is the entire lesson. We are going to measure it and refuse to look away from it.
Reward without a reward head
# The reward is a KNOWN function of the state (PushT: drive the T-block to the# target pose), so we do NOT learn a reward head — we compute it analytically from# the DECODED observation. That is the honest crux: in imagination this reads the# block pose the world model HALLUCINATED, the very dims 3.1 measured it gets wrong.# The same function scores the real trajectory in eval, so the gap is pure state# divergence, not two different reward definitions.POS_SCALE, ANG_SCALE = 0.5, float(np.pi) # from PushTEnv.step: reward = -0.5*(pos_err/0.5 + ang_err/pi) def reward_from_obs(norm_obs): """Shaped PushT reward from a normalized observation (any leading batch shape). De-normalize, then read tee_xy and tee_yaw (target is the fixed origin).""" obs = norm_obs * std_t + mean_t pos_err = torch.linalg.vector_norm(obs[..., 2:4], dim=-1) # ||tee_xy - target(0,0)|| ang_err = torch.abs(torch.atan2(obs[..., 4], obs[..., 5])) # |wrap(tee_yaw - 0)| return -0.5 * (pos_err / POS_SCALE + ang_err / ANG_SCALE), pos_err, ang_errDreamer learns a reward head — a network that predicts reward from the latent —
because in a pixel game it has no other access to it. We don't need one: the PushT
reward is a known function of the state (drive the block's pose to the origin), so we
compute it analytically from the decoded observation. That is not a shortcut past
the hard part — it is the hard part, in plain sight. The reward reads tee_xy and
tee_yaw straight out of the frame the world model imagined. When the imagined block
is wrong, the imagined reward is wrong, and the actor optimizes the wrong thing. A
learned reward head would inherit exactly the same poison from exactly the same
place; making the reward analytic just puts the poison where you can see it.
The actor and critic live on the model's state
# The actor and critic both read the world-model STATE [h, z] — the same features# available in a dream AND on the real robot (there we recover [h,z] by filtering# the posterior on real frames). A tanh-squashed Gaussian keeps actions in the# env's [-1, 1] and stays reparameterized, so the imagined return differentiates# straight through the (frozen) dynamics into the policy — Dreamer's analytic# actor gradient, no REINFORCE estimator needed at this scale. class ActorCritic(nn.Module): def __init__(self, state_dim, act_dim, hidden_dim): super().__init__() self.actor = nn.Sequential(nn.Linear(state_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, act_dim)) self.log_std = nn.Parameter(torch.zeros(act_dim)) self.critic = nn.Sequential(nn.Linear(state_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, 1)) def act(self, state, sample=True): """Return (action in [-1,1], base-Gaussian entropy). sample=False acts on the mean — the deterministic policy used for BOTH imagined and real eval.""" mean = self.actor(state) std = torch.exp(self.log_std) pre = mean + std * torch.randn_like(mean) if sample else mean entropy = (0.5 + 0.5 * np.log(2 * np.pi) + self.log_std).sum() return torch.tanh(pre), entropy def value(self, state): return self.critic(state).squeeze(-1)The actor and critic both read the world model's state [h, z] — the deterministic
recurrent state and the latent. That choice matters for a reason that only pays off
at deployment: [h, z] is available in a dream (roll the prior) and on the
real robot (filter the posterior on real frames). Same features, both worlds, so the
same policy runs in each. The actor is a tanh-squashed Gaussian, which keeps actions
in the env's [-1, 1] and stays reparameterized — so the imagined return can
differentiate straight through the (frozen) dynamics into the policy. That is
Dreamer's analytic actor gradient; at this scale we need no REINFORCE estimator.
Step 1: learn the simulator (unchanged from 3.1)
# Step 1: learn the simulator, exactly as in 3.1. Reconstruction + dynamics loss.torch.manual_seed(args.seed) # model init reproducible, independent of data collectionwm = WorldModel(OBS_DIM, ACT_DIM, args.latent_dim, args.hidden_dim, args.embed_dim).to(device)wm_opt = torch.optim.Adam(wm.parameters(), lr=args.wm_lr)global_step = 0for epoch in range(args.wm_epochs): epoch_recon, num_batches = 0.0, 0 for batch in torch.randperm(len(train_obs), generator=shuffle_gen).split(args.batch_size): recon_loss, dyn_loss = wm.observe(train_obs[batch], train_act[batch]) wm_opt.zero_grad() (recon_loss + dyn_loss).backward() wm_opt.step() epoch_recon += recon_loss.item() num_batches += 1 if args.rerun: rr.set_time("step", sequence=global_step) rr.log("wm/recon", rr.Scalars([recon_loss.item()])) rr.log("wm/dyn", rr.Scalars([dyn_loss.item()])) global_step += 1 if epoch % 15 == 0 or epoch == args.wm_epochs - 1: print(f"[wm] epoch {epoch:3d} recon {epoch_recon / num_batches:.5f}")Nothing new here — it is 3.1's world model, trained on scripted-expert sequences with the same reconstruction-plus-dynamics loss. We spend the fewest words on it because you already built it; it is the substrate, not the point. When it is trained, we freeze it.
Step 2: learn a policy inside the frozen dream
# Step 2: FREEZE the world model and learn a policy INSIDE it. Every gradient the# actor sees comes from a dream — start from real filtered states, then roll the# frozen prior forward under the actor's own actions, decode, and score the reward.# The lambda-return (Dreamer) mixes the imagined rewards with the critic's bootstrap;# the actor maximizes it by analytic gradient, the critic regresses onto it.for p in wm.parameters(): p.requires_grad_(False) # frozen: gradients flow THROUGH it to the actions, never INTO itstate_dim = args.hidden_dim + args.latent_dimagent = ActorCritic(state_dim, ACT_DIM, args.hidden_dim).to(device)agent_opt = torch.optim.Adam(agent.parameters(), lr=args.actor_lr)ctx_obs = train_obs[:, : args.context + 1] # frames to warm the start state fromctx_act = train_act[:, : args.context] def lambda_return(rewards, values): """Dreamer's lambda-return: G_t = r_t + gamma*((1-lam)*V_{t+1} + lam*G_{t+1}), bootstrapped at the horizon by the critic. rewards: (H,B); values: (H+1,B).""" out, last = [], values[-1] for t in reversed(range(len(rewards))): last = rewards[t] + args.gamma * ((1 - args.lam) * values[t + 1] + args.lam * last) out.append(last) return torch.stack(out[::-1]) # (H, B) imag_reward_hist = float("nan")for it in range(args.imag_iters): idx = torch.randint(len(train_obs), (args.imag_batch,), generator=shuffle_gen) with torch.no_grad(): h, z = wm.warm(ctx_obs[idx], ctx_act[idx]) # real start states; detached from the actor graph states, rewards, entropies = [torch.cat([h, z], dim=-1)], [], [] for _ in range(args.imag_horizon): action, entropy = agent.act(states[-1], sample=True) h, z, obs_pred = wm.step(h, z, action) rewards.append(reward_from_obs(obs_pred)[0]) entropies.append(entropy) states.append(torch.cat([h, z], dim=-1)) # Stop-gradients keep the two objectives from cross-contaminating (canonical Dreamer): # the critic regresses V(state) onto the return, so its input states are detached (it # must not reshape the actor); the actor maximizes the return, so the critic bootstrap # inside it is detached (it must not inflate the critic). The actor's gradient reaches # it only through the imagined REWARD path. One optimizer is fine once they no longer cross. values = agent.value(torch.stack(states).detach()) # (H+1, B) — grad to the critic only returns = lambda_return(rewards, values.detach()) # (H, B) — the reward path carries the actor's analytic grad actor_loss = -(returns.mean() + args.ent_coef * torch.stack(entropies).mean()) critic_loss = 0.5 * F.mse_loss(values[:-1], returns.detach()) agent_opt.zero_grad() (actor_loss + critic_loss).backward() nn.utils.clip_grad_norm_(agent.parameters(), 100.0) agent_opt.step() imag_reward_hist = torch.stack(rewards).mean().item() # mean per-step imagined reward under current actor if args.rerun: rr.set_time("imag_iter", sequence=it) rr.log("imag/reward_per_step", rr.Scalars([imag_reward_hist])) rr.log("imag/critic_loss", rr.Scalars([critic_loss.item()])) if it % max(args.imag_iters // 5, 1) == 0 or it == args.imag_iters - 1: print(f"[imag] iter {it:4d} imagined reward/step {imag_reward_hist:+.4f} " f"critic {critic_loss.item():.4f}")This is the imagination loop. Sample real start states, filter them to (h, z), then
roll the frozen model forward for imag_horizon steps under the actor's own
actions — decode each dreamed state, score its reward, and stop. That sequence of
imagined rewards and critic values becomes a lambda-return (Dreamer's
bias/variance-controlled target); the actor maximizes it by analytic gradient, and the
critic regresses onto it. Every gradient the actor ever sees comes from a dream. The
one honest caveat in the code: the analytic gradient through a frozen GRU is
high-variance, so the learning rate is deliberately small (--actor_lr 1e-4) — a
larger one diverges on some seeds. On the default config this whole chapter runs in
~0.22 min on a CPU laptop (measured; T4 not yet measured). It is deliberately
tiny — the lesson is the mechanism and the gap, not a competitive agent.
And it works, in the sense that matters for the setup: the imagined return rises. Across seeds the actor learns to drive the imagined block from its ~0.17 m spawn down to about 0.01 m from the target — in the dream, it parks the block. The policy is learning. Hold that thought.
The headline, measured: the imagination gap
# The headline. Deploy the SAME deterministic policy in two worlds and compare.# IMAGINED: warm (h,z) on the start frame, then roll the prior forward under the# actor for eval_horizon steps — a pure dream, decoded, reward from the decode.# REAL: filter (h,z) on true frames each step, act, step the true PushT sim. The# reward uses the SAME function on the TRUE state. Success is the env's own flag.# If imagined return >> real return, the actor learned to please a dream whose block# dynamics 3.1 measured are wrong. If they agree, pushing transferred. Either is the# lesson; we do not tune until the story is clean.agent.eval()imag_returns, real_returns, real_success = [], [], []imag_final_pos, real_final_pos = [], []for ep in range(args.eval_episodes): seed = EVAL_SEED_BASE + args.seed + ep env = PushTEnv() obs0 = env.reset(seed) norm0 = torch.tensor((obs0 - obs_mean) / obs_std, device=device, dtype=torch.float32).unsqueeze(0) with torch.no_grad(): # IMAGINED rollout: one-frame warm, then dream forward h, z = wm._init(1) z = wm.posterior(torch.cat([h, wm.encoder(norm0)], dim=-1)) dream_r, dream_pos = [], None for _ in range(args.eval_horizon): action, _ = agent.act(torch.cat([h, z], dim=-1), sample=False) h, z, obs_pred = wm.step(h, z, action) r, pos_err, _ = reward_from_obs(obs_pred) dream_r.append(r.item()) dream_pos = pos_err.item() imag_returns.append(float(np.mean(dream_r))) imag_final_pos.append(dream_pos) with torch.no_grad(): # REAL rollout: filter on true frames, act in the true sim h, z = wm._init(1) obs, real_r, done = obs0, [], False for _ in range(args.eval_horizon): norm = torch.tensor((obs - obs_mean) / obs_std, device=device, dtype=torch.float32).unsqueeze(0) z = wm.posterior(torch.cat([h, wm.encoder(norm)], dim=-1)) action, _ = agent.act(torch.cat([h, z], dim=-1), sample=False) real_r.append(reward_from_obs(norm)[0].item()) obs, _, done, info = env.step(action[0].cpu().numpy()) h = wm.gru(torch.cat([z, action], dim=-1), h) if done: break final_norm = torch.tensor((obs - obs_mean) / obs_std, device=device, dtype=torch.float32) real_final_pos.append(reward_from_obs(final_norm)[1].item()) real_success.append(float(info["success"])) real_returns.append(float(np.mean(real_r))) imag_return = float(np.mean(imag_returns)) # mean per-step reward the actor BELIEVES it earnsreal_return = float(np.mean(real_returns)) # mean per-step reward it ACTUALLY earnsgap = imag_return - real_return # >0 => imagination is rosier than realitysuccess_rate = float(np.mean(real_success))print(f"\n=== the imagination gap ({args.eval_episodes} held-out starts, " f"{args.eval_horizon}-step rollouts) ===")print(f"IMAGINED return/step (in the dream) : {imag_return:+.4f} final tee-dist {np.mean(imag_final_pos):.3f} m")print(f"REAL return/step (true PushT) : {real_return:+.4f} final tee-dist {np.mean(real_final_pos):.3f} m")print(f"gap (imagined - real) : {gap:+.4f} " + ("<- imagination is DELUDED: the policy earns its reward on a hallucinated block" if gap > 0.02 else "<- imagined and real roughly agree at this scale"))print(f"real task success rate : {success_rate:.2f} " "(the shaped reward can rise while success stays ~0 — the block never actually parks)")if args.rerun: rr.log("eval/imagined_return", rr.Scalars([imag_return])) rr.log("eval/real_return", rr.Scalars([real_return]))Now deploy the same deterministic policy in two worlds and compare. Imagined: warm the state on the start frame, roll the prior forward under the actor, score the decoded reward — a pure dream. Real: filter the state on true frames, act, and step the real PushT sim, scoring the same reward function on the true state. Measured across seeds 0-2 at the default config:
- Imagined return/step: ~-0.18. Final imagined block-to-target distance ~0.01 m — the dream is convinced the policy solved the task.
- Real return/step: ~-0.39. Final real block-to-target distance ~0.16 m — the block barely moved from where it spawned.
- The gap is +0.17 to +0.24 return/step, and it is positive on every seed. Real task success rate: 0.00.
Read that again. The policy did not fail to train — it trained beautifully, in imagination. It learned to move the pusher and to sweep the imagined block to the target. But the world it trained in is wrong exactly where the reward looks, so the skill it acquired is skill at pleasing a hallucination. In the real sim the pusher moves and the block does not follow the way the dream promised, and the T never parks. Imagination is only as good as your world model — and yours, from 3.1, learned the pusher and not the block. This is not a bug to fix in 3.2; it is the measured, honest sequel to the finding you already made.
A longer dream is not a better plan
The obvious hope is that a longer imagination horizon — letting the actor dream
further before it scores its plan — would hand it a better plan and rescue the real
performance. It does neither, and the investigation exercise has you measure exactly
why. Sweep imag_horizon from 5 to 30 and two numbers refuse to move: the dream parks
the block at every horizon (imagined final block-distance ~0.01–0.02 m whether the
actor imagines 5 steps or 30), and reality stays exactly as stuck (real block-distance
~0.16 m, real success 0.00 at every horizon). The delusion is horizon-invariant.
Rolling a wrong model further does not buy real skill and does not deepen the
delusion either — the model is simply wrong about the block wherever you read it, at
step 5 and at step 30 alike. The ceiling is the world model's block dynamics, and no
length of rollout raises it. That is why the thing that would actually close this gap
is a better world model — which, per 3.1, costs compute: the Scale Lab, and the
map's "why world models eat compute" thesis applied now to acting.
Report and what carries forward
metrics = { "eval_episodes": args.eval_episodes, "eval_horizon": args.eval_horizon, "final_imagined_reward_per_step_train": round(imag_reward_hist, 6), "gamma": args.gamma, "hidden_dim": args.hidden_dim, "imag_horizon": args.imag_horizon, "imag_iters": args.imag_iters, "imagination_gap": round(gap, 6), "imagined_final_tee_dist": round(float(np.mean(imag_final_pos)), 6), "imagined_return_per_step": round(imag_return, 6), "latent_dim": args.latent_dim, "real_final_tee_dist": round(float(np.mean(real_final_pos)), 6), "real_return_per_step": round(real_return, 6), "real_success_rate": round(success_rate, 6), "seed": args.seed, "smoke": bool(args.smoke),}(args.out / "metrics.json").write_text(json.dumps(metrics, indent=2, sort_keys=True) + "\n")print(f"\nmetrics: {args.out / 'metrics.json'} (imagination gap {metrics['imagination_gap']:+.4f} " f"return/step; real success {success_rate:.2f})")if args.rerun: print(f"recording: {args.out / 'dreamer.rrd'} — open it with: rerun {args.out / 'dreamer.rrd'}")The metrics record the imagined and real return, the gap, the block-to-target distance in each world, and the real success rate — the numbers the exercises reproduce. You have now built the full Dreamer skeleton by hand: a world model (3.1) and an actor-critic that learns inside it (3.2). And you have measured the honest boundary of the idea — that a policy is only ever as good as the dream it was raised in. That boundary is not a discouragement; it is the map. When you read the real DreamerV3 (see "read the real thing"), the machinery that closes this gap — a far better world model, a learned reward and continue head, a stochastic latent with KL balancing — is exactly the machinery you can now see the need for, because you measured what its absence costs.
What we cut
- The learned reward and continue heads. Real Dreamer predicts reward and episode-continuation from the latent; we compute reward analytically and imagine a fixed horizon. This is honest here (the reward is a known function of state) and it puts the world model's error in plain view, but it means we skip the machinery that lets Dreamer act on pixels with no state access.
- The stochastic latent and KL balancing (cut in 3.1, still cut here) — so the imagined rollout has no calibrated uncertainty, which is one thing a real Dreamer would use to avoid over-trusting a shaky dream.
- A world model good enough to act in. The deepest cut, and the point. We act inside the deliberately-tiny 3.1 model, measure the gap, and name the fix (a better model = more compute = the Scale Lab) rather than paying for it free-tier.
Read the real thing
meta.yaml pins danijar/dreamerv3 at e3f0224. Three files hold the loop you
just hand-built — read them against it.
imagine(): rolling the prior under the actor. Our imagination loop is
dreamer.py#train_actor: warm (h, z) from real frames, then step the frozen model
forward under agent.act for imag_horizon ticks, collecting each decoded reward.
Production does exactly this in dreamerv3/rssm.py, method imagine — it nj.scans
the prior transition (self._prior(deter), no posterior, no observation) forward
under a policy callable, the same "roll the prior on the actor's own actions" move.
The wrapper is dreamerv3/agent.py's imag_loss, which calls
self.dyn.imagine(starts, policyfn, H, ...) with policyfn sampling the actor. What
they add: the rollout is a jitted scan over batched latents starting from every
real step in a replay buffer, refilled as new experience arrives — where ours dreams
off a handful of warmed starts, theirs dreams off thousands, continuously.
Lambda-returns, and where the actor gradient comes from. Our lambda_return in
dreamer.py#train_actor is the textbook recursion G_t = r_t + γ((1-λ)V_{t+1} + λG_{t+1}), and our actor maximizes it by an analytic gradient — the reparameterized
tanh-Gaussian differentiates straight through the frozen GRU. Production keeps the
lambda-return (dreamerv3/agent.py, lambda_return — the same recursion with a
continue-masked discount) but does not use that analytic dynamics-gradient. Its
imagine stop-gradients the state into the policy (policy(sg(carry))), and
imag_loss trains the actor by REINFORCE: policy_loss = -(logpi * sg(adv_normed) + actent * ents), advantages normalized by a running scale. This is the one place our
chapter and DreamerV3 genuinely part ways: DreamerV1 backpropped through the dynamics
as we do; v3 switched to the score-function estimator because it survives discrete
actions and long unrolls where the pathwise gradient explodes — the very
high-variance we bought down with actor_lr 1e-4. Same target, different way of
pushing it into the policy.
The learned reward head vs our analytic reward. Our dreamer.py#reward reads
tee_xy/tee_yaw straight out of the decoded frame — no reward head, because PushT's
reward is a known function of state. Production can't do that: on pixels it has no
state, so dreamerv3/agent.py learns one, self.rew = embodied.jax.MLPHead(...),
trained by losses['rew'] = self.rew(inp, 2).loss(obs['reward']), alongside a
continue head self.con predicting episode termination (we imagine a fixed horizon
instead). But swapping analytic for learned changes nothing about the lesson you
measured: both read reward off the world model's imagined state, so a model wrong
about the block poisons either one. The imagination gap is a property of the dream,
not of where the reward comes from — which is exactly why rolling further (the
horizon investigation) never closed it.
Read next: open dreamerv3/agent.py, find imag_loss, and read down to the
lambda_return call right inside it. That method is the entire actor-critic-in-imagination
loop you built by hand — read it knowing the return is computed the same way you
computed it, and only the actor's gradient path differs.