Model-Based Reinforcement Learning: Learning and Planning with World Models

Model-free agents learn purely from trial and error and often need millions of interactions — acceptable in a fast simulator, unacceptable for a real robot or a system affecting people. Model-based RL learns a model of how the world works — how states change and what rewards follow — and uses it to plan or to generate imagined experience. Done well, it is dramatically more sample-efficient. Done badly, the agent exploits the model's mistakes.

What is a model?#

A learned model predicts dynamics and rewards:

It is trained by supervised learning on observed transitions — regression for deterministic dynamics, or probabilistic models (e.g. networks outputting Gaussian means and variances) to capture stochasticity.

Using the model#

1. Dyna: learning from imagined experience#

Sutton's Dyna architecture (1990) interleaves:

• real experience → update the value function (direct RL) and update the model;
• planning: sample simulated transitions from the model and apply the same Q-learning update.

Each real step can drive many imagined updates.

import random
import numpy as np

def dyna_q(env, episodes=50, n_planning=20, alpha=0.1, gamma=0.95, eps=0.1, seed=0):
    rng = random.Random(seed)
    Q = np.zeros((env.observation_space.n, env.action_space.n))
    model = {}                                            # (s, a) -> (r, s2, terminal)
    for ep in range(episodes):
        s, _ = env.reset(seed=seed + ep); done = False
        while not done:
            a = env.action_space.sample() if rng.random() < eps else int(np.argmax(Q[s]))
            s2, r, term, trunc, _ = env.step(a); done = term or trunc
            Q[s, a] += alpha * (r + gamma * (0 if term else Q[s2].max()) - Q[s, a])   # direct RL
            model[(s, a)] = (r, s2, term)                                          # model learning
            for _ in range(n_planning):                                            # planning
                (ps, pa), (pr, ps2, pterm) = rng.choice(list(model.items()))
                Q[ps, pa] += alpha * (pr + gamma * (0 if pterm else Q[ps2].max()) - Q[ps, pa])
            s = s2
    return Q
# e.g. dyna_q(gym.make("FrozenLake-v1", is_slippery=False))

With planning steps, Dyna typically reaches good policies in far fewer real episodes than plain Q-learning.

2. Planning at decision time: model-predictive control (MPC)#

At each step, use the model to simulate candidate action sequences over a short horizon, choose the best, execute only the first action, then replan. Sampling-based optimisers such as the cross-entropy method make this practical. MPC re-plans constantly, correcting for model errors — widely used in robotics and industrial control.

3. Tree search with a model#

Games provide perfect models (the rules); search methods such as MCTS use them (next lecture).

The central problem: model error#

Learned models are imperfect, especially far from the training data. A planner optimising against the model will exploit its errors, finding actions that look great in imagination but fail in reality. Errors also compound over long imagined rollouts. Remedies:

• Probabilistic ensembles: train several models; disagreement measures uncertainty (epistemic). PETS (Chua et al., 2018) combined probabilistic ensembles with MPC and matched model-free performance with far fewer samples.
• Short rollouts: MBPO (Janner et al., 2019) generates only short model rollouts branching from real states and trains SAC on the mix of real and imagined data — trusting the model only near data.
• Penalise uncertainty in planning, or return to real data frequently.

Latent world models#

For high-dimensional observations like images, predicting future pixels is hard and wasteful. World models learn a compact latent state and dynamics in that space:

• World Models (Ha & Schmidhuber, 2018): a VAE for frames, an RNN for latent dynamics, and a small controller trained "in the dream".
• Dreamer (Hafner et al., 2020–2023): learns a recurrent state-space model from pixels and trains an actor–critic entirely on imagined latent trajectories; DreamerV3 used a single set of hyperparameters across diverse domains, including collecting diamonds in Minecraft from scratch.
• MuZero (Schrittwieser et al., 2020): learns a model that predicts only what matters for planning — rewards, values and policies — not observations, and plans with MCTS in the learned latent space; it matched AlphaZero in Go, chess and shogi without being given the rules, and set strong results on Atari.

Model-based vs model-free#