AI-NATIVE SERIES · WORLD MODELS FROM CS229, EPISODE 1 OF 1 · AGE 20

World Models Start in Six Lines of CS229: Learn the Wolf, Guard the Flock (for Age 20)

By Paul Jialiang Wu · agentic-portfolio-lovat.vercel.app · 2026-10-16 · 18 min read

The one question: What is a world model small enough to build, run and understand, and how does it lead to Dreamer, MuZero and Genie?

This piece does NOT teach the whole CS229 course or train a neural world model; it builds the tabular one in section 19.3 and maps the road up. The family drill it hands to children belongs to Shepherd's Academy of War, part 5.

What this essay is about: world models, starting from six lines

This is an essay about world models: a machine's small, learned copy of how its world behaves, used to plan before it acts. It is written for engineers of about twenty who want to do research on world models, not just read about them. You need a browser, Python, and about half an hour.

Many labs are betting that robotics, games and physical AI will run on world models. NVIDIA's Cosmos platform puts the case plainly: a physical AI system needs "a digital twin of the world, the world model" (NVIDIA (2025), Cosmos World Foundation Model Platform for Physical AI, arXiv 2501.03575). Most explanations start at the big end, with 11-billion-parameter video models. This one starts at the small end, with six lines in the Stanford CS229 lecture notes by Tengyu Ma and Andrew Ng (Ma & Ng, CS229 Lecture Notes (Stanford, Spring 2026)). In section 19.3, "Learning a model for an MDP", those six lines contain a complete world model: watch, count, plan, act, repeat.

The essay was built with playbook-math, from a hand-picked slice of a 278-page set of notes: about 30 pages. Each was scored on three things: how strongly it pulls an engineer toward research, how much it matters for world models, and whether a family could play it.

The running thing it was built on is the Shepherd's World Model, a free page where a shepherd learns a wolf's secret habits in order to guard two flocks. The picture comes from Shepherd's Academy of War, a family academy where children learn how an attack works so that they can protect. The wolf is an animal, and no person is the enemy.

Activity 1 · Invent the shepherd's method first

The six lines

CS229 section 19.3 starts from a Markov decision process (MDP): states, actions, rewards, and transition probabilities that say where each action tends to lead. Usually you are not given those probabilities. You must estimate them from experience, and the notes give the estimate as a ratio (eq. 19.5). Each transition probability, the chance of landing in s′ after taking action a in state s, is the number of times that happened, divided by the "#times we took action a in state s" (Ma & Ng, CS229 Lecture Notes, eq. 19.5 (p. 233)). A pair never tried gets a uniform guess.

Then the notes put learning and planning together in a loop of a few lines:

1. Start with any policy.
2. Repeat: - (a) run it for some trials; - (b) update the transition counts; - (c) run value iteration on the estimated model; - (d) switch to the greedy policy for the new values.

Step (c) is the planning half, so here it is plainly. The value V(s) of a state is how much reward you can expect from there on. Value iteration repeats one update until the numbers stop changing (§19.2): V(s) ← R(s) + γ · max over actions a of Σ over next states s′ of P̂_sa(s′) · V(s′). Here γ (0.95 in the demo) discounts the future. A policy is a rule for which action to take in each state, and the greedy policy simply takes the action with the best value. On night 1, the shepherd that knows nothing believes every wolf move is equally likely, so it guards every cell equally, which means it guards none.

That is the whole world model. A world model is a small copy of how the world behaves, learned from what you saw and used to plan before you act. This one has a memory of what happened (the counts), a model built from that memory (the ratios), and a planner that uses the model before acting (value iteration).

Activity 2 · Watch the shepherd learn

With seed 1, the shepherd that never learns plans forever with "every move is equally likely". Here is how it compares with the learning shepherd, over its last 30 of 60 nights. That is the second press: the score after 60 nights minus the score after 30, divided by 30. (Computed; open the page with seed 1 to reproduce.)

The last row is the lesson inside the lesson. A wanderer has almost no habit to learn, so the world model has almost nothing to add. Against it, a shepherd that learns does about as well as one that assumes every wolf move is equally likely and never updates (93% against 88%). Learning cannot add a pattern the wolf does not have. A world model is only as useful as the structure in the world it models. That sentence is worth a research career. In my view, much of the argument about whether large models "have world models" is really an argument about how much structure their training data forced them to capture. That is an opinion, not a finding.

The road up: from six lines to world foundation models

Later work keeps some lines of this loop and changes others. Kalman's filter assumes the model is given and learns nothing. Genie learns a simulator but does not plan. MuZero learns only the parts planning needs. The table says which part each one keeps, and the caveat each still carries.

Belief instead of position.

• The shepherd sees the wolf. Most robots do not; they see noisy clues.
• CS229 section 20.4 (LQG) handles this: keep a belief over the hidden state with a Kalman filter, then control with LQR. Kalman's 1960 paper states the deep link in one sentence: "The filtering problem is shown to be the dual of the noise-free regulator problem." (Kalman (1960), A New Approach to Linear Filtering and Prediction Problems, J. Basic Eng. 82).
• Estimating and acting turn out to be mirror images. That is the kind of surprise that makes people become researchers.

Learn, then plan, interleaved.

• Sutton's Dyna (1991) is this loop as an architecture, decades before these notes. In his words: "Planning is incremental and can use the probabilistic and ofttimes incorrect world models generated by learning processes." (Sutton (1991), Dyna, ACM SIGART Bulletin 2).
• "Ofttimes incorrect" is the honest heart of the field. The shepherd's early tables are wrong too.

Dream inside the model.

• Ha and Schmidhuber (2018) replaced the table with neural networks and wrote: "We can even train our agent entirely inside of its own hallucinated dream generated by its world model, and transfer this policy back into the actual environment." (Ha & Schmidhuber (2018), World Models, arXiv 1803.10122).
• In the demo, value iteration on the learned table is exactly that: the shepherd trains its plan inside its own model of the wolf. The dream button just shows one such imagined night, where the grey wolf moves the way the shepherd thinks it moves. A dream is only as good as the model: the shepherd's early dreams are wrong.

Learn only what planning needs.

• MuZero (2020) "learns an iterable model that produces predictions relevant to planning" (Schrittwieser et al. (2020), MuZero, Nature 588): policy, value and reward, not a full picture of the world.
• It did not learn the rules of chess. It learned what it needed to win.

One recipe for many domains.

• DreamerV3, published in Nature in 2025, used one world-model recipe across very different tasks. It was described as the first algorithm, "to our knowledge", to collect diamonds in Minecraft from scratch without human data or curricula (Hafner et al. (2025), Mastering diverse control tasks through world models, Nature 640).
• Keep the hedge when you repeat it.

Foundation world models.

• Google DeepMind's Genie has "11B" parameters and is described by its authors as a foundation world model (Bruce et al. (2024), Genie, arXiv 2402.15391).
• Genie 3 stays consistent for "a few minutes" in real time (Google DeepMind (2025), Genie 3 blog), not forever.
• OpenAI presented video generation as "a promising path towards building general purpose simulators of the physical world" (OpenAI (2024), Video generation models as world simulators, abstract as given on the MIT CSAIL seminar page). That is a research bet, not a proven world simulator, and the video model takes no actions.

The CS229 notes also hold the two engines inside these big models:

• the variational autoencoder (section 11.5; Kingma and Welling, arXiv 1312.6114, Kingma & Welling (2013), Auto-Encoding Variational Bayes), which compresses what the agent sees into a small code;
• diffusion models (section 14; Ho, Jain and Abbeel, arXiv 2006.11239, Ho, Jain & Abbeel (2020), Denoising Diffusion Probabilistic Models), which generate the frames of video world simulators.

They are the next rungs, not this one.

Activity 3 · Write the counting half yourself, then finish it

import random
random.seed(1)
N, SHEEP = 7, 6                       # a corridor of 7 cells; the sheep are in the last one
def wolf_step(w):                     # the TRUE habit, hidden from the shepherd
    return min(w + 1, SHEEP) if random.random() > 0.2 else max(w - 1, 0)
counts = [[0] * N for _ in range(N)]  # counts[w][w2]: how often the wolf went from w to w2
visits = [0] * N
def p_hat(w, w2):                     # CS229 eq. 19.5; never seen -> uniform over the wolf's moves
    if visits[w] == 0:
        return 0.5 if abs(w2 - w) == 1 or (w in (0, SHEEP) and w2 == w) else 0.0
    return counts[w][w2] / visits[w]
for night in range(200):              # watch the wolf and count
    w = 0
    while w != SHEEP:
        w2 = wolf_step(w)
        counts[w][w2] += 1
        visits[w] += 1
        w = w2
for w in range(SHEEP):
    print(w, "-> forward", round(p_hat(w, w + 1), 2), " back", round(p_hat(w, max(w - 1, 0)), 2))
E = [0.0] * N                         # plan: how many steps until the wolf reaches the sheep? (the shepherd's deadline)
for _ in range(500):                  # value iteration on expected steps-to-sheep under P-hat
    E = [0.0 if w == SHEEP else 1 + sum(p_hat(w, w2) * E[w2] for w2 in range(N)) for w in range(N)]
print("expected steps for the wolf to reach the sheep from cell 0:", round(E[0], 2))

It prints learned forward probabilities between 0.77 and 0.84 for a true value of 0.8, and an expected 9.32 steps from the first cell (computed with seed 1). Be clear about what it is. The counting is CS229's eq. 19.5. The last loop evaluates the learned chain, which is the planning step with no choice in it. Give the shepherd actions (stay, step toward the wolf) and take the best one each round, and it becomes value iteration: that is your exercise, and it is the demo's engine. Then break it on purpose:

• change the 0.2 to 0.5;
• give it 5 nights instead of 200;
• make the wolf's habit depend on the night.

Each change is a real research question in miniature. How much data does a world model need? What happens when the world drifts? What does planning on a wrong model cost?

Why this is where research starts

The most inspiring pages in CS229 for a would-be researcher are not about world models at all. They are section 8.2, on double descent. For decades the textbook picture said that past a point, bigger models overfit and test error rises. Then measurements showed the opposite past a further point. Belkin and colleagues reported that "increasing model capacity beyond the point of interpolation results in improved performance" (Belkin, Hsu, Ma & Mandal (2019), Reconciling modern machine-learning practice and the classical bias-variance trade-off, PNAS 116 (arXiv 1812.11118)). Nakkiran and colleagues titled their follow-up Deep Double Descent: Where Bigger Models and More Data Hurt (Nakkiran et al. (2019), Deep Double Descent, arXiv 1912.02292), because more data could even make things worse for a while.

That is the research move in its plainest form: run the thing, measure, be surprised, rebuild the theory. The shepherd demo is small enough to do it tonight. The wanderer row above is already a surprise worth explaining.

Same story, five exits

Truth, love and beauty in this essay

• Truth: every number reproduces. The demo is seeded, the code prints its results, and each claim about a famous paper carries its own hedge ("to our knowledge", "a few minutes", "a promising path").
• Love: the essay hands you a child to teach and a family drill to run. It meets the half-right beliefs about world models with what they got right.
• Beauty: six lines hold the whole idea. Later world models keep one or more of those lines and swap the table for a network billions of times bigger.

One lens is my own, and I label it as mine. Shepherd's Academy of War takes its picture from David's report to Saul: "Your servant was keeping his father’s sheep; and when a lion or a bear came, and took a lamb out of the flock, I went out after him, and struck him, and rescued it out of his mouth…" (1 Samuel 17:34–35, WEB). A shepherd who learns the predator's habits in order to protect the flock is, in my faith, a picture of love that does its homework. The checklist above stands without that belief.

Mistakes that are half right

• "A world model is a video generator." Some world models generate video, but the CS229 one is a table of counts and a planner. The defining feature is planning with a learned model, not pictures.
• "Bigger world models always help." Only if the world has structure to capture. The wanderer row is the small proof: remember the shepherd who never learned and still scored 88%.
• "MuZero learned the game." It learned what planning needed (policy, value, reward), which is less, and smarter.

Now teach it

Why it matters: five whys

The problem: an engineer wants to work on world models and does not know where to start, because most headline papers start at billions of parameters.

1. Why do the headline papers start big? The headline results are at scale: Genie's 11B parameters (Bruce et al. (2024), Genie, arXiv 2402.15391), and DreamerV3 in Nature (Hafner et al. (2025), Mastering diverse control tasks through world models, Nature 640).
2. Why does starting big block understanding? At scale you cannot see which part does the work: the model, the planner, or the data.
3. Why not? Because the parts are entangled in one network. In the CS229 loop they are separate lines you can switch off one at a time (Ma & Ng, CS229 Lecture Notes, eq. 19.5 (p. 233)).
4. Why does switching parts off matter? It is how you find what matters. The never-learning baseline and the wanderer row are exactly that kind of experiment, and Dyna names the honest caveat: learned models are "ofttimes incorrect" (Sutton (1991), Dyna, ACM SIGART Bulletin 2).
5. Root cause: world models are taught top-down. The fix is to build a small complete one, measure it against a baseline, break it on purpose, and then climb the lineage knowing what each step added.

Operate: run it this week

1. Read the teach-ahead line.
2. Sketch the shepherd's method on paper (Activity 1, 5 minutes).
3. Run the Shepherd's World Model on all four habits and fill in your own copy of the table (10 minutes).
4. Run the counting half in Python, finish it with the shepherd's actions, then break it on purpose (15 minutes).
5. Run the family drill with a child and post the outcome.

Feedback: how you know it worked

• You can state the CS229 section 19.3 loop from memory in four steps.
• You can explain why the wanderer row barely improves.
• You changed one assumption in the code and can say what broke and why.

Simulation-based learning has strong evidence in higher education: g = 0.85 across 145 studies (Chernikova et al. (2020), Simulation-based learning in higher education, Review of Educational Research 90). That is for students using simulations, not for children, and not for building one, so treat the family drill as a test, not a promise. Studies suggest young children build causal maps of their world: "2- to 4-year-old children construct new causal maps" (Gopnik et al. (2004), A theory of causal learning in children, Psychological Review 111). The drill asks whether that shows up when they predict a wolf.

Feedforward: tripwires, decided now

• If your learned model beats the baseline by less than the noise between seeds, you have not learned anything yet. Run more seeds before you claim a result, and always check whether the shepherd who never learns already wins.
• If you call something a world model, say which line it runs. Counting, planning, dreaming or generating video are different claims.
• If a famous result is quoted without its hedge, restore the hedge ("to our knowledge", "a few minutes").

The 5S self-score

Sources

1. Ma, T. & Ng, A. CS229 Lecture Notes, Stanford, Spring 2026 (section 19.3, pp. 232–234; 20.4; 11.5; 14; 8.2). cs229.stanford.edu
2. Craik, K. (1943). The Nature of Explanation, as quoted in Quanta Magazine (2025). quantamagazine.org
3. Kalman, R. E. (1960). A New Approach to Linear Filtering and Prediction Problems. J. Basic Eng. 82. Crossref
4. Sutton, R. S. (1991). Dyna. ACM SIGART Bulletin 2. OpenAlex
5. Ha, D. & Schmidhuber, J. (2018). World Models. arXiv:1803.10122
6. Schrittwieser, J. et al. (2020). MuZero. Nature 588. nature.com
7. Hafner, D. et al. (2025). Mastering diverse control tasks through world models. Nature 640. nature.com
8. Bruce, J. et al. (2024). Genie: Generative Interactive Environments. arXiv:2402.15391
9. Google DeepMind (2025). Genie 3. deepmind.google
10. NVIDIA (2025). Cosmos World Foundation Model Platform for Physical AI. arXiv:2501.03575
11. OpenAI (2024). Video generation models as world simulators (abstract via MIT CSAIL). csail.mit.edu
12. Kingma, D. P. & Welling, M. (2013). Auto-Encoding Variational Bayes. arXiv:1312.6114
13. Ho, J., Jain, A. & Abbeel, P. (2020). Denoising Diffusion Probabilistic Models. arXiv:2006.11239
14. Belkin, M., Hsu, D., Ma, S. & Mandal, S. (2019). Reconciling modern machine-learning practice and the classical bias–variance trade-off. arXiv:1812.11118
15. Nakkiran, P. et al. (2019). Deep Double Descent. arXiv:1912.02292
16. Chernikova, O. et al. (2020). Simulation-based learning in higher education. Review of Educational Research 90. doi:10.3102/0034654320933544
17. Gopnik, A. et al. (2004). A theory of causal learning in children. Psychological Review 111. PubMed
18. Kobayashi, K. (2024). Interactive learning effects of preparing to teach and teaching. Educational Psychology Review. ERIC EJ1414484
19. 1 Samuel 17:34–36, World English Bible. bible-api.com

Related