Model-Based Reinforcement Learning

A Markov decision process is defined by states $s$, actions $a$, a transition function $P(s'|s,a)$, a reward function $R(s,a)$, and a discount $\gamma$. Model-free methods like Deep Q-Networks (DQN) never estimate $P$ or $R$ explicitly; they learn $Q(s,a)$ or a policy $\pi(a|s)$ directly from sampled transitions $(s,a,r,s')$. Model-based RL instead fits approximations $\hat{P}_\theta(s'|s,a)$ and $\hat{R}_\theta(s,a)$ — collectively "the model" — from observed transitions, and then uses this model computationally, without needing new real interactions, to improve the policy. There are two broad ways to exploit a learned model. The first is planning: at decision time (or in the background), the agent uses the model to simulate candidate action sequences and picks the one with highest predicted return, using techniques such as random shooting, the cross-entropy method, model-predictive control (MPC), or tree search. The second is background planning / Dyna-style learning: the model generates synthetic transitions that are fed into an ordinary model-free update (e.g., Q-learning or an actor-critic update), effectively multiplying the amount of training data available per real environment step. Many modern systems blend both: they plan short rollouts to pick actions, and also use the model to generate extra training targets. The central quantity governing whether any of this is worthwhile is model bias and how it compounds. If the one-step prediction error of the model is $\epsilon$ (in some suitable metric, e.g., total variation distance between $\hat P$ and $P$), then the error in a $k$-step rollout's state distribution grows, in the worst case, linearly to quadratically in $k$ for small $\epsilon$, and can be much worse if the model is used to roll out a changing policy trained partly on its own hallucinations — the so-called distribution shift problem, since the model is only accurate on states the real environment actually visits. Managing this compounding error — via short rollout horizons, uncertainty-aware ensembles, or model-predictive replanning — is the defining engineering challenge of MBRL, and separates naive approaches from state-of-the-art ones like MBPO and Dreamer.

Step 1 — Collect real data. The agent interacts with the true environment using some (possibly exploratory) policy, storing transitions $(s_t, a_t, r_t, s_{t+1})$ in a replay buffer $\mathcal{D}_{\text{real}}$.

Step 2 — Fit the dynamics model. A parametric model $\hat P_\theta(s'|s,a)$, typically a neural network (feedforward, recurrent, or a latent-variable "world model"), is trained by supervised learning to predict $s_{t+1}$ (or the delta $s_{t+1}-s_t$) and $r_t$ from $(s_t,a_t)$, minimizing a loss such as negative log-likelihood under a Gaussian output distribution:

Predicting a distribution (mean and variance) rather than a point estimate lets the agent quantify its own uncertainty, which is crucial for controlling compounding error. Many systems fit an ensemble of $B$ such models to capture epistemic uncertainty (disagreement between models) in addition to aleatoric uncertainty (the predicted variance of each model).

Step 3 — Use the model. Two families of usage:

• Planning / control: Sample or optimize action sequences $a_{t:t+H}$ by rolling them forward through $\hat P_\theta$ and scoring cumulative predicted reward; execute the first action; replan at the next step (Model Predictive Control, MPC). Trajectory optimizers such as the Cross-Entropy Method (CEM) or gradient-based shooting methods are typical optimizers.
• Background data generation (Dyna-style): From states in $\mathcal{D}_{\text{real}}$, sample short synthetic rollouts of length $k$ using $\hat P_\theta$ and the current policy, add these synthetic transitions to a separate buffer $\mathcal{D}_{\text{model}}$, and train the policy/value function with an off-the-shelf model-free algorithm using a mix of real and synthetic data.

Step 4 — Iterate. Because the policy changes, and because the model was only trained on old data, the model must be periodically re-fit on newly collected real transitions, and the loop repeats. This outer loop — collect real data, fit model, plan or generate synthetic data, improve policy, collect more real data — is the skeleton of essentially every MBRL algorithm. ```text Algorithm: Generic Model-Based RL (Dyna-style, e.g., MBPO) Input: real env E, policy π_φ, model ensemble {P̂_θ^1,...,P̂_θ^B}, rollout length k, real buffer D_real, model buffer D_model initialize π_φ, model ensemble θ, D_real ← ∅, D_model ← ∅ for iteration = 1, 2, ... do

# ---- (A) collect real experience ---- for t = 1..N_real do a ← π_φ(s) # possibly with exploration noise s', r ← step(E, s, a) D_real ← D_real ∪ {(s,a,r,s')} s ← s' end for

# ---- (B) fit dynamics model on REAL data only ---- for epoch = 1..E_model do sample minibatch B ⊂ D_real update θ by gradient descent on negative log-likelihood loss end for

# ---- (C) generate synthetic rollouts (background planning) ---- for m = 1..N_rollouts do s0 ← sample a real state from D_real # start on-distribution s ← s0 for step = 1..k do # SHORT horizon: caps error a ← π_φ(s) pick model j uniformly from ensemble # or aggregate ensemble s', r ← P̂_θ^j(s, a) # sample from predicted distribution D_model ← D_model ∪ {(s,a,r,s')} s ← s' end for end for

# ---- (D) improve policy with model-free updates on mixed data ---- for update = 1..N_policy_updates do sample minibatch from D_model (mostly) and D_real (some) update π_φ, value/Q functions via SAC / Q-learning / actor-critic step end for end for ``` For pure planning agents (no policy network at all, e.g., PETS), step (D) is replaced by direct action selection: ```text Algorithm: Model Predictive Control with a learned model (e.g., PETS) at each real time step t with current state s_t: initialize a distribution over action sequences a_{t:t+H} (e.g., mean/std for CEM) for CEM iteration = 1..M do sample P candidate action sequences from the distribution for each candidate sequence: roll it forward through the model ensemble from s_t (optionally average predictions over ensemble members, or sample a different member per particle, to hedge against model error) compute predicted cumulative reward refit the distribution to the top-E elite sequences by predicted return end for execute only a_t, the FIRST action of the best sequence observe real s_{t+1}; append (s_t,a_t,r_t,s_{t+1}) to D_real; replan next step ``` The "replan every step, execute only the first action" pattern is the essence of MPC: it bounds how far the agent trusts its own model before checking back against reality.

Time and space. Let $|S|$ or feature dimension $d_s$ denote state size, $d_a$ action size, $B$ the number of ensemble models, $H$ the planning horizon, $P$ the number of candidate action sequences per CEM iteration, and $M$ the number of CEM iterations.

• Model training: one gradient step over a minibatch costs $O(\text{batch size} \times (d_s+d_a) \times \text{width})$ for a feedforward network — essentially standard neural-network training cost, independent of environment complexity beyond $d_s,d_a$.
• Planning at each real step (MPC/CEM): each rollout of horizon $H$ requires $H$ forward passes through the model; there are $P$ candidates per CEM iteration and $M$ refinement iterations, so per real environment step planning costs $O(M \cdot P \cdot H \cdot B)$ model forward passes (the $B$ factor if evaluating the full ensemble per step, though many implementations sample one ensemble member per particle to keep cost $O(M\cdot P\cdot H)$). This is paid per real time step, so it trades wall-clock compute for real-world sample efficiency — the central bargain of MBRL.
• Dyna-style background rollouts: generating $N_{\text{rollouts}}$ synthetic trajectories of length $k$ costs $O(N_{\text{rollouts}} \cdot k)$ model evaluations, but this is far cheaper than $k$ real environment steps whenever real interaction (physical time, robot wear, simulator cost) dominates a single neural-network forward pass — the typical regime in robotics and industrial control.
• Space: an ensemble of $B$ networks costs $O(B \times \text{network size})$; the synthetic buffer $D_{\text{model}}$ can be kept small (short rollouts, frequently overwritten) since stale synthetic data becomes increasingly wrong as the policy changes, unlike real data which remains ground truth forever.

Why the loop is correct (in the limit). If $\hat P_\theta \to P$ and $\hat R_\theta \to R$ (the model converges to the true dynamics and reward, which happens under standard supervised-learning consistency conditions given enough diverse real data and sufficient model capacity), then: 1. Any planning procedure that computes $\arg\max$ expected return under $\hat P_\theta,\hat R_\theta$ over horizon $H$ converges to the optimal $H$-step lookahead policy under the true MDP, by continuity of the Bellman expectation operator in $P$ and $R$. 2. Any policy-improvement step performed on synthetic data drawn from $\hat P_\theta \approx P$ is, to first order, equivalent to performing that same step on real data, so Dyna-style updates converge to the same fixed point as pure model-free learning, but require far fewer real samples to reach a given model accuracy, since the same real transition informs both the model fit and any downstream planning.

The honest caveat — compounding model error. The convergence argument above hinges on $\hat P_\theta \approx P$ on the states the agent actually visits. But the model is trained on states visited by past policies, while planning/rollouts query it on states induced by the current, changing policy — a distribution-shift problem structurally identical to the one in imitation learning. A useful bound (informal): if the per-step total-variation model error is bounded by $\epsilon$, and the model is used to roll out $k$ steps under a policy whose value function has per-step reward bounded by $R_{\max}$, then the return-difference between planning in the model versus the real environment is bounded by

reflecting that model error accumulates quadratically with rollout length in the worst case (each of the $k$ steps can compound errors from all $k-1$ previous ones). This is the formal justification for two dominant design choices in modern MBRL: (a) keep synthetic rollouts short (MBPO typically uses $k=1$–$15$ steps, always starting from real states) so error stays small, and (b) replan frequently in MPC so the agent never commits to a long open-loop plan whose errors have had time to compound; it only trusts the model's very first predicted step before checking reality again.

Consider a 1-D "cart on a frictionless line" toy control task: state $s$ is position $x$ and velocity $v$; action $a\in[-1,1]$ is applied force; true dynamics are $v_{t+1}=v_t+0.1a_t$, $x_{t+1}=x_t+0.1v_t$; reward is $r_t=-x_t^2$ (stay near the origin). Suppose the learned model has a small systematic bias: it predicts $\hat v_{t+1}=v_t+0.1a_t\cdot(1.0)$ correctly for velocity, but slightly overestimates position sensitivity: $\hat x_{t+1}=x_t+0.11v_t$ (an 10% overestimate — a modest, realistic learned-model error). Start at $x_0=1.0, v_0=0.0$, and apply a constant corrective action $a_t=-0.5$ for 5 steps in both the true environment and the model, to see how far predictions diverge.

Even though the velocity prediction is exact at every step (the model has zero error there), the position error accumulates every step because each step's position update depends on the velocity accumulated from all previous (slightly mis-integrated) positions through the coupled dynamics — this is compounding error in miniature. Over only 5 steps the divergence is small (0.005), but the position error grows roughly linearly here because the bias is a constant multiplicative factor on velocity's contribution to position; in higher-dimensional, nonlinear, chaotic dynamics (e.g., contact-rich robotic manipulation), the same mechanism produces errors that grow much faster, often practically exponentially, over just tens of steps. This is precisely why an MBRL agent using this model for MPC with replanning every step would only ever act on the $t=0\to1$ prediction (error $0.000$) before re-measuring the real state, while an agent naively planning an open-loop 5-step trajectory in the model and executing it blindly would end up $0.005$ off in position — small here, but the qualitative lesson generalizes: short-horizon, frequently-corrected use of a model is dramatically more robust than long-horizon open-loop trust in it.

Model-based RL underlies systems where real-world trials are expensive or dangerous. In robotics, MPC with learned dynamics models lets manipulators and legged robots adapt to new objects or terrain within minutes rather than days of trial and error. In industrial process control (chemical plants, HVAC, data-center cooling), learned dynamics models allow safe, sample-efficient optimization without risking costly real-world failures during exploration. In autonomous driving simulation pipelines, learned world models are used to generate vast amounts of synthetic driving scenarios for training and stress-testing perception-to-control stacks. Game-playing agents such as MuZero extend the model-based idea by learning a model not of raw observations but of value-relevant abstract dynamics, which it then searches with a variant of Monte Carlo Tree Search — showing how a learned model, once available, naturally invites the kind of explicit lookahead planning that model-free agents cannot perform. The Dreamer family of agents demonstrates that agents can be trained almost entirely "in imagination," inside a learned latent world model, before ever touching the real environment again — the logical extreme of the sample-efficiency argument that motivates MBRL in the first place.

Ask the tutor: How does the MBPO algorithm formally bound the number of synthetic rollout steps needed to guarantee policy improvement over pure model-free learning?