Model-Based Reinforcement Learning

AlphaGo needed around 5 million self-play games to surpass human Go players. A well-tuned model-based agent tackling the same problem can learn a comparable policy with orders of magnitude fewer environment interactions, because it spends most of its compute on a learned simulator rather than the real world. That asymmetry is the central argument for model-based reinforcement learning (MBRL).

The fundamental split: model-free vs. model-based

In model-free RL, the agent learns a value function or policy entirely from samples of real experience. Each gradient step consumes data collected by actually executing actions. The environment is a black box: you query it, observe (s, a, r, s'), and move on. No structure is assumed.

Model-based RL adds one object: a dynamics model M that approximates the environment's transition and reward functions:

M : (s_t, a_t) --> (s_{t+1}, r_t)

Once trained, M can be queried millions of times per second without touching the real environment. This changes the data budget problem entirely. Real interactions are expensive (robot wear, simulation wall-time, API costs); synthetic rollouts through M are cheap.

The trade-off is model bias: errors in M compound over long rollouts, turning a small prediction mistake into a completely fictitious trajectory. Every design decision in MBRL is, in some form, an attempt to exploit the sample efficiency of the model while limiting the damage its errors cause.

How agents use a learned model

There are three broad strategies, and production systems often combine them.

1. Dyna-style background planning

Sutton's Dyna architecture (1991) remains the cleanest illustration. The agent collects real experience, updates M, then generates k synthetic transitions from M and uses those to update the policy or value function alongside real data. The ratio of synthetic to real updates is a hyper-parameter that controls the sample-efficiency / bias trade-off.

for each real step:
    observe (s, a, r, s')
    update M with (s, a, r, s')
    for k in range(K):
        s_sim  = sample from replay buffer
        a_sim  = policy(s_sim)
        r_sim, s'_sim = M(s_sim, a_sim)
        update policy/value with (s_sim, a_sim, r_sim, s'_sim)

MBPO (Janner et al., NeurIPS 2019) formalised this with a theoretical bound: branching short synthetic rollouts (length 1-5) from real data limits the compounding of model error while providing enough imagined experience to accelerate learning. On MuJoCo locomotion tasks, MBPO reaches model-free asymptotic performance using roughly 20-40x fewer environment samples.

2. Latent-space imagination

World Models (Ha and Schmidhuber, 2018) and the Dreamer family (Hafner et al., 2019) learn a compact latent representation of the state and train the dynamics model entirely in that compressed space. Instead of predicting pixel-level observations, the model predicts the next latent vector. Policy gradients can then be backpropagated through the differentiable imagined rollout:

loss = -sum_t gamma^t * r_theta(z_t, a_t)

where z_{t+1} = f_phi(z_t, a_t)   # latent transition model
      a_t     = pi_psi(z_t)        # policy in latent space

Because the imagined trajectories are differentiable, you get analytic gradients through the model rather than relying on high-variance Monte Carlo estimates. DreamerV2 reached human-level performance on 55 Atari games using latent imagination, without ever planning in pixel space.

3. Model predictive control (MPC) with a learned model

Here the model is used purely for planning, not policy training. At each real step, the agent optimises a sequence of actions over a short horizon H by repeatedly simulating candidate trajectories through M and selecting the sequence with the highest predicted return. No explicit policy is stored; the plan is computed online.

at each timestep t:
    for each candidate action sequence a_{t:t+H}:
        simulate trajectory: z_0 = encoder(s_t)
                             z_{i+1} = M(z_i, a_{t+i})
        compute predicted return
    execute a_t from the best sequence
    replan at t+1

Cross-Entropy Method (CEM) and random shooting are popular planning algorithms here. The advantage is zero policy-overfitting; the cost is that planning is compute-intensive at inference time.

What the dynamics model looks like

Modern MBRL uses neural network models, but the architecture choice matters.

Probabilistic ensembles (train N independently initialised networks, use disagreement as a proxy for epistemic uncertainty) are the current default for state-based continuous control. Disagreement between ensemble members is a reliable signal for when a proposed rollout has drifted outside the training distribution - a natural stop criterion for synthetic rollouts.

When it falls down

Compounding model error. Every dynamics prediction introduces noise. Over a T-step rollout, errors from step 1 persist and interact with errors at step 2, 3, ... T. For T = 20, a model with 2% per-step error can produce completely fictitious states. Short rollout limits (k = 1-5 in MBPO) mitigate this but also reduce the planning horizon.

Distribution shift at plan execution. The model is trained on data collected by a particular behavioural policy. When the learning policy improves and visits new states, the model is extrapolating rather than interpolating. This is called the Dyna trap: the policy learns to exploit model errors rather than to solve the true task, producing high imagined returns that do not transfer to reality.

Reward model vs. dynamics model. Often the reward function is itself learned alongside the dynamics. Misspecified reward models introduce a second source of hallucination. An agent can discover actions that yield high predicted reward with low predicted uncertainty - because the reward model is confidently wrong in a narrow region of state space.

Partial observability and aliasing. Standard MBRL assumes the observed state is Markov. When it is not (pixel observations, multi-agent settings, hidden variables), the deterministic model is misspecified from first principles. Recurrent latent models (RSSM) recover Markov structure in latent space, but the quality of the learned latent depends heavily on the encoder's inductive biases.

High-dimensional, discontinuous dynamics. Learned neural models smooth out sharp discontinuities. Contact dynamics in robotics, collision events, and abrupt mode switches (a ball bouncing, a latch snapping) are systematically underpredicted. This is one reason model-based robotics still lags behind model-free methods in contact-rich manipulation, despite the sample efficiency advantage.

Catastrophic forgetting of the model. As the policy explores new regions, the model needs to be updated. Naive gradient descent forgets old dynamics, destabilising the policy. Replay buffers help but add memory and require careful prioritisation.

Further reading

• Janner, Fu, Zhang, Levine. "When to Trust Your Model: Model-Based Policy Optimization." NeurIPS 2019. https://arxiv.org/abs/1906.08253
• Ha and Schmidhuber. "World Models." 2018. https://arxiv.org/abs/1803.10122
• Hafner, Lillicrap, Ba, Norouzi. "Dream to Control: Learning Behaviors by Latent Imagination." ICLR 2020. https://arxiv.org/abs/1912.01603
• Wang et al. "Benchmarking Model-Based Reinforcement Learning." 2019. https://arxiv.org/abs/1907.02057