World Models · Challenge 1

The Shifting Maze

A Dyna agent has learned its world perfectly. Then the world changes. Will it notice?

~45 minpredict, run, design

5 partsA → E, getting harder

Hand iniPad photo or text

Same Dyna-Q as the chapter: ε = 0.1, γ = 0.95, +1 for reaching G. One change: the world is deterministic, so α = 1. The agent trusts each observation completely, which is also why its model is so sure of itself.Same Dyna-Q agent as in the chapter, two new mazes. In each, the wall moves partway through training, and you have to predict what the agent does about it.

Part A · warm-upThe blocking maze

Figure C1.1.The agent learns the 10-step route on the left. At step 1,000 its gap is walled up and a new gap opens at the far end. Dyna-Q runs with $n = 10$ planning steps.

Predict: which curve shows Dyna-Q's cumulative reward over 3,000 steps? (The dashed line marks the change.)

B. D is a trap: cumulative reward can never go down, because the rewards are all ≥ 0. After the change, the agent keeps walking into the new wall. Each bump corrects one entry of its model, but its Q-values still point at the old route, so it has to "unlearn" its way out, square by square. This is an optimistic error (the model believes in a gap that has gone), and optimistic errors fix themselves. It just takes time. Press Lock in and watch it happen.

Picked a curve?

Part B · the main eventThe shortcut maze

Figure C1.2.Now the reverse. The agent learns the long 16-step route. At step 3,000 a second gap opens: a 10-step shortcut. Nothing gets blocked; the world only gets better. Dyna-Q runs with $n = 50$ planning steps.

Predict: within the remaining 3,000 steps, does Dyna-Q ($n=50$) discover and start using the shortcut?

Run it and see. The why is Part B's written question below.

Committed?

B · writtenExplain●●●

In 2–3 sentences: why does more planning ($n = 50$) not help here, when it helped so much in the chapter's race? Use the words "model" and "data" in your answer.

No hidden solution here. Write it down and send it to me.

Part C · designDesign a fix

You're Rich Sutton. Your agent is stuck with a pessimistic model that will never correct itself. Change the algorithm so it finds the shortcut, while still exploiting what it has learned. Write your design down before opening any hint. One rule: you can't just raise ε. That explores everywhere, forever, and you pay for it on every single step.

Hint 1 · a question back

Which (state, action) pairs is the model least entitled to be confident about, and is there a number the agent could keep track of that measures that?

Hint 2 · a nudge

The world might change. A transition you tried 5 steps ago is probably still accurate. One you last tried 3,000 steps ago is stale. Staleness grows with time.

Hint 3 · the key idea

Make stale transitions look attractive in imagination. Since the agent optimises its imagined reward, it'll plan routes that go and check stale places. You're turning pessimism into optimism on purpose, in exactly the places where the model might be out of date.

Hint 4 · Sutton's answer (Dyna-Q+)

Track $\tau(s,a)$, the number of real steps since $(s,a)$ was last actually tried. During planning only, use the bonus-augmented reward

$$\tilde r \;=\; \hat r + \kappa\sqrt{\tau(s,a)}.$$

Actions never tried from a visited state are allowed in planning too, modelled as "stay put, reward 0" plus the bonus. The real reward stays untouched. Only the dream is tilted toward curiosity.

Now test it. The simulator below runs plain Dyna-Q against Dyna-Q with your bonus $\kappa$. Find a κ that discovers the shortcut without wrecking performance before the change.

Part D · numericalHow stale is stale?

DNumerical●●●

Take Dyna-Q+ with $\kappa = 10^{-3}$ and γ = 0.95. The agent is at the start state, where its learned value is that of the long route: the goal is 16 steps away, so $V(s) = \gamma^{15} \approx 0.463$.

One action from this state, "up", has never been tried. In planning, it's modelled as "bump the wall, stay in $s$, reward 0", plus the bonus.

How many real steps $\tau$ must pass before planning makes "up" the greedy action in $s$?

Hint (only if stuck)

The naive answer compares the bonus with $V(s)$ directly, and gives about 214,000 steps. It's wrong. What does Q-learning's target look like for an action whose imagined next state is the same state?

Solution

The imagined transition loops back to $s$, so its planning target is $b + \gamma\max_{a'}Q(s,a') = b + \gamma V(s)$, with $b = \kappa\sqrt{\tau}$. The action becomes greedy when that exceeds $V(s)$:

$$b + \gamma V > V \iff b > (1-\gamma)V = 0.05 \times 0.463 \approx 0.0232.$$

So $\sqrt{\tau} > 23.2$, giving $\tau \gtrsim 536$ steps. That's 400× sooner than the naive answer. The self-loop means the agent gets its own value back for free, so the bonus only has to beat the discount loss $(1-\gamma)V$, not $V$ itself.

The general lesson: in RL, a small per-step bonus is effectively multiplied by the horizon $1/(1-\gamma)$. Remember this when you meet entropy bonuses in Dreamer.

Part E · frontierToward Dreamer

EOpen●●●

Dreamer works on pixels, where (from Chapter 1) you essentially never see the same frame twice.

(i) Why does the $\kappa\sqrt{\tau(s,a)}$ bonus stop making sense there?
(ii) Propose a replacement signal that a neural world model could compute, which is high exactly where the model is likely to be wrong.
(iii) What's one way your signal could fool itself?

Compare with the literature (after you've written yours)

(i) $\tau(s,a)$ is a count over exact states, and with continuous, never-repeating states every count is 0 or 1.

(ii) One well-known answer is ensemble disagreement (Plan2Explore, Sekar et al. 2020). Train several dynamics heads, and use the variance of their predictions as an intrinsic reward inside imagination. Where they disagree, the data was thin. Others: prediction error (curiosity, Pathak et al. 2017) and learned density models (pseudo-counts).

(iii) Prediction error is fooled by noise: a TV showing static is forever unpredictable, so a curious agent will stare at it (the "noisy-TV problem"). Disagreement handles this better, because all ensemble members learn to predict the same average, so their disagreement shrinks even where the noise doesn't.

Plain DreamerV3 uses neither. It relies on an entropy bonus on its actor, and on the fact that it keeps retraining its model on fresh data. We'll see later whether that's enough.

Send me (photo of your iPad page or plain text):

Your Part A and B predictions, before running anything, and whether they were right.
Part B written answer.
Part C: your own design, and the deepest hint rung you opened (or none).
The κ range that worked in the simulator.
Part D number and Part E (i)–(iii).