The Complete Guide to Model-Based RL

What You'll Master

Why Model-Based?

You're training a robot arm to stack wooden blocks. Each attempt takes about 10 seconds: reach, grasp, lift, place, check, reset. Your model-free RL algorithm — say, SAC — needs around one million environment steps. That's 116 days of non-stop robot operation for a single stacking skill.

In simulation, a million steps is a coffee break. On a real robot, it's a death sentence for your project timeline, your motors, and your grad student's sanity.

What if the robot could practice in its imagination?

The Chess Grandmaster Analogy

Think about how a chess grandmaster plays. They don't physically move pieces to evaluate a position. They simulate moves mentally — "if I take the bishop, they'll push the pawn, I'll castle..." — and only execute moves they've already vetted in their head. The physical board sees maybe 40 moves. The mental board sees thousands.

That's the core idea. Instead of learning purely from real interactions (expensive), you learn a dynamics model — a learned simulator of how the world works — and then plan or practice inside that simulator.

Sample Efficiency: The Numbers

10-100x less data. That's not a minor improvement. It's the difference between "works in simulation only" and "works on a real robot in your lab."

The Dynamics Model

The core building block. We want to learn a function that predicts what happens next:

This is supervised learning. We collect transitions (s_t, a_t, s_t+1) from the real environment, then train a neural network to predict s_t+1 from (s_t, a_t).

Training Objective

Symbol by symbol:

• N — number of transitions in our dataset (replay buffer).

• fθ(si, ai) — model's predicted next state. This is a d-dimensional vector (same dimension as the state space).

• si' — the actual next state observed from the real environment.

• || · ||2 — squared L2 norm. Sum of squared differences across all state dimensions.

Architecture

A standard dynamics model is a feedforward neural network:

python
import torch.nn as nn

class DynamicsModel(nn.Module):
    def __init__(self, state_dim, action_dim, hidden=256):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim + action_dim, hidden),
            nn.ReLU(),
            nn.Linear(hidden, hidden),
            nn.ReLU(),
            nn.Linear(hidden, hidden),
            nn.ReLU(),
            nn.Linear(hidden, state_dim)  # predict next state
        )

    def forward(self, state, action):
        x = torch.cat([state, action], dim=-1)
        return self.net(x)  # predicted s_{t+1}

Delta vs. Absolute Prediction

In practice, predicting the change in state works better than predicting the absolute next state:

Data Collection

Where do we get (s, a, s') transitions? Three options, often used in combination:

• Random policy: Execute random actions, collect transitions. Cheap but only covers states near the initial distribution.

• Current policy: Roll out whatever policy we currently have. Covers relevant states but is expensive.

• Replay buffer: Store ALL transitions ever collected. Train model on full history. Gets richer over time.

Version 0.5: Random Shooting MPC

The simplest possible use of a dynamics model. No policy network, no gradients, no value function. Just brute-force search:

That's it. Sample random plans, simulate them in your head, pick the best first action. Re-plan every step.

Hand Calculation Example

Version 1.0: Dyna

Random shooting has no policy — it replans from scratch every step. What if we used the model to generate extra training data for a policy?

Dyna (Sutton, 1991) is the foundational idea: mix real transitions with imagined transitions, train a policy on both.

The key insight: every real transition is used twice. Once directly for Q-learning. Once to improve the model, which then generates K additional training transitions. With K=10, each real environment step effectively gives you 11 learning updates.

Data Amplification Factor

Model Error Propagation

Here's the fundamental problem with MBRL. Your model is never perfect. And when you roll out multiple steps, errors compound exponentially.

The Compounding Error Theorem

What the Numbers Mean

Version 1.5: Short Horizon + Replan

The error bound tells us: keep rollouts short. But short horizons mean myopic plans. The solution: Model Predictive Control (MPC) — plan short, replan often.

MPC limits error propagation to H steps. If H = 5, the worst-case error is bounded by our H=5 calculation (manageable), regardless of the total episode length.

Cross-Entropy Method (CEM)

Random shooting wastes samples. CEM iteratively narrows the search toward good action sequences:

Hand Calculation: CEM

The Full MPC Loop

Ensemble Models & Uncertainty

A single model gives you a point prediction. It can't tell you "I'm not sure about this region." Ensembles fix this by training multiple models and measuring their disagreement.

Symbol by symbol:

• K — number of models in ensemble (typically 5-7).

• fθk — the k-th model, trained independently.

• μ(s,a) — average prediction. Use this as the "best guess" next state.

• σ2(s,a) — variance across models. High = "we haven't seen this region" = don't trust the prediction.

PETS: Probabilistic Ensemble Trajectory Sampling

Pessimism Principle

When models disagree, the agent should be conservative. This emerges automatically from PETS (high-disagreement actions have high-variance returns, which CEM avoids), but can be made explicit:

Version 2.0: Model-Based Policy Optimization

MPC has a fatal limitation: no policy. It replans from scratch every step, which is computationally expensive at deployment. What if we used the model to train a reusable policy?

MBPO (Janner et al., 2019) is the state-of-the-art answer. Use short model rollouts to generate data, then train a model-free algorithm (SAC) on that data.

Why Short Rollouts (h = 1 to 5)?

The error propagation bound from Section 05 tells us: with L=1.5 and ε=0.01, a 5-step rollout has ~15% state error. That's borderline acceptable for training a policy. A 20-step rollout? 6000% error — pure noise. MBPO explicitly bounds rollout length.

Monotonic Improvement Guarantee

MBPO comes with a theoretical guarantee: under certain conditions, the policy improves monotonically (never gets worse). The key bound:

In English: the true policy performance is at least as good as the model-predicted performance minus a term proportional to the maximum model error. Keep model error small → model-predicted improvement transfers to reality.

Comparison: MPC vs. MBPO

When Models Fail

Three distinct failure modes plague MBRL. Understanding them is essential for debugging real systems.

Failure Mode 1: Objective Mismatch

Think of it this way: you train a weather model to predict temperature across the entire Earth equally well. But your user only cares about whether it'll rain in San Francisco tomorrow. Global MSE might be low, but SF predictions could be terrible.

Failure Mode 2: Model Exploitation

This is the same phenomenon as overfitting in supervised learning, but more dangerous: the policy actively searches for model weaknesses.

Failure Mode 3: Compounding Errors in Multi-Step

We derived this in Section 05. But here's the insidious part: it's not just that predictions get noisy. The errors are correlated. If the model consistently underpredicts friction, it'll predict the robot slides further than reality at every step. After 10 steps, the model thinks the robot is in a completely different place — systematically wrong, not randomly wrong.

Defenses

Case Study: Real Robot Learning

Let's put it all together. Real robot manipulation with MBRL: pushing objects to target locations using a 7-DOF arm.

The Setup

• State (17D): 7 joint positions + 7 joint velocities + 3D object position.

• Action (7D): Desired joint velocity for each motor.

• Reward: -||object_pos - target_pos|| (negative distance to target).

• Horizon: 100 steps per episode (5 seconds at 20Hz control).

The Recipe

Results Comparison

10x less data than SAC. 100x less than PPO. Slightly lower success rate, but achievable in a single lab session rather than days of continuous robot operation.

Summary & Cheat Sheet

Version Comparison

Decision Tree: When to Use Model-Based RL

Key Equations Reference

Connections & What's Next

Where MBRL Lives in the RL Landscape

The World Models Frontier

MBRL's natural evolution: instead of predicting raw state s_t+1, learn a latent dynamics model that predicts in a compressed representation. This handles high-dimensional observations (images) and enables longer-horizon planning by operating in a smoother space. Key systems: Dreamer (Hafner et al.), IRIS, and foundation world models.

Open Problems

• Contact-rich manipulation: Discontinuous dynamics break smooth model assumptions.

• Multi-task transfer: Can one world model serve many tasks?

• Online adaptation: How fast can models adapt when dynamics change (e.g., robot picks up heavy object)?

• Scaling to vision: Latent models work, but sample efficiency gains diminish at scale.

"The map is not the territory — but a good map saves you from walking off cliffs." — Adapted from Korzybski