Model-Based RL: Learning World Models

On 23 August 2023, Chandrayaan-3's lander Vikram had roughly fifteen minutes to go from an orbital velocity of about 1.68 km/s to a dead stop on the lunar surface, correcting its own trajectory the entire way, with no possibility of a human on Earth intervening — the round-trip signal delay to the Moon is over two seconds, too slow for real-time joystick control of a spacecraft decelerating that fast. A model-free reinforcement learning agent solves this class of problem by trial and error: try an action, see what happens, adjust. That approach is what learns to play Atari from scratch. It is not an option for a lunar landing, because "try an action and see what happens" means crashing a one-of-a-kind spacecraft that took years to build. ISRO's guidance engineers could not let the vehicle learn by doing on the Moon itself. What they could do — and did, exhaustively, after Chandrayaan-2's 2019 lander failed to survive its own final descent — was build a detailed model of how the vehicle responds to thruster commands under lunar gravity and atmosphere-free drag, and rehearse thousands of descent trajectories inside that simulated model before committing to the one real, irreversible attempt. That is the exact idea this chapter formalizes: instead of learning a policy purely from real interaction, an agent can learn a model of the environment's dynamics, then use that model to plan.

From model-free control to a learned world model

Recall the Markov Decision Process framework: an environment is described by states S, actions A, a transition function P(s'|s,a) giving the probability of landing in state s' after taking action a in state s, a reward function R(s,a), and a discount factor γ. Everything you have studied so far in this course under Q-learning and policy gradients is model-free: the agent never explicitly writes down P or R. It samples real transitions from the environment and updates a value function Q(s,a) or a policy π(s) directly, treating the environment as a black box it only ever queries by acting in it.

Model-based reinforcement learning adds one deliberate extra step. The agent uses its stream of real experience — tuples (s, a, r, s') — to fit an explicit approximation of the environment's dynamics: a learned transition model P̂(s'|s,a) and a learned reward model R̂(s,a), together called the world model. Once that model exists, the agent no longer needs to touch the real environment to reason about the consequences of an action — it can query the model instead. Querying a model is usually orders of magnitude cheaper than querying reality: a simulated thruster burn costs microseconds of compute; a real one costs fuel, spacecraft wear, and in the worst case the mission. This is the entire motivation for the extra machinery: sample efficiency. Every real interaction is expensive, dangerous, or irreversible in exactly the domains where model-based RL earns its keep — robotics, spacecraft guidance, industrial control, and, as the worked example below shows, any setting where you cannot afford to "just try it" many times before deciding.

What "using the model" means: planning

Having a model is only useful if the agent does something with it. The general term for that is planning: computing or improving a value function or policy using the model instead of, or in addition to, real experience. Three concrete ways this shows up:

Value iteration on the learned model. If P̂ and R̂ are known (even approximately), the agent can run the Bellman backup Q̂(s,a) = R̂(s,a) + γ Σ_s' P̂(s'|s,a) max_a' Q̂(s',a') directly, without generating a single new real sample — it is dynamic programming applied to an estimated MDP instead of the true one.

Model Predictive Control (MPC). Instead of solving the whole MDP, the agent uses the model to simulate a handful of action sequences a short number of steps into the future, picks whichever sequence looks best under the model, executes only the first action for real, observes the true outcome, and re-plans from the new state. This is standard in robotics and is close to what a descent-guidance system does at each control cycle.

Dyna-style planning (Sutton, 1990). This is the architecture the diagram below depicts, and it is the cleanest way to see model-based RL as an addition to, not a replacement for, what you already know. The agent keeps updating Q(s,a) from real transitions exactly as in ordinary Q-learning. Separately, every real transition it observes is also used to update the world model P̂, R̂. Then, between real steps, the agent draws additional simulated transitions from that model and runs the same Q-update rule on them. The real environment supplies ground truth; the model supplies volume.

Worked example: learning a model of death-over batting outcomes

Consider a T20 team's data-analytics unit trying to decide, at a fixed point in the death overs, whether to instruct the batter to play Aggressive or Defensive against a given bowler. They cannot run a controlled experiment against a real bowler in a real match hundreds of times — every real trial is one over of one real match. What they have are five recorded real overs where the batter played Aggressive against this bowler in this situation:

Playing Defensive against this bowler, from past data, is essentially deterministic: a steady single, reward +2, every time.

Step 1 — learn the model. From the five Aggressive trials: 2 out of 5 ended in a wicket, 3 out of 5 in a good start. The maximum-likelihood transition model is P̂(Wicket|Aggressive) = 2/5 = 0.4 and P̂(Good|Aggressive) = 3/5 = 0.6. The reward model is deterministic per outcome: R̂(Wicket) = −5, R̂(Good) = +8. This is exactly what "learning a world model" means in the tabular case — counting.

Step 2 — plan with the model. Both outcomes are terminal (the over-ending decision is a one-shot choice here), so the Bellman backup for Aggressive collapses to a plain expectation:

Q_model(Aggressive) = P̂(Wicket)·R̂(Wicket) + P̂(Good)·R̂(Good)
                     = 0.4·(−5) + 0.6·(8)
                     = −2.0 + 4.8
                     = 2.8

Aggressive (2.8) beats Defensive (2.0) under the learned model — the analytics unit can recommend Aggressive, and it has extracted that decision from a model built with only five real overs, with no further real trials needed.

Now compare this to what a purely model-free Q-learning agent would report from the same five samples, updating incrementally with the standard rule Q ← Q + α(r − Q) at a constant step size α = 0.5, starting from Q₀ = 0:

# Model learning: fit P_hat and R_hat by counting
rewards_seen = [-5, 8, 8, -5, 8]          # 5 real trials, in order
outcomes     = ["W", "G", "G", "W", "G"]

from collections import Counter
counts = Counter(outcomes)                       # Counter({'G': 3, 'W': 2})
n = len(outcomes)
P_hat = {s: c / n for s, c in counts.items()}    # {'W': 0.4, 'G': 0.6}
R_hat = {"W": -5, "G": 8}

# Planning: one exact Bellman backup using the learned model
Q_model = sum(P_hat[s] * R_hat[s] for s in P_hat)
print(round(Q_model, 4))          # 2.8

# Compare: model-free Q-learning, same 5 samples, constant step size
alpha, Q = 0.5, 0.0
for r in rewards_seen:
    Q += alpha * (r - Q)
print(round(Q, 5))                # 4.09375

Tracing the second loop by hand confirms the printed value: Q₁ = 0 + 0.5(−5−0) = −2.5; Q₂ = −2.5 + 0.5(8−(−2.5)) = 2.75; Q₃ = 2.75 + 0.5(8−2.75) = 5.375; Q₄ = 5.375 + 0.5(−5−5.375) = 0.1875; Q₅ = 0.1875 + 0.5(8−0.1875) = 4.09375. Notice that 2.8 is exactly the plain sample mean of the five rewards ((−5+8+8−5+8)/5 = 2.8), because with deterministic per-outcome rewards, the model-based expectation over counted probabilities is the sample mean. The model-free trace of 4.09375 is not the sample mean at all — constant-step-size Q-learning implicitly weights recent samples more heavily than old ones (a geometric decay of (1−α) per step back), so the two good starts at the end of the sequence dominate the estimate. Same five real overs, two different numbers, and the model-based one is the statistically efficient one. This is not a contrived quirk — it is exactly why Dyna-Q's model-learning step is valuable even before any simulated planning happens: fitting P̂, R̂ and backing up through them uses every data point equally, while naive incremental updates do not.

The second half of the Dyna-Q advantage is what the diagram's lower loop shows: once P̂, R̂ exist, the agent can draw as many additional simulated (s,a,r,s') samples as compute allows — sampling a Wicket 40% of the time and a Good start 60% of the time from the model — and run extra planning backups on them, refining estimates for this and neighbouring states without bowling a single additional real over. Real experience is what keeps the model honest; simulated experience is what makes the value estimates converge fast and cheaply once the model is decent.

The misconception: "model-based" does not mean "given a simulator"

The most common confusion at this point is thinking a chess or Go engine that searches a game tree using Monte Carlo Tree Search — like early AlphaGo — is doing model-based reinforcement learning because it "uses a model to plan ahead." It is not, in the technical sense this chapter defines. AlphaGo's tree search uses the actual, exactly known rules of Go as its model: given a board state and a move, the next state is not estimated, it is computed exactly from the rules. That is planning with a known model — closer to classical search than to model-based RL. Model-based RL specifically means the agent does not have access to the true P and R and must estimate P̂, R̂ from its own experience, the way the death-overs example above estimated dismissal probability from five real overs because nobody handed the analytics unit the bowler's true dismissal rate.

DeepMind's MuZero (2019) is the cleaner example of genuine model-based RL applied to games: it plans with MCTS exactly like AlphaZero, but it never sees the rules of chess, Go, or Atari — it learns an internal, abstract transition and reward model purely from experience and plans inside that learned representation. The correction matters because a learned model is never exact, and planning against an imperfect model can actively mislead the agent — a failure mode called model bias or model exploitation: the planner finds an action sequence that the model predicts is excellent precisely because the model has a blind spot there, and the real environment then punishes the agent for trusting it. This is a genuinely open problem in the field, not a footnote — it is the reason methods like PILCO (Deisenroth & Rasmussen, using Gaussian processes to keep an honest estimate of model uncertainty) and MBPO (Janner et al., 2019, which limits planning to short rollouts branching off real states rather than long imagined trajectories) exist: both are ways of not trusting the learned model further than the data supports it.

When the extra machinery pays off

Model-based RL is not strictly better than model-free RL — it trades real-world sample efficiency for extra computation and the risk of model bias. In a domain where real interaction is nearly free and fast, such as an Atari emulator running at thousands of frames per second on a GPU cluster, there is little to gain from learning a model of a simulator you can already query directly for free — this is why the landmark model-free results (DQN, PPO) target exactly those domains. Model-based methods earn their complexity where each real trial is slow, costly, or dangerous: PILCO learning to balance a physical cart-pole in under twenty real trials where model-free methods need thousands; Ha and Schmidhuber's 2018 "World Models" training a car-racing controller almost entirely inside a learned, compressed latent-space "dream" of the track before ever driving for real; and, at the scale this chapter opened with, any guidance system where the one real attempt has to work the first time.

Active recall

Attempt these before reading the answers.

1. In one sentence, what does a model-based RL agent compute that a model-free agent never computes?
2. A robot arm is tested on a new material and produces four real trials of a "grip and lift" action: outcomes are Success (reward +10), Success (+10), Slip (−3), Success (+10). Estimate P̂ and use it to compute the planned expected value of the action.
3. Using constant step size α = 0.5 starting at Q₀ = 0, trace model-free Q-learning on the same four rewards in the same order as question 2. Does it match your answer to question 2? Why or why not?
4. True or false, with justification: "AlphaZero, which plans using Monte Carlo Tree Search over the known rules of chess, is an example of model-based reinforcement learning."
5. Name one concrete real-world reason a robotics team would choose a model-based approach (like PILCO) over a model-free approach (like vanilla Q-learning), tying it to the sample-efficiency argument made in this chapter.

Answers

1. A model-free agent updates Q(s,a) or π(s) straight from real transitions and never represents P(s'|s,a) or R(s,a) explicitly. A model-based agent additionally fits estimates P̂(s'|s,a) and R̂(s,a) from experience and can use them to plan — compute value estimates or simulate extra transitions — without further real interaction.

2. Counts: 3 Success, 1 Slip, out of 4 trials, so P̂(Success) = 3/4 = 0.75, P̂(Slip) = 1/4 = 0.25. Planned value = 0.75·(10) + 0.25·(−3) = 7.5 − 0.75 = 6.75.

3. Q₁ = 0 + 0.5(10−0) = 5; Q₂ = 5 + 0.5(10−5) = 7.5; Q₃ = 7.5 + 0.5(−3−7.5) = 2.25; Q₄ = 2.25 + 0.5(10−2.25) = 6.125. This gives 6.125, close to but not equal to 6.75, because the single Slip trial happened third rather than last — its influence has partly decayed by the final update. It does not match exactly because constant-step-size Q-learning weights the most recent sample more than earlier ones, while the model-based expectation weights every observed trial equally regardless of order.

4. False. AlphaZero's tree search uses the exact, known rules of chess to generate the next board state — that is planning with a given model, not a learned one. Model-based RL specifically requires the agent to estimate P̂ and R̂ from its own experience because the true dynamics are unknown to it, as MuZero does by learning an internal model instead of using chess's known rules.

5. Real robot trials are slow (each one takes real seconds to minutes to execute and reset) and physically wear the hardware or risk damaging it, so a method that extracts a usable policy from a handful of real trials — as PILCO does by explicitly modeling and planning over the arm's learned dynamics — is far cheaper than a model-free method that needs thousands of real trials to converge.

Practice Exercises

Now it is time to practice! Complete these challenges to solidify your understanding:

• Exercise 1: Write a short program that demonstrates the core concept from this chapter. Test it with at least 3 different inputs.
• Exercise 2: Find a real-world example where model-based rl: learning world models is used in an Indian company (like TCS, Infosys, Flipkart, or ISRO). Write a paragraph explaining the connection.
• Exercise 3: Create a mind-map connecting model-based rl: learning world models to at least 3 other topics you have studied.

Key Takeaways — Summary and Recap

Let us recap what we covered: the core ideas behind model-based rl: learning world models, how they connect to real-world applications, and why they matter for your journey in computer science. Remember these key points as you move forward. For competitive exam preparation (CBSE, JEE, BITSAT), focus on understanding the WHY behind each concept, not just the WHAT.