Model-Based Deep RL

The agent learns a small simulator of the world from real tries, then practices thousands of imagined moves inside it before each real move.

Ensemble mean 1.50 with epistemic variance 0.000267 below gate 0.05 is accepted into imagined training

Why Does This Exist?

Model free control treats the world as a black box and needs millions of real tries. On hardware that means wear, time, and crashes. Each transition updates one value or gradient and throws away what it reveals about how the world moves.

Model based deep RL keeps that structure. It fits neural nets to transition and reward from the same real logs, then rolls out imagined trajectories on GPU. Supervised fitting is far cheaper than trial and error, so real samples drop by an order or more. This page shows the ensemble guard and one accepted imagined step. For the plain idea see model based RL, for the seeding trick see Dyna, and for pixel worlds see learning latent dynamics.

You need one idea first. Aleatoric noise is randomness in the world itself. Epistemic doubt is the model not knowing some region yet. Only the second one tells you to stop imagining.

Think of It Like This

A chess player thinking before touching a piece

A novice learns only from tournament games where every blunder costs a match. A stronger player pictures five replies ahead, drops bad lines in the head, and plays one tested move on the board.

When the pictured position feels unfamiliar, the strong player picks a safe solid line rather than a sharp trap built on a guess.

Where the analogy stops: chess rules are fixed and known. A world model starts with no physics and learns from noisy logs, so its imagined lines can invent moves that do not exist.

How It Actually Works

Let $s$ be state, $a$ action, and $\mathcal{D}_{\text{env}}$ real transitions $(s, a, r, s')$ . Let $\hat{p}_\theta(s' \mid s, a)$ be the learned transition and $\hat{r}_\psi(s, a)$ the learned reward. Many nets predict a Gaussian over the change:

$\hat{p}_\theta(s' \mid s, a) = \mathcal{N}(s + \mu_\theta(s, a), \Sigma_\theta(s, a))$

Here $\mu_\theta$ is predicted mean change and $\Sigma_\theta$ is predicted noise. Training minimizes Gaussian log loss on $\mathcal{D}_{\text{env}}$ .

One net hallucinates outside its data, and the policy chases the hallucination. The fix is a bootstrap ensemble of $B$ nets with different starts and data bags. Let $\mu_i$ and $\sigma_i^2$ be member $i$ mean and noise, with average mean $\bar{\mu}$ :

Aleatoric noise: $\bar{\sigma}^2 = \frac{1}{B} \sum_i \sigma_i^2$ , the world randomness.
Epistemic doubt: $\frac{1}{B} \sum_i (\mu_i - \bar{\mu})^2$ , the member disagreement.

High disagreement means unknown ground. Short branching rollouts keep errors small: seed a real state from $\mathcal{D}_{\text{env}}$ , imagine only $k$ steps with $k$ from 1 to 15, store in $\mathcal{D}_{\text{model}}$ , and train the actor critic on a mix with most weight on imagined data plus some real anchor.

Worked example

Three nets see $s = 1.0$ and $a = 0.5$ and report means $1.48$ , $1.52$ , $1.50$ with noise $0.010$ , $0.020$ , $0.015$ :

Mean $\bar{\mu} = (1.48 + 1.52 + 1.50) / 3 = 1.50$ .
Aleatoric $\bar{\sigma}^2 = (0.010 + 0.020 + 0.015) / 3 = 0.015$ .
Epistemic $[(1.48 - 1.50)^2 + (1.52 - 1.50)^2 + 0] / 3 = 0.0008 / 3 \approx 0.000267$ .
Total $0.015 + 0.000267 = 0.015267$ .

Gate $\tau = 0.05$ . Since $0.000267$ sits far below it, the imagined step $(1.0, 0.5, 1.50)$ joins $\mathcal{D}_{\text{model}}$ . The diagram shows the same three dots, the 1.50 mean, and the open gate.

Code

The snippet splits the same two uncertainties and applies the gate.

1def ensemble_stats(means: list[float], noises: list[float]) -> tuple[float, float, float]:2    """Mean, aleatoric noise, and epistemic disagreement."""3    b = len(means)4    mean = sum(means) / b5    aleatoric = sum(noises) / b6    epistemic = sum((m - mean) ** 2 for m in means) / b7    return (round(mean, 4), round(aleatoric, 6), round(epistemic, 6))8
9
10if __name__ == "__main__":11    print(ensemble_stats([1.48, 1.52, 1.5], [0.01, 0.02, 0.015]))12    print(0.000267 < 0.05)

1# -> expected output:2(1.5, 0.015, 0.000267)3True

Watch Out For

Rolling a thousand imagined steps and believing the fantasy

The symptom: imagined returns climb while the real robot falls on step one. Feeding each prediction back in compounds error about as $k^2$ , so after 15 to 20 steps the states are unphysical yet the critic trusts them.

The fix: branch short runs from real states, schedule $k$ from 1 upward only as validation steadies, and cut any branch the moment epistemic disagreement passes your gate.

The Quick Version

Model based deep RL fits transition and reward nets with supervised loss, then trains the policy mostly on imagined rollouts.
A bootstrap ensemble splits world noise $\bar{\sigma}^2$ from model doubt measured by member disagreement.
In the example the mean is $1.50$ , aleatoric $0.015$ , epistemic $0.000267$ , total $0.015267$ .
Doubt below gate $0.05$ accepts the step, while far regions spike doubt and halt the branch.
Short branches from real states bound drift, which is where the sample savings come from.