← Back to portfolio

Shepherd's World Model

Two flocks, one shepherd, one fast wolf. The shepherd does not know where the wolf goes. It watches, counts what the wolf does, plans with what it has learned, and guards the sheep. That is a world model: a small copy of how the world behaves, learned from experience and used to plan. The loop is section 19.3 of Stanford's CS229 notes (Ma & Ng). Made for World Models Start in Six Lines of CS229.

1 · Pick the wolf's secret habit. The shepherd is not told which one you picked.

🐑 two flocks🧑‍🌾 shepherd🐺 wolf■ what the shepherd thinks the wolf will do next

2 · Let the nights run.

Wolves turned away

0 of 0

Night 0. The shepherd knows nothing about this wolf yet, so it guesses that every move is equally likely.

How good is the learned model?

A dream rolls the wolf forward using only what the shepherd has learned, not the real habit. Early dreams are wrong. Later ones look like the real wolf. That is the idea behind Ha & Schmidhuber's "World Models" (2018), in miniature.

3 · The math, exactly as the shepherd runs it

  1. Count (CS229 eq. 19.5): P̂(wolf goes w → w′) = #times the wolf moved w → w′ ÷ #times the wolf was at w. Never seen? Uniform over the wolf's legal moves (the notes say uniform over all states; a wolf cannot teleport).
  2. Plan (value iteration, §19.2): V(s) ← R(s) + γ · maxa Σs′ P̂sa(s′) V(s′) over every (shepherd cell, wolf cell) pair: 64 × 64 = 4,096 states, up to 5 actions, γ = 0.95. The wolf takes two steps per turn, so the shepherd must be where the wolf is going, not where it is.
  3. Act: move to the neighbouring cell with the best value. Repeat every night, counts accumulating (the algorithm on CS229 p. 233–234).

Rewards: +10 when the wolf is turned away (it steps onto the shepherd, or the shepherd steps onto it), −10 if it reaches a flock, −0.1 per step otherwise, and 0 if the night ends after 30 steps.

Learned vs true: where the wolf goes from the cell it is in now