Model-Based Reinforcement Learning

Background

Model-based RL predates model-free: dynamic programming (Bellman, 1957) requires a known model
AlphaGo and AlphaZero use model-based planning (MCTS) combined with learned value functions
Sample efficiency: model-based methods achieve similar performance with 10-100x fewer environment interactions

The Concept

Model-based RL learns a model of the environment (predicts next state and reward given current state and action), then uses it for planning. The agent can simulate many hypothetical trajectories in the model before committing to real actions, dramatically improving sample efficiency.

How it Links

World models are the deep learning version of model-based RL dynamics models
Dyna architecture (Sutton, 1991) combines model-based planning with model-free learning - both used simultaneously
MPC (Model Predictive Control) is the industrial control analog of model-based RL

Concrete Example

Chess: AlphaZero does not memorize openings or endgames. Instead, it uses a model (the rules of chess) to simulate thousands of move sequences via MCTS for each position, evaluating outcomes with a learned value network. It finds strong moves by planning ahead - not by pure pattern matching.

Prerequisites

Foundations of RL, Multilayer Perceptrons, The Math You Need

Learning outcomes

Explain the Dyna architecture and how it interleaves real experience with model-generated samples for improved sample efficiency.
Describe the World Models (Ha & Schmidhuber) framework: VAE vision, MDN-RNN memory, and a controller trained entirely in imagination.
Explain MBPO rollout length scheduling and the bias-variance tradeoff for model-generated data.
Derive how MuZero learns a latent dynamics model without explicit knowledge of rewards or transition rules, and contrast it with AlphaZero.
Compare model-free and model-based methods across sample efficiency, wall-clock time, and the compounding error problem.

TL;DR

Model-based RL learns an explicit model of the environment (transitions and rewards) and uses it to generate additional training data or to plan ahead, dramatically improving sample efficiency — at the cost of compounding errors when the model is imperfect. MuZero takes this further by learning a latent dynamics model that is sufficient for planning but never explicitly predicts raw observations.

Intuition

Model-free algorithms like DQN and PPO are powerful but data-hungry: every improvement to the policy requires real interactions with the environment. Model-based methods learn to predict “what would happen if I did X”, then use that world model either as a cheap simulator (Dyna), as a planning oracle (AlphaZero MCTS), or as a full training environment (World Models, Dreamer). The tradeoff is that learned models are imperfect, and errors compound over long rollouts.

Insight

The central insight of model-based RL: one real transition can teach the agent many imagined transitions. If the model is accurate, imagined rollouts are nearly as good as real ones and vastly cheaper. The challenge is knowing when to trust the model.

The Dyna Architecture

Sutton’s Dyna [1] (1990) is the simplest unifying framework: the agent alternates between (a) real experience, which updates both the value function and the model, and (b) simulated experience, where the model is queried for randomly chosen $(s, a)$ pairs and those hallucinated transitions are used for additional Q-learning updates.

repeat:
  1. Take action a in real env, observe (s, a, r, s')
  2. Update Q(s, a) with real (s, a, r, s')
  3. Update model M(s, a) -> (r, s')
  4. Repeat n times:
       sample random (s_sim, a_sim) seen before
       r_sim, s'_sim = M(s_sim, a_sim)
       update Q(s_sim, a_sim) with (r_sim, s'_sim)

The parameter $n$ controls the ratio of imagined to real updates. When the model is perfect, large $n$ accelerates learning enormously. When the model is wrong, spurious gradients hurt. Dyna is the ancestor of almost every modern model-based approach.

World Models (Ha & Schmidhuber 2018)

Ha and Schmidhuber’s World Models paper [2] decomposes the agent into three modules:

V (Vision): a Variational Autoencoder (VAE) that compresses each observation frame $o_t$ to a compact latent $z_t \in \mathbb{R}^{32}$ .
M (Memory): a Mixture Density Network RNN (MDN-RNN) that predicts $P(z_{t+1} \mid z_t, a_t, h_t)$ where $h_t$ is the RNN hidden state.
C (Controller): a small linear controller $a_t = W_c [z_t, h_t]$ with weights trained by CMA-ES in the imagination of the M model.

The key result: when the temperature of the MDN-RNN is turned up (making the dream harder and more random), a controller trained entirely in imagination can transfer to the real environment and achieve competitive scores on car racing and vizdoom tasks.

Warning

Hallucination in the dream: if the controller learns to exploit model inaccuracies (a “dream-hacking” strategy) it may fail in the real environment. Regularising the MDN-RNN temperature controls this.

MBPO: When to Trust Your Model

Janner et al. (2019) [3] formalise the bias-variance tradeoff for model-generated rollouts. With an ensemble of neural network dynamics models, MBPO generates short rollouts of length $k$ from states visited by the real policy, then adds these to the SAC replay buffer alongside real transitions.

MBPO rollout length and model error bound

Let $\epsilon_m$ be the expected one-step model error. After a $k$ -step rollout the compounding error grows as:

\mathbb{E}[\text{return error}] \leq \frac{2r_{\max} \gamma^k \epsilon_\pi}{(1-\gamma)^2} + \frac{2r_{\max} \gamma \epsilon_m k}{(1-\gamma)^2}

The first term is the policy approximation error carried forward; the second term is the compounding model error, which grows linearly in $k$ . Janner et al. show that there exists an optimal rollout length $k^*$ that minimises the total return error. Empirically, short rollouts ( $k = 1$ to $5$ ) work best early in training, with $k$ increasing as the model quality improves. MBPO achieves near-SAC sample efficiency on MuJoCo locomotion tasks with roughly 20-50x fewer real environment steps.

AlphaZero: MCTS with Learned Value and Policy

AlphaZero [5] combines:

A deep network $f_\theta(s) = (p, v)$ that outputs a policy prior $p$ over moves and a value $v \in [-1, 1]$ estimating the expected outcome from $s$ .
Monte Carlo Tree Search (MCTS): each node stores visit count $N(s,a)$ , value sum $W(s,a)$ , and action value $Q(s,a) = W(s,a)/N(s,a)$ . The PUCT selection rule balances exploration and exploitation:

a^* = \arg\max_a \left[ Q(s,a) + c_{\text{puct}} \cdot p(a \mid s) \cdot \frac{\sqrt{\sum_b N(s,b)}}{1 + N(s,a)} \right]

The network is trained entirely from self-play games: the MCTS visit distribution becomes the policy target, and the game outcome becomes the value target.

AlphaZero mastered chess, shogi, and Go from scratch — no human games, no domain-specific features. The learned value function replaces the handcrafted evaluation functions used in classical engines.

MuZero: Planning Without a Known Model

MuZero [4] removes the requirement that the environment model is known (which was the case for AlphaZero’s game tree). Instead, three networks are learned jointly:

Representation $h_\theta$ : maps the real observation $o_t$ to a hidden state $s^0$ .
Dynamics $g_\theta$ : predicts next hidden state and immediate reward given a hidden state and action: $(r^k, s^k) = g_\theta(s^{k-1}, a^k)$ .
Prediction $f_\theta$ : maps hidden state to policy prior $p^k$ and value $v^k$ , identical to AlphaZero’s head.

How MuZero learns a latent dynamics model without an explicit reward function

MuZero’s training objective is:

\mathcal{L}(\theta) = \sum_{k=0}^{K} \left[ \ell^r(r_t^k, z_{t+k}) + \ell^v(v_t^k, G_{t+k}^{\lambda}) + \ell^p(p_t^k, \pi_{t+k}) \right]

where $z_{t+k}$ is the target reward (from real play), $G_{t+k}^\lambda$ is the bootstrapped n-step return, and $\pi_{t+k}$ is the MCTS visit distribution.

Crucially:

There is no reconstruction loss on raw observations. The hidden states $s^k$ are never required to predict pixels — they only need to be sufficient for accurate reward prediction, value estimation, and policy improvement.
The dynamics network $g_\theta$ operates entirely in hidden space. Transitions are latent-to-latent, not pixel-to-pixel.
Reward is predicted by the dynamics network as a side output, so the agent does learn rewards — it just does not require them to be pre-specified or analytically available. The key difference from AlphaZero is that MuZero learns what the reward function is from experience.

This design means MuZero can be applied to visually rich environments (Atari) where specifying the dynamics model analytically is impossible, while retaining MCTS-quality planning. It achieved superhuman play on all 57 Atari games and matched AlphaZero on Go/Chess/Shogi.

AlphaZero vs MuZero: Key Differences

Property	AlphaZero	MuZero
Environment model	Known (game rules hard-coded)	Learned latent dynamics model
Reward function	Known (game outcome ±1)	Predicted by dynamics network from experience
Planning space	Real game states	Latent hidden states
Domains	Perfect-information board games	Board games + Atari + more
Observation	Board state	Raw pixels or feature vectors
Reconstruction loss	N/A	None — purely predictive hidden states

Dreamer: Latent Imagination for Continuous Control

Dreamer [6] extends the World Models idea to continuous-action tasks. A Recurrent State Space Model (RSSM) learns a compact world model; a policy and value function are then trained entirely in imagination by backpropagating through the differentiable world model. DreamerV2 and DreamerV3 further improved stability and generalisation, achieving strong results on Atari and DM Control.

Model-Free vs Model-Based: The Full Comparison

Dimension	Model-Free (DQN, PPO, SAC)	Model-Based (Dyna, MBPO, Dreamer, MuZero)
Sample efficiency	Low — needs many real interactions	High — leverages imagined rollouts
Wall-clock per update	Fast (no model inference needed)	Slower (model rollouts add overhead)
Compounding error	None	Grows with rollout length; mitigated by short rollouts or ensemble uncertainty
Planning horizon	Implicit in value function	Explicit MCTS or imagination rollouts
Handles partial observability	With memory (LSTM)	RSSM / belief-state models (Dreamer)
Works in stochastic envs	Yes	Yes, but model must capture stochasticity (MDN-RNN, RSSM)
Transfer / sim2real	Requires real env interactions	Real model can be queried without full env
Hallmark algorithms	DQN, PPO, SAC, TD3	Dyna, MBPO, AlphaZero, MuZero, DreamerV3

Real-World Applications

Robotics sim2real: MBRL trains policies in a fast physics simulator, then transfers to the real robot. The model absorbs the simulation gap and can be used to plan corrective actions online.

Game playing: AlphaZero demonstrated superhuman Chess/Shogi/Go; MuZero extended this to Atari without hard-coded rules. Both rely on MCTS planning over a learned (or known) model.

Autonomous driving planning: Planning layers in self-driving stacks (Waymo, Wayve) learn predictive models of other agents and use them to simulate possible futures before committing to a trajectory.

Drug discovery / scientific planning: MBRL frameworks are used to plan experiments in combinatorial spaces where real evaluations are expensive (wet lab synthesis, quantum chemistry).

Warning

Compounding error is the central challenge: a 1% per-step model error becomes 40% error over a 50-step horizon ( $0.99^{50} \approx 0.60$ ). Short rollouts (MBPO), ensemble uncertainty (PETS), or latent-only models that skip pixel reconstruction (MuZero) are the main mitigations.

Pitfalls

Warning

Model exploitation: the policy can learn to exploit model inaccuracies rather than learn to perform well in the real environment. Ensemble disagreement penalties and pessimistic value estimates (MOPO) partially address this.

Warning

Distributional shift in the model: the dynamics model is trained on data collected by older policies; it may be inaccurate for states visited by the current, more capable policy. Retraining the model frequently (MBPO) helps but adds computational cost.

Tip

When to use model-based RL: (1) real-world interactions are expensive (robotics, drug discovery, data-center cooling), (2) you have access to a differentiable simulator (physics engines), or (3) you need planning over longer horizons than a value function can implicitly capture.

Pop Quiz

Check yourself

Score: 0 / 5(0 answered)

1. In the Dyna architecture, what is the purpose of the simulated experience (model queries) step?
2. Why does model error compound over a k-step rollout?
3. How does MuZero differ from AlphaZero with respect to the environment model?
4. In AlphaZero, what serves as the policy target during network training?
5. MBPO achieves better sample efficiency than pure model-free SAC primarily because:

Interview Q&A

Why does model error compound, and how do modern model-based algorithms mitigate this?onsite

Each step of a model-based rollout applies the learned dynamics $\hat{f}$ to the current (possibly already incorrect) state. If the per-step error in predicted reward or next state is $\epsilon$ , after $k$ steps the error in the cumulative return is $O(k \cdot \epsilon)$ at minimum, and in practice the errors feed back into subsequent states so the divergence grows faster. The main mitigations are:

Short rollouts (MBPO $k = 1$ – $5$ ): limits how much error can accumulate.
Ensemble disagreement: train an ensemble of $N$ dynamics models; high variance across ensemble members signals unreliable regions and can be used as a pessimistic penalty (MOPO).
Latent-only models (MuZero, Dreamer): by never predicting raw observations, the model is not penalised for pixel-level inaccuracies that do not affect planning quality.
Dyna-style interleaving: keep real data dominant in the replay buffer; use model data only for supplementary updates.

How does MuZero differ from AlphaZero, and what does MuZero's dynamics network actually learn?staff

AlphaZero requires the transition function (game rules) and reward function (game outcome) to be provided explicitly. MCTS is run using these known rules; the network learns only the value and policy heads. MuZero removes this assumption: it learns a representation network $h_\theta(o_t) \to s^0$ , a dynamics network $g_\theta(s^{k-1}, a^k) \to (r^k, s^k)$ , and a prediction network $f_\theta(s^k) \to (p^k, v^k)$ . Planning proceeds by running MCTS entirely in the latent space using $g_\theta$ .

The dynamics network learns whatever hidden-state encoding makes reward prediction, value estimation, and policy prediction accurate. Crucially, there is no reconstruction loss — the hidden states need not decode back to pixels. The model is trained end-to-end by unrolling $K$ imagined steps and minimising the sum of reward prediction loss, value loss, and policy loss against targets from real play. Because reward prediction is part of the training objective, MuZero does learn the reward structure of the environment from experience — it just does not require it to be analytically specified beforehand.

Describe the Dyna framework and explain why it can be seen as a unifying framework for model-free and model-based RL.phone

Dyna (Sutton 1990) maintains both a Q-value function (like a model-free agent) and a learned model of the environment. After each real transition $(s, a, r, s')$ , it: (1) updates $Q(s,a)$ with the real TD target, (2) updates the model with the real transition, and (3) performs $n$ synthetic Q-learning updates by sampling stored $(s, a)$ pairs and querying the model for imagined $(r, s')$ . When $n = 0$ , Dyna reduces to pure model-free Q-learning. When the model is perfect and $n \to \infty$ , it approaches pure planning. Intermediate $n$ blends both. This makes Dyna a spectrum rather than a single algorithm, which is why it is seen as a unifying framework.

References

Sutton, R. S., and Barto, A. G. (2018). Reinforcement Learning: An Introduction (2nd ed.). MIT Press. ISBN 978-0-262-03924-6
Ha, D., and Schmidhuber, J. (2018). World Models. arXiv preprint. arXiv:1803.10122
Janner, M., Fu, J., Zhang, M., and Levine, S. (2019). When to Trust Your Model: Model-Based Policy Optimization. NeurIPS 2019 (arXiv preprint). arXiv:1906.08253
Schrittwieser, J., Antonoglou, I., Hubert, T., Simonyan, K., Sifre, L., Schmitt, S., Guez, A., Lockhart, E., Hassabis, D., Graepel, T., et al. (2020). Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model. Nature 588, 604-609. doi:10.1038/s41586-020-03051-4
Silver, D., Hubert, T., Schrittwieser, J., Antonoglou, I., Lai, M., Guez, A., Lanctot, M., Sifre, L., Kumaran, D., Graepel, T., et al. (2018). A general reinforcement learning algorithm that masters chess, shogi, and Go through self-play. Science 362(6419), 1140-1144. doi:10.1126/science.aar6404
Hafner, D., Lillicrap, T., Ba, J., and Norouzi, M. (2019). Dream to Control: Learning Behaviors by Latent Imagination. ICLR 2020 (arXiv preprint). arXiv:1912.01603
Mnih, V., Kavukcuoglu, K., Silver, D., Rusu, A. A., Veness, J., Bellemare, M. G., Graves, A., Riedmiller, M., Fidjeland, A. K., Ostrovski, G., et al. (2015). Human-level control through deep reinforcement learning. Nature 518, 529-533. doi:10.1038/nature14236
Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O. (2017). Proximal Policy Optimization Algorithms. arXiv preprint. arXiv:1707.06347

Foundations of RL Offline RL & RLHF