Model-Based Policy Optimization (MBPO)

Why Does This Exist?

Model free methods such as Soft Actor-Critic reach strong final performance but need millions of real environment steps. On a physical robot that cost is paid in hours, hardware wear, and broken parts. Model based methods answer by learning a simulator and training inside it, yet long imagined trajectories drift: each predicted step feeds its own small error into the next one, and after hundreds of steps the policy is optimizing hallucinated rewards. That failure is called model exploitation.

MBPO, introduced by Janner et al. in 2019, keeps the sample savings of a learned model while bounding that drift. The trick is to never roll out far. Every imagined trajectory is a short branch of length $k$, usually 1 to 5 steps, that starts from a real state already seen in the environment. Short branches keep the compounding error term linear in $k$ instead of quadratic in the horizon, and a mixed replay buffer keeps gradients grounded in reality. This page covers that mechanism and its worked numbers. The broader family is in model-based deep RL, and the imagination machinery it builds on is in Dreamer.

The one idea you need first

A dynamics model $\hat{p}(s_{t+1}, r_t \mid s_t, a_t)$ predicts the next state and reward from the current state $s_t$ and action $a_t$. MBPO learns several of them as an ensemble, and the disagreement between members measures how uncertain the model is. When members agree, the prediction is trustworthy. When they disagree, the branch stops.

Think of It Like This

How It Actually Works

Three parts work together: an ensemble of probabilistic dynamics models, branched rollout generation, and a mixed replay buffer feeding an off policy optimizer.

1. Ensemble of probabilistic dynamics models

MBPO trains $B$ neural networks, typically 7, each predicting a Gaussian over the next state change and reward:

Here $\mu$ is the predicted mean and $\Sigma$ is the predicted spread. The spread of one member captures the randomness of the environment itself. The disagreement across the $B$ members captures the model's own ignorance: large disagreement means this region is far from the training data.

2. Branched rollouts from real states

Each iteration samples a batch of real states $s_0$ from the environment buffer $\mathcal{D}_{\text{env}}$, then simulates $k$ steps with the current policy $\pi$. At every imagined step one ensemble member is picked at random to predict the outcome, and the transition is stored in a model buffer $\mathcal{D}_{\text{model}}$. If ensemble disagreement passes a threshold, the branch ends early.

3. Mixed batches for the optimizer

Policy and value updates draw about 95 percent of each mini batch from $\mathcal{D}_{\text{model}}$ and 5 percent from $\mathcal{D}_{\text{env}}$. The imagined majority supplies the sample efficiency. The real minority keeps the optimizer anchored to the physical world.

Why short branches bound the error

Let $\epsilon_{\text{model}}$ be the single step model error and $\gamma$ the discount factor. Under branched rollouts the compounding part of the return discrepancy bound grows as

which is linear in $k$. Long rollouts from a start state pay a factor near $1/(1 - \gamma)^2$ instead, which at $\gamma = 0.99$ multiplies errors by about 10,000.

Worked example

Seed a branch with $\epsilon_{\text{model}} = 0.05$, $\gamma = 0.99$, and horizon $k = 2$:

Seed from the real buffer: $s_0 = [1.00, 0.50]$. The policy outputs $a_0 = 0.20$. The ensemble mean predicts $\hat{s}_1 = [1.15, 0.45]$ with reward $\hat{r}_0 = 1.20$, and member disagreement sits far below the safety threshold, so the transition $(s_0, a_0, \hat{r}_0, \hat{s}_1)$ enters $\mathcal{D}_{\text{model}}$.

From $\hat{s}_1$ the policy outputs $a_1 = 0.40$. The ensemble predicts $\hat{s}_2 = [1.32, 0.38]$ with reward $\hat{r}_1 = 1.50$, again with low disagreement, and the second transition is stored. The branch ends because $k = 2$ is reached.

Now the bounds. Since $1 - \gamma = 0.01$:

Keeping $k$ at 2 holds the bound at 10.00 instead of 100.00, a 90 percent reduction, which is why the policy gradients stay trustworthy.

Code

This snippet computes the bound for several horizons and runs the same 2 step branch as the worked example, so you can check the arithmetic yourself.

1def branch_error(k: int, eps: float = 0.05, gamma: float = 0.99) -> float:2    return (k * eps) / (1.0 - gamma)3
4
5states = [[1.00, 0.50], [1.15, 0.45], [1.32, 0.38]]6rewards = [1.20, 1.50]7
8print(round(branch_error(1), 2))    # -> 5.09print(round(branch_error(2), 2))    # -> 10.010print(round(branch_error(20), 2))   # -> 100.011print(states[1])                    # -> [1.15, 0.45]12print(sum(rewards))                 # -> 2.7

The imagined return over the branch is $1.20 + 1.50 = 2.70$, and the bound at $k = 2$ reads 10.0 against 100.0 for a 20 step rollout from the same model.

Watch Out For

The Quick Version

• MBPO trains mostly on short imagined branches of 1 to 5 steps, each seeded from a real environment state.
• Short branches make compounding model error grow linearly in $k$, replacing the near quadratic blowup of long rollouts.
• An ensemble of probabilistic dynamics models supplies predictions plus a disagreement signal that ends shaky branches early.
• Mini batches mix about 95 percent imagined transitions with 5 percent real ones, pairing sample efficiency with grounding.
• At single step error 0.05 and discount 0.99, the bound reads 5.00 for $k = 1$, 10.00 for $k = 2$, and 100.00 for $k = 20$.

What to Read Next

• Model-Based Deep RL
• Dreamer: World Models for RL
• Soft Actor-Critic (SAC)