Model-Based Reinforcement Learning

Learning Before Acting

Model-free reinforcement learning learns values or policies directly from experience. Model-based reinforcement learning also learns, or is given, a model of what happens next:

\widehat p_\phi(s',r,d\mid s,a),

where $d$ indicates termination. The model can answer counterfactual questions without another real interaction:

\text{What would happen if we took }a\text{ in }s\text{?}

This separates expensive data collection from reusable computation. One real transition can train the model once, then support many imagined action sequences, policy updates, or value targets.

A model need not reproduce every detail of the world. For control, it only needs to preserve the consequences that affect rewards and decisions. It may predict observations,

p_\phi(o_{t+1}\mid o_t,a_t),

physical states,

p_\phi(s_{t+1}\mid s_t,a_t),

or compact latent states,

p_\phi(z_{t+1}\mid z_t,a_t).

A visually realistic simulator can therefore be less useful than a small, smooth model that retains task-relevant dynamics.

Let a behavior policy $\pi_\beta$ collect

\mathcal D =\{(s_i,a_i,r_i,d_i,s_i')\}_{i=1}^{N}.

Model-based learning alternates between

\boxed{ \text{collect real data} \longrightarrow \text{fit a model} \longrightarrow \text{improve behavior in the model} \longrightarrow \text{collect new data} }.

The model can improve decisions in several ways:

optimize an action sequence at decision time;
train a policy on imagined trajectories;
generate synthetic replay transitions;
extend a value target for a few steps;
learn a representation in which control is easier.

The tempting idea is to fit a simulator and then treat it as the real environment. The central difficulty is that the policy is optimized specifically to find trajectories the model predicts are good—including trajectories that look good only because the model is wrong.

Learning a One-Step Dynamics Model

For deterministic dynamics, predict the next state with

\hat s_{t+1}=f_\phi(s_t,a_t).

The usual regression objective is

\boxed{ \mathcal L_{\text{dyn}}(\phi) =\frac1N\sum_{i=1}^{N} \left\|s_i'-f_\phi(s_i,a_i)\right\|_2^2 }.

When adjacent states are similar, predicting the change can be easier:

\widehat{\Delta s}_t=g_\phi(s_t,a_t), \qquad \hat s_{t+1}=s_t+g_\phi(s_t,a_t).

The residual form builds persistence into the parameterization instead of asking the network to relearn it.

Squared error is a likelihood objective

Suppose the model has fixed isotropic Gaussian noise:

p_\phi(s'\mid s,a) =\mathcal N\!\left(s';\mu_\phi(s,a),\sigma^2I\right).

Then maximum likelihood produces exactly the squared-error estimator:

\boxed{ \arg\max_\phi\sum_i\log p_\phi(s_i'\mid s_i,a_i) = \arg\min_\phi\sum_i \left\|s_i'-\mu_\phi(s_i,a_i)\right\|_2^2 }.

See derivation

For one $d$ -dimensional transition,

-\log p_\phi(s'\mid s,a) =\frac d2\log(2\pi\sigma^2) +\frac{1}{2\sigma^2} \left\|s'-\mu_\phi(s,a)\right\|_2^2.

The first term and the positive scale

Summing over the dataset therefore leaves

\arg\min_\phi \sum_i\left\|s_i'-\mu_\phi(s_i,a_i)\right\|_2^2.

Stochastic dynamics require distributions

A deterministic model averages all possible futures. When the transition is multimodal, that average may be a state the environment can never produce. We instead learn

p_\phi(s'\mid s,a)

by minimizing negative log-likelihood:

\boxed{ \mathcal L_{\text{NLL}}(\phi) =-\sum_{i=1}^{N}\log p_\phi(s_i'\mid s_i,a_i) }.

This is empirical forward-KL minimization:

\min_\phi \mathbb E_{(s,a)\sim d^{\pi_\beta}} D_{\mathrm{KL}} \left( p(\cdot\mid s,a) \,\|\, p_\phi(\cdot\mid s,a) \right).

See derivation

Expanding the divergence gives

D_{\mathrm{KL}}(p\|p_\phi) =\mathbb E_{s'\sim p} \left[ \log p(s'\mid s,a) -\log p_\phi(s'\mid s,a) \right].

The true-dynamics term

Replacing the remaining expectation with dataset samples yields negative log-likelihood.

A diagonal heteroscedastic Gaussian predicts both mean and variance:

p_\phi(s'\mid s,a) =\mathcal N\!\left( s';\mu_\phi(s,a), \operatorname{diag}(\sigma_\phi^2(s,a)) \right).

Its per-transition loss is

\mathcal L_i =\frac12\sum_{j=1}^{d} \left[ \frac{(s'_{ij}-\mu_{ij})^2}{\sigma_{ij}^2} +\log\sigma_{ij}^2 +\log 2\pi \right].

The variance down-weights unpredictable residuals, while $\log\sigma^2$ prevents the model from explaining every error with infinite noise.

A complete one-step world model may factor as

p(s',r,d\mid s,a) =p_\phi(s'\mid s,a) p_\rho(r\mid s,a) p_\eta(d\mid s,a),

with joint loss

\mathcal L =\mathcal L_{\text{dyn}} +\lambda_r\mathcal L_{\text{reward}} +\lambda_d\mathcal L_{\text{termination}}.

Termination must be modeled explicitly: imagining rewards beyond a true terminal state can create arbitrarily optimistic plans.

One-Step Accuracy Is Not Rollout Accuracy

The model is trained on real states,

(s_t,a_t)\sim d^{\pi_\beta},

but an imagined rollout feeds predictions back into the model:

\hat s_{t+1}=f_\phi(\hat s_t,a_t).

An early error changes every later input. Even a small local error can therefore become a large trajectory error.

Assume deterministic true dynamics $f^\star$ , a learned model $f_\phi$ , and

e_t=\|\hat s_t-s_t\|.

If the learned model is $L_f$ -Lipschitz in state and has one-step error at most $\epsilon$ , then

\boxed{ e_{t+1}\le L_fe_t+\epsilon }.

Consequently,

\boxed{ e_t\le\epsilon\sum_{k=0}^{t-2}L_f^k }.

See derivation

Add and subtract the model prediction from the true state:

\begin{aligned} e_{t+1} &=\|f_\phi(\hat s_t,a_t)-f^\star(s_t,a_t)\|\\ &\le \|f_\phi(\hat s_t,a_t)-f_\phi(s_t,a_t)\| +\|f_\phi(s_t,a_t)-f^\star(s_t,a_t)\|\\ &\le L_fe_t+\epsilon. \end{aligned}

Unrolling from $e_1=0$ gives

e_t\le\epsilon\sum_{k=0}^{t-2}L_f^k = \begin{cases} \epsilon(t-1),&L_f=1,\\[4pt] \epsilon\dfrac{L_f^{t-1}-1}{L_f-1},&L_f\ne1. \end{cases}

Neutral dynamics produce linear error growth. Locally expansive dynamics, with $L_f>1$ , can amplify errors exponentially.

Return error can grow quadratically

Suppose every relevant transition satisfies

D_{\mathrm{TV}} \left( p(\cdot\mid s,a), p_\phi(\cdot\mid s,a) \right) \le\epsilon,

and $|r(s,a)|\le R_{\max}$ . For a fixed policy and horizon $H$ ,

\boxed{ |J_p(\pi)-J_\phi(\pi)| =O(R_{\max}\epsilon H^2) }.

See derivation

Let $d_t^p$ and $d_t^\phi$ be the real and model state distributions. Each transition adds at most $\epsilon$ total variation, while applying the same Markov kernel cannot increase it. Hence

D_{\mathrm{TV}}(d_t^p,d_t^\phi) \le(t-1)\epsilon.

The expected reward mismatch at time $t$ is bounded by

\left| \mathbb E_{d_t^p,\pi}[r] -\mathbb E_{d_t^\phi,\pi}[r] \right| \le2R_{\max}(t-1)\epsilon.

Summing over the horizon,

\begin{aligned} |J_p(\pi)-J_\phi(\pi)| &\le2R_{\max}\epsilon \sum_{t=1}^{H}(t-1)\\ &=R_{\max}\epsilon H(H-1). \end{aligned}

With discounting, replace the sum by

\sum_{t=1}^{H}\gamma^{t-1}(t-1),

which scales as $\gamma/(1-\gamma)^2$ in the infinite-horizon limit.

The same local model error affects the next reward and every later state distribution. Long imagined rollouts are therefore expensive in bias, even when they are cheap in computation.

Optimization finds model errors

The fit controls prediction error under the data occupancy:

\mathbb E_{(s,a)\sim d^{\pi_\beta}} [\ell_\phi(s,a)].

The optimized policy cares about another occupancy:

\mathbb E_{(s,a)\sim d^{\pi_\phi}} [\ell_\phi(s,a)].

In general,

d^{\pi_\phi}\ne d^{\pi_\beta}.

Low data density causes inaccurate extrapolation; the planner discovers an optimistic error; the new policy moves toward it; and the model is queried still farther outside its support. Unlike ordinary covariate shift, the optimizer actively searches for the mistake.

A conservative policy update can limit one source of shift:

\mathbb E_{s\sim d^{\pi_\beta}} D_{\mathrm{KL}} \left( \pi_\phi(\cdot\mid s) \,\|\, \pi_\beta(\cdot\mid s) \right) \le\delta.

This does not guarantee similar state occupancies—small action changes can have large effects—but it reduces how far one model-learning iteration tries to move.

Uncertainty Separates Noise from Ignorance

Aleatoric uncertainty is irreducible randomness in the transition:

s'\sim p(\cdot\mid s,a).

It remains even with infinite data. The covariance emitted by a probabilistic model primarily represents this conditional noise.

Epistemic uncertainty comes from insufficient knowledge of the dynamics. It should be large in poorly covered regions and shrink when relevant data is collected. A single neural network can extrapolate confidently, so its output variance is not a reliable measure of epistemic uncertainty.

If $\theta$ denotes uncertain model parameters, the Bayesian predictive distribution is

p(s'\mid s,a,\mathcal D) =\int p(s'\mid s,a,\theta) p(\theta\mid\mathcal D)d\theta.

Its predictive covariance decomposes as

\boxed{ \operatorname{Cov}(S') = \underbrace{\mathbb E_\theta[\Sigma_\theta]}_{\text{aleatoric}} + \underbrace{\operatorname{Cov}_\theta(\mu_\theta)}_{\text{epistemic}} }.

See derivation

For a scalar prediction, add and subtract the conditional mean $\mathbb E[S'\mid\theta]$ inside the squared deviation from the marginal mean. Expanding the square gives

\operatorname{Var}(S') =\mathbb E_\theta [\operatorname{Var}(S'\mid\theta)] +\operatorname{Var}_\theta (\mathbb E[S'\mid\theta]).

The vector form replaces variances with covariance matrices.

Bootstrap ensembles

A practical approximation trains $K$ models with different bootstrap samples, initializations, or minibatch orders:

\{p_{\phi_1},\ldots,p_{\phi_K}\}.

For deterministic predictions $\mu_k(s,a)$ ,

\bar\mu(s,a)=\frac1K\sum_{k=1}^{K}\mu_k(s,a),

\boxed{ \Sigma_{\text{epi}}(s,a) =\frac1K\sum_{k=1}^{K} (\mu_k-\bar\mu)(\mu_k-\bar\mu)^\top }.

For probabilistic members with covariance $\Sigma_k$ ,

\Sigma_{\text{total}} \approx \frac1K\sum_k\Sigma_k +\frac1K\sum_k (\mu_k-\bar\mu)(\mu_k-\bar\mu)^\top.

Ensembles are only a coarse posterior approximation, but disagreement often identifies unsupported state-action pairs well enough to improve control.

Uncertainty needs a decision rule

Estimating expected return is risk neutral:

J_{\text{mean}}(A)=\mathbb E[J(A)].

It is not automatically conservative. A planner must decide how uncertainty changes the objective:

\boxed{ J_{\text{LCB}}(A) =\mathbb E[J(A)] -\lambda\sqrt{\operatorname{Var}(J(A))} },

or, for directed exploration,

J_{\text{UCB}}(A) =\mathbb E[J(A)] +\lambda\sqrt{\operatorname{Var}(J(A))}.

A stepwise pessimistic reward can use

\tilde r(s,a)=r(s,a)-\lambda u(s,a),

where $u(s,a)$ is ensemble disagreement. The penalty should target epistemic uncertainty: penalizing irreducible noise indiscriminately makes the controller risk averse rather than merely cautious about model error.

Rolling out only the ensemble mean is also unsafe. When members predict distinct modes, their average may be physically impossible. Better choices sample complete model trajectories, propagate particles, evaluate lower return quantiles, or optimize a worst-case ensemble return.

Planning with a Learned Model

Given deterministic dynamics and initial state $s_1$ , open-loop planning solves

\boxed{ a_{1:H}^* =\arg\max_{a_{1:H}} \sum_{t=1}^{H}\gamma^{t-1}r(s_t,a_t) \quad \text{s.t.} \quad s_{t+1}=f_\phi(s_t,a_t) }.

For stochastic dynamics, the objective becomes

a_{1:H}^* =\arg\max_{a_{1:H}} \mathbb E_{p_\phi} \left[ \sum_{t=1}^{H}\gamma^{t-1}r(s_t,a_t) \right].

A fixed sequence cannot respond to the realized future state. Model predictive control turns it into feedback:

\boxed{ \text{observe} \longrightarrow \text{plan }H\text{ steps} \longrightarrow \text{execute the first action} \longrightarrow \text{replan} }.

Frequent replanning corrects stochastic outcomes, disturbances, and model mismatch before they compound across the entire episode.

Random shooting

The simplest black-box planner samples candidate sequences

A^{(i)}\sim q_0(A), \qquad i=1,\ldots,N,

rolls each one through the model, and selects

A^*=\arg\max_i\widehat J(A^{(i)}).

It requires no gradients and parallelizes easily. But an $H$ -step plan with $d_a$ -dimensional actions is an $Hd_a$ -dimensional search problem, so uniform sampling deteriorates quickly.

Cross-entropy method

The cross-entropy method, or CEM, adapts the proposal distribution toward high-return sequences. Start with

q_{\eta_0}(A)=\mathcal N(A;\mu_0,\Sigma_0).

At each iteration, sample candidates, score them in the model, retain an elite set $\mathcal E$ , and fit the next proposal:

\boxed{ \eta_{j+1} =\arg\max_\eta \sum_{A^{(i)}\in\mathcal E_j} \log q_\eta(A^{(i)}) }.

For a Gaussian proposal,

\mu_{j+1} =\frac1M\sum_{A^{(i)}\in\mathcal E_j}A^{(i)},

\Sigma_{j+1} =\frac1M\sum_{A^{(i)}\in\mathcal E_j} (A^{(i)}-\mu_{j+1})(A^{(i)}-\mu_{j+1})^\top.

See derivation

The CEM update is maximum likelihood on the elite samples. For a Gaussian, differentiating

\sum_{i\in\mathcal E} \log\mathcal N(A^{(i)};\mu,\Sigma)

with respect to $\mu$ and setting the result to zero gives the elite sample mean. Substituting that mean and optimizing $\Sigma$ gives the elite sample covariance.

\textbf{CEM Model Predictive Control}

\begin{aligned} &\text{Observe the current state }s_t.\\[0.25em] &\text{Initialize }q_\eta(a_{t:t+H-1}).\\[0.25em] &\text{Repeat for }J\text{ CEM iterations:}\\[0.25em] &\qquad \text{sample }N\text{ action sequences from }q_\eta,\\[0.25em] &\qquad \text{roll them out through one or more learned models},\\[0.25em] &\qquad \text{score expected or pessimistic return},\\[0.25em] &\qquad \text{refit }q_\eta\text{ to the top }M\text{ sequences}.\\[0.25em] &\text{Execute the first action of the best sequence.}\\[0.25em] &\text{Observe the next state and replan.} \end{aligned}

Sampling-based planning works with non-smooth rewards, discrete events, and black-box models. Its price is repeated model evaluation, poor scaling with long action sequences, and continued exposure to model bias.

Learning a Policy Through the Model

Planning spends computation at deployment. An alternative amortizes that computation into a policy

a_t=\pi_\theta(s_t).

With deterministic differentiable dynamics,

s_{t+1}=f_\phi(s_t,a_t),

the model return

J_\phi(\theta) =\sum_{t=1}^{H}\gamma^{t-1}r(s_t,a_t)

is a differentiable computation graph.

Let

u_t=\frac{ds_t}{d\theta}.

The state sensitivity follows

\boxed{ u_{t+1}=A_tu_t+B_t },

where

A_t =\frac{\partial f_\phi}{\partial s_t} +\frac{\partial f_\phi}{\partial a_t} \frac{\partial\pi_\theta}{\partial s_t},

B_t =\frac{\partial f_\phi}{\partial a_t} \frac{\partial\pi_\theta}{\partial\theta}.

See derivation

Because $a_t=\pi_\theta(s_t)$ ,

\frac{da_t}{d\theta} =\frac{\partial\pi_\theta}{\partial\theta} +\frac{\partial\pi_\theta}{\partial s_t}u_t.

Applying the chain rule to $s_{t+1}=f_\phi(s_t,a_t)$ gives

\begin{aligned} u_{t+1} &=\frac{\partial f_\phi}{\partial s_t}u_t +\frac{\partial f_\phi}{\partial a_t} \frac{da_t}{d\theta}\\ &=A_tu_t+B_t. \end{aligned}

Starting from $u_1=0$ ,

u_t =\sum_{k=1}^{t-1} \left( \prod_{j=k+1}^{t-1}A_j \right)B_k.

Therefore,

\frac{dJ_\phi}{d\theta} =\sum_{t=1}^{H}\gamma^{t-1} \left[ \frac{\partial r_t}{\partial s_t}u_t +\frac{\partial r_t}{\partial a_t} \left( \frac{\partial\pi_\theta}{\partial\theta} +\frac{\partial\pi_\theta}{\partial s_t}u_t \right) \right].

The products of closed-loop Jacobians explain the curse of chaos. Singular values below one make long-horizon gradients vanish; values above one make them explode. Contact events and chaotic dynamics can make tiny early action changes produce discontinuous later outcomes.

Stochastic dynamics can be reparameterized as

s_{t+1}=g_\phi(s_t,a_t,\xi_t), \qquad \xi_t\sim p(\xi),

and a stochastic policy as

a_t=h_\theta(s_t,\epsilon_t), \qquad \epsilon_t\sim p(\epsilon).

Backpropagation then treats sampled noise as an external input. This pathwise estimator can have low variance for short, smooth rollouts, but it retains the long-Jacobian problem and depends strongly on the accuracy of model derivatives.

Policy gradients need only model samples

We can instead sample imagined trajectories and apply the score-function estimator from policy gradients:

\boxed{ \nabla_\theta J_\phi(\theta) =\mathbb E_{\tau\sim\pi_\theta,p_\phi} \left[ \sum_{t=1}^{H} \nabla_\theta\log\pi_\theta(a_t\mid s_t) G_t \right] }.

With a critic,

\nabla_\theta J_\phi(\theta) =\mathbb E \left[ \sum_t \nabla_\theta\log\pi_\theta(a_t\mid s_t) A_\phi^\pi(s_t,a_t) \right].

This estimator does not differentiate through transition Jacobians and naturally supports discrete actions. It trades that stability for higher sampling variance. Both estimators remain biased when the learned model is wrong.

Short Rollouts Spend Less Model Accuracy

Full imagined episodes demand global model accuracy. Short rollouts initialized from real replay states demand only local accuracy.

If a $k$ -step rollout begins at $s_0\sim\mathcal B_{\text{real}}$ , the previous error argument suggests bias of order

O(\epsilon k^2)

rather than $O(\epsilon H^2)$ . Sampling many real start states also covers different points in the episode without predicting the path needed to reach each one.

Dyna

Dyna uses the model as an additional source of transitions. A real transition and a synthetic transition can drive the same Q-learning update:

\boxed{ Q(s,a)\leftarrow Q(s,a) +\alpha \left[ r+\gamma\max_{a'}Q(s',a')-Q(s,a) \right] }.

The distinction is only where $(r,s')$ came from.

\textbf{Deep Dyna with Short Model Rollouts}

\begin{aligned} &\text{Initialize model }p_\phi,\text{ off-policy learner, and real buffer }\mathcal B.\\[0.25em] &\text{Repeat:}\\[0.25em] &\qquad \text{collect real transitions and append them to }\mathcal B,\\[0.25em] &\qquad \text{fit }p_\phi\text{ on }\mathcal B,\\[0.25em] &\qquad \text{sample real start states }s_0\sim\mathcal B,\\[0.25em] &\qquad \text{roll the current policy through }p_\phi\text{ for }k\text{ steps},\\[0.25em] &\qquad \text{store synthetic transitions in }\mathcal B_{\text{model}},\\[0.25em] &\qquad \text{update the learner from }\mathcal B\cup\mathcal B_{\text{model}},\\[0.25em] &\qquad \text{refresh stale synthetic data after refitting the model}. \end{aligned}

Old imagined transitions were produced by an older model and policy. Keeping them indefinitely mixes incompatible synthetic environments, so model buffers are usually short lived.

Model-based value expansion

Instead of storing an entire synthetic rollout, use it to extend a target for $k$ steps:

\boxed{ \widehat V_k(s_t) =\sum_{j=0}^{k-1}\gamma^j\hat r_{t+j} +\gamma^kV_{\bar\omega}(\hat s_{t+k}) }.

For an action-value critic,

\widehat Q_k(s_t,a_t) =\sum_{j=0}^{k-1}\gamma^j\hat r_{t+j} +\gamma^k \mathbb E_{a\sim\pi} Q_{\bar\omega}(\hat s_{t+k},a).

The rollout length trades two approximations:

\boxed{ \text{short }k \Rightarrow \text{more bootstrap error}, \qquad \text{long }k \Rightarrow \text{more model error} }.

Model-Based Policy Optimization, or MBPO, applies the same principle to off-policy actor-critic learning. It trains on a mixture

\mathcal L_Q =\lambda_{\text{real}} \mathbb E_{\mathcal B_{\text{real}}}[\ell_Q] +\lambda_{\text{model}} \mathbb E_{\mathcal B_{\text{model}}}[\ell_Q],

with short synthetic rollouts rooted at real states. The model need not be globally correct; it only needs to improve local targets enough to justify its bias.

Latent World Models

A state-space model predicts directly in a compact physical state. When the agent receives images, audio, or partial observations, that state may be unavailable. A latent world model jointly learns an internal state and its dynamics:

z_t\sim q_\psi(z_t\mid o_{\le t},a_{<t}),

z_{t+1}\sim p_\phi(z_{t+1}\mid z_t,a_t).

The representation should discard unpredictable or irrelevant detail while retaining what is needed for rewards, termination, future prediction, values, and actions. The control problem is unchanged, but the state through which the agent predicts and plans must now be learned.

In addition to latent dynamics, a useful world model predicts task-relevant outcomes:

p_\phi(z_{t+1}\mid z_t,a_t), \qquad p_\phi(r_t\mid z_t,a_t).

Pixel-perfect reconstruction is neither necessary nor sufficient. The latent state must preserve the information needed to choose well.

Variational representation learning

For a single observation $x$ , introduce an encoder $q_\psi(z\mid x)$ for the latent model

p_\phi(x,z)=p(z)p_\phi(x\mid z).

The log likelihood has the lower bound

\boxed{ \log p_\phi(x) \ge \mathbb E_{q_\psi(z\mid x)} [\log p_\phi(x\mid z)] -D_{\mathrm{KL}} \left(q_\psi(z\mid x)\|p(z)\right) }.

See derivation

Start by inserting the encoder into the marginal likelihood:

\begin{aligned} \log p_\phi(x) &=\log\int q_\psi(z\mid x) \frac{p_\phi(x,z)}{q_\psi(z\mid x)}dz\\ &\ge \mathbb E_{q_\psi} \left[ \log p_\phi(x,z)-\log q_\psi(z\mid x) \right]. \end{aligned}

Using $p_\phi(x,z)=p(z)p_\phi(x\mid z)$ produces the reconstruction term minus the prior KL. The exact gap is

\log p_\phi(x)-\mathcal L =D_{\mathrm{KL}} \left(q_\psi(z\mid x)\|p_\phi(z\mid x)\right) \ge0.

For a Gaussian encoder, sampling is reparameterized as

\epsilon\sim\mathcal N(0,I), \qquad z=\mu_\psi(x)+\sigma_\psi(x)\odot\epsilon,

so gradients can pass through the sampled latent.

A sequential latent model

A useful controlled state-space model factorizes as

\begin{aligned} p_\phi(z_{1:H},o_{1:H},r_{1:H}\mid a_{1:H}) &=p(z_1) \prod_{t=1}^{H-1} p_\phi(z_{t+1}\mid z_t,a_t)\\ &\quad\times \prod_{t=1}^{H} p_\phi(o_t\mid z_t) p_\phi(r_t\mid z_t,a_t). \end{aligned}

Its sequential evidence lower bound is

\boxed{ \begin{aligned} \mathcal L_{\text{seq}} &=\mathbb E_q \left[ \sum_t\log p_\phi(o_t\mid z_t) +\sum_t\log p_\phi(r_t\mid z_t,a_t) \right]\\ &\quad- D_{\mathrm{KL}} \left( q_\psi(z_{1:H}\mid o,a) \,\|\, p_\phi(z_{1:H}\mid a) \right). \end{aligned} }.

If the posterior and prior both factor across transitions, the trajectory KL becomes

\begin{aligned} D_{\mathrm{KL}}(q\|p) &=D_{\mathrm{KL}} \left(q_\psi(z_1\mid o_1)\|p(z_1)\right)\\ &\quad+ \sum_{t=1}^{H-1} \mathbb E_{q(z_t)} D_{\mathrm{KL}} \left( q_\psi(z_{t+1}\mid z_t,o_{\le t+1},a_{\le t}) \,\|\, p_\phi(z_{t+1}\mid z_t,a_t) \right). \end{aligned}

Each transition term trains the predictive prior to match the posterior state inferred after seeing the next observation.

A filtering posterior,

q_\psi(z_t\mid o_{1:t},a_{1:t-1}),

uses all information available online. A single-frame encoder $q_\psi(z_t\mid o_t)$ is simpler but cannot resolve perceptual ambiguity with history. A deterministic encoder cannot represent belief uncertainty and may collapse unless reconstruction, reward, contrastive, or variance objectives preserve useful structure.

Control in imagination

Once trained, the agent can imagine entirely in latent space:

\tilde z_{t+1}\sim p_\phi(\cdot\mid\tilde z_t,\tilde a_t), \qquad \tilde r_t\sim p_\phi(\cdot\mid\tilde z_t,\tilde a_t).

Latent MPC optimizes short action sequences, executes one action, observes the real next frame, updates the posterior state, and replans. A latent actor-critic instead learns

\pi_\theta(a\mid z), \qquad Q_\omega(z,a),

from encoded real transitions and short imagined rollouts. Both avoid decoding pixels during every imagined step.

The failure mode is representation bias: a latent model can predict observations well while discarding a small feature crucial for future reward. Control-related losses must therefore shape the state alongside reconstruction and dynamics consistency.

Latent models replace high-dimensional observation rollouts with compact imagined trajectories:

\boxed{ o_{\le t} \xrightarrow{q_\psi} z_t \xrightarrow{p_\phi,\,\pi} z_{t+1:t+H} \xrightarrow{r_\phi,\,V} \text{decision} }.

Reconstruction preserves observable information, transition KL terms align posterior and predictive states, and reward or value losses emphasize task-relevant structure. None is sufficient alone.

A latent model can reconstruct images while discarding a small feature crucial for control, or predict rewards while learning brittle dynamics outside the policy distribution. Its quality must therefore be judged by multi-step prediction, uncertainty calibration, and the performance of planning or actor-critic learning in the latent space.

The best world model is not the one that renders the most realistic future. It is the smallest belief state that preserves the consequences needed to choose well.

Choosing How to Use the Model

Model-based methods exchange environment interactions for modeling and planning compute.

Use direct planning or MPC when the action dimension and horizon are moderate, the model is fast, and frequent replanning is possible. Use pathwise optimization when dynamics and rewards are smooth and differentiable. Use policy gradients on model samples when derivatives are unavailable. Use Dyna, value expansion, or MBPO when the model is locally useful but unreliable over full episodes.

Across these choices, the same safeguards recur:

evaluate multi-step rollout and control error, not only one-step validation loss;
distinguish environmental noise from lack of data;
limit policy movement into unsupported regions;
prefer short rollouts and frequent replanning as uncertainty grows;
refresh synthetic transitions when the model changes.

The central tradeoff is

\boxed{ \text{longer imagination} \Rightarrow \text{more synthetic signal} +\text{ more accumulated model bias} }.

Model-based RL is not the act of replacing the environment with a flawless simulator. It is the discipline of using an imperfect model only where its counterfactual predictions are worth more than their error.