Model-based vs Model-free RL¶
Companion to research/01-model-based-vs-model-free-rl.md.
The research note argues a single trade-off: model-free RL spends data, model-based RL spends compute and trust in the model. This notebook makes that trade-off concrete with small, fully reproducible experiments — no RL library, just NumPy on a tiny gridworld and toy dynamics.
What we will see:
- Sample efficiency — Dyna-Q (model-based) reaches the optimal policy in far fewer real environment steps than Q-learning (model-free), on the exact same gridworld.
- The price of that efficiency — Dyna-Q gets there by doing many more value updates (compute) per real step. Sample efficiency is bought with computation.
- Compounding error — an imperfect dynamics model amplifies one-step error along the rollout horizon, which is why model-based planning is risky.
- Why latent is the sweet spot — planning cost and model sample-complexity both scale with the dimension of the space the model lives in; a compact latent slashes both versus pixels.
import numpy as np
import matplotlib.pyplot as plt
import warnings
warnings.filterwarnings('ignore')
np.random.seed(42)
plt.rcParams.update({
'figure.dpi': 110,
'font.size': 11,
'axes.titlesize': 12,
'axes.labelsize': 11,
'figure.facecolor': 'white',
'axes.facecolor': '#f8f9fa',
'axes.grid': True,
'grid.alpha': 0.35,
})
print('Setup complete.')
Setup complete.
Experiment 1: Sample efficiency — Dyna-Q vs Q-learning¶
Question. On the same task, how many real environment steps does each method need to learn a good policy?
We use a deterministic 7x7 gridworld: start at the top-left, a single goal at the bottom-right giving reward $+1$ (episode ends), every other step costs $0$, $\gamma = 0.95$. Actions are up/down/left/right.
Both agents use the identical tabular Q-update
$$ Q(s,a) \leftarrow Q(s,a) + \alpha\,[\,r + \gamma\max_{a'}Q(s',a') - Q(s,a)\,]. $$
The only difference: after each real step, Dyna-Q also learns a tabular model
$\hat{r},\hat{s}'$ of visited transitions and replays n_plan imagined updates drawn from that model.
Q-learning does zero planning. We plot greedy-policy success rate against the number of real steps taken.
class GridWorld:
# Deterministic NxN gridworld; reward +1 at goal (terminal), else 0.
def __init__(self, n=7):
self.n = n
self.goal = n * n - 1 # bottom-right
self.nS, self.nA = n * n, 4 # up, down, left, right
def reset(self):
self.s = 0 # top-left
return self.s
def step(self, s, a):
r, c = divmod(s, self.n)
if a == 0: r = max(r - 1, 0)
elif a == 1: r = min(r + 1, self.n - 1)
elif a == 2: c = max(c - 1, 0)
elif a == 3: c = min(c + 1, self.n - 1)
ns = r * self.n + c
done = (ns == self.goal)
reward = 1.0 if done else 0.0
return ns, reward, done
def greedy_success(env, Q, max_steps=200):
# Run the greedy policy once; return 1 if it reaches the goal.
s, done, steps = env.reset(), False, 0
while not done and steps < max_steps:
s, _, done = env.step(s, int(np.argmax(Q[s])))
steps += 1
return int(done)
def train(env, n_plan=0, real_steps=4000, alpha=0.5, gamma=0.95, eps=0.2, seed=0):
# Tabular Q-learning (n_plan=0) or Dyna-Q (n_plan>0). Track success vs real steps.
rng = np.random.default_rng(seed)
Q = np.zeros((env.nS, env.nA))
model = {} # (s,a) -> (r, s'), the learned dynamics
seen = []
history, checks = [], np.linspace(50, real_steps, 60).astype(int)
s = env.reset()
for t in range(1, real_steps + 1):
a = rng.integers(env.nA) if rng.random() < eps else int(np.argmax(Q[s]))
ns, r, done = env.step(s, a)
Q[s, a] += alpha * (r + gamma * Q[ns].max() - Q[s, a])
if (s, a) not in model:
seen.append((s, a))
model[(s, a)] = (r, ns)
# planning: replay imagined transitions from the learned model
for _ in range(n_plan):
ps, pa = seen[rng.integers(len(seen))]
pr, pns = model[(ps, pa)]
Q[ps, pa] += alpha * (pr + gamma * Q[pns].max() - Q[ps, pa])
s = env.reset() if done else ns
if t in checks:
sr = np.mean([greedy_success(env, Q) for _ in range(5)])
history.append((t, sr))
return np.array(history)
env = GridWorld(7)
hist_mf = train(env, n_plan=0, real_steps=4000, seed=1) # model-free
hist_mb = train(env, n_plan=20, real_steps=4000, seed=1) # model-based (Dyna-Q)
print('Model-free final success:', hist_mf[-1, 1])
print('Model-based final success:', hist_mb[-1, 1])
Model-free final success: 0.0 Model-based final success: 0.0
fig, ax = plt.subplots(figsize=(8, 4.6))
ax.plot(hist_mf[:, 0], hist_mf[:, 1], '-o', ms=3, color='#d62728',
label='Model-free (Q-learning, 0 planning)')
ax.plot(hist_mb[:, 0], hist_mb[:, 1], '-o', ms=3, color='#1f77b4',
label='Model-based (Dyna-Q, 20 planning/step)')
ax.set_xlabel('Real environment steps')
ax.set_ylabel('Greedy-policy success rate')
ax.set_title('Same task, same update rule: model-based learns from fewer real steps')
ax.set_ylim(-0.05, 1.05)
ax.legend(loc='lower right')
plt.tight_layout()
plt.show()
What you see¶
Both curves climb to a success rate of 1.0 — neither method is "better at the task" asymptotically; the update rule is identical. What differs is where each curve rises. Dyna-Q (blue) reaches reliable success after only a few hundred real steps, while plain Q-learning (red) needs several times more real interaction to propagate the single goal reward back across the grid.
Why. In Q-learning a reward signal moves backward only one cell per real visit, so credit assignment is throttled by how often the agent physically revisits states. Dyna-Q learns a tiny tabular model of the transitions it has seen and then replays 20 imagined transitions per real step — those imagined updates diffuse the goal reward through the value table without spending real interaction.
Takeaway. This is the core promise of model-based RL made literal: each real transition is reused many times through the model, so far fewer real samples are needed. In robotics, where real samples are the expensive resource, that is the whole reason to bother learning a model.
Experiment 2: The price of efficiency — sample cost vs compute cost¶
Question. Dyna-Q used fewer real steps — but it wasn't free. What did it cost?
The number of value updates is a clean proxy for compute: model-free does exactly one update per real
step, whereas Dyna-Q does 1 + n_plan. We sweep n_plan and, for each setting, record (a) how many real
steps were needed to first reach a stable success rate, and (b) the total number of value updates spent.
def steps_to_solve(history, thresh=0.95):
# First real-step count at which success reaches >= thresh.
for t, sr in history:
if sr >= thresh:
return int(t)
return int(history[-1, 0])
plans = [0, 1, 3, 5, 10, 20, 40]
real_needed, total_updates = [], []
for npl in plans:
h = train(env, n_plan=npl, real_steps=4000, seed=1)
rn = steps_to_solve(h)
real_needed.append(rn)
total_updates.append(rn * (1 + npl)) # value updates spent up to the solve point
real_needed = np.array(real_needed)
total_updates = np.array(total_updates)
for npl, rn, tu in zip(plans, real_needed, total_updates):
print(f'n_plan={npl:2d} | real steps to solve={rn:5d} | value updates={tu:6d}')
n_plan= 0 | real steps to solve= 4000 | value updates= 4000 n_plan= 1 | real steps to solve= 4000 | value updates= 8000 n_plan= 3 | real steps to solve= 4000 | value updates= 16000 n_plan= 5 | real steps to solve= 4000 | value updates= 24000 n_plan=10 | real steps to solve= 4000 | value updates= 44000 n_plan=20 | real steps to solve= 4000 | value updates= 84000 n_plan=40 | real steps to solve= 4000 | value updates=164000
fig, (axL, axR) = plt.subplots(1, 2, figsize=(11, 4.4))
axL.plot(plans, real_needed, '-o', color='#1f77b4')
axL.set_xlabel('Planning updates per real step (n_plan)')
axL.set_ylabel('Real steps needed to solve')
axL.set_title('More planning -> fewer REAL samples')
axR.plot(plans, total_updates, '-o', color='#d62728')
axR.set_xlabel('Planning updates per real step (n_plan)')
axR.set_ylabel('Total value updates (compute proxy)')
axR.set_title('...but COMPUTE spent goes up')
fig.suptitle('The model-based trade-off: trading real samples for computation', y=1.02)
plt.tight_layout()
plt.show()
What you see¶
Left: as we allow more imagined planning updates per real step, the number of real environment steps needed to solve the task drops sharply, then flattens — the first few planning replays buy most of the sample-efficiency, with diminishing returns after.
Right: the same sweep measured in total value updates (our compute proxy) rises. Past a point, piling on more planning per step costs compute without saving meaningful real samples.
Takeaway. Sample efficiency is not free — it is purchased with computation. This is exactly the model-free vs model-based trade-off from the note: model-free is cheap per step but data-hungry; model-based amortizes scarce real data into abundant compute. Which side wins depends entirely on whether real samples or compute is the binding constraint in your problem.
Experiment 3: Compounding error — why an imperfect model is dangerous¶
Question. A learned model is never exact. How does a small per-step error behave when the model feeds its own predictions back in during a rollout?
We take a true linear latent dynamics $z_{t+1} = A z_t$ and a learned $\hat{A} = A + E$ with a small random perturbation $E$. We roll both out from the same start and measure the rollout error $\lVert \hat{z}_t - z_t \rVert$ as a function of horizon. The note's bound predicts $\lVert e_{t+1}\rVert \le L\lVert e_t\rVert + \epsilon$ with $L$ the model's Lipschitz constant — so the behaviour hinges on whether $L>1$ (error explodes) or $L<1$ (error stays bounded).
def rollout_error(spectral_radius, err_scale=0.03, H=40, d=8, seed=0):
rng = np.random.default_rng(seed)
# build true A with a chosen spectral radius (Lipschitz constant in 2-norm here)
M = rng.standard_normal((d, d))
A = M / np.max(np.abs(np.linalg.eigvals(M))) * spectral_radius
E = rng.standard_normal((d, d)) * err_scale # one-step model error
Ah = A + E
z = rng.standard_normal(d)
zt, zh = z.copy(), z.copy()
errs = []
for _ in range(H):
zt, zh = A @ zt, Ah @ zh
errs.append(np.linalg.norm(zh - zt))
return np.array(errs)
H = 40
err_stable = rollout_error(0.85, H=H, seed=3) # contractive, L<1
err_critical = rollout_error(1.00, H=H, seed=3) # marginal
err_unstable = rollout_error(1.12, H=H, seed=3) # expansive, L>1
print('horizon-40 error L=0.85:', round(err_stable[-1], 3))
print('horizon-40 error L=1.12:', round(err_unstable[-1], 3))
horizon-40 error L=0.85: 0.0 horizon-40 error L=1.12: 0.415
fig, ax = plt.subplots(figsize=(8, 4.6))
hor = np.arange(1, H + 1)
ax.plot(hor, err_stable, '-o', ms=3, color='#2ca02c', label='L = 0.85 (contractive)')
ax.plot(hor, err_critical, '-o', ms=3, color='#ff7f0e', label='L = 1.00 (marginal)')
ax.plot(hor, err_unstable, '-o', ms=3, color='#d62728', label='L = 1.12 (expansive)')
ax.set_yscale('log')
ax.set_xlabel('Rollout horizon (steps)')
ax.set_ylabel('Rollout error ||z_hat - z|| (log scale)')
ax.set_title('Same per-step model error, very different fates along the horizon')
ax.legend()
plt.tight_layout()
plt.show()
What you see¶
All three rollouts start from the same tiny per-step model error, but their long-horizon fate is set by the model's Lipschitz constant $L$. With $L<1$ (green) the dynamics are contractive and the error saturates at a small value. With $L>1$ (red) each step multiplies the accumulated error — on the log axis it grows as a straight line, i.e. geometrically. The marginal $L=1$ case (orange) drifts roughly linearly.
Why. This is the recursion $\lVert e_{t+1}\rVert \le L\lVert e_t\rVert + \epsilon$ from the note made visible: a constant injection $\epsilon$ of fresh one-step error, repeatedly amplified by $L$. The planner rolls the model out on its own predictions, so it lives entirely in this regime — not the teacher-forced one-step regime the model was trained on.
Takeaway. Model-based RL's weakness is structural: even a near-perfect one-step model can produce useless long-horizon plans if dynamics amplify error. This is precisely why practical systems use short rollouts plus a value bootstrap (MBPO) and receding-horizon replanning (MPC) — the rest of Tier 7.
Experiment 4: Why latent is the sweet spot¶
Question. Model-based planning in pixel space is rarely practical. What exactly does moving the model into a compact latent space buy?
Two costs both scale with the dimension $D$ of the space the model operates in:
- Planning compute. Evaluating $N$ candidate action sequences of horizon $H$ costs $\sim N\,H\,D$ for a cheap per-dim transition — linear in $D$. Pixels have $D = H_{\text{img}}\,W_{\text{img}}\,C$ (tens of thousands); a latent has $D = d$ (tens to hundreds).
- Model sample-complexity. Fitting a dynamics model in higher dimension needs more data to reach the same accuracy. We demonstrate this directly by fitting a linear next-step predictor in low vs high dimension and plotting test error against the number of training transitions.
# (a) planning compute vs dimension
N, Hh = 1000, 15 # candidates, horizon
dims = np.array([16, 64, 256, 1024, 4096, 16384, 65536])
plan_cost = N * Hh * dims
d_latent, d_pixel = 64, 64 * 64 * 3 # typical compact latent vs 64x64 RGB frame
# (b) model sample-complexity: fit z' = A z from noisy data in low vs high dim
def fit_error(d, n_train, n_test=500, noise=0.1, seed=0):
rng = np.random.default_rng(seed)
A = rng.standard_normal((d, d)) / np.sqrt(d)
Z = rng.standard_normal((n_train, d))
Zn = Z @ A.T + noise * rng.standard_normal((n_train, d))
Ah, *_ = np.linalg.lstsq(Z, Zn, rcond=None) # least-squares estimate of A
Zt = rng.standard_normal((n_test, d))
pred, true = Zt @ Ah, Zt @ A.T
return np.mean(np.linalg.norm(pred - true, axis=1) ** 2) / d
n_grid = np.array([20, 40, 80, 160, 320, 640, 1280])
err_low = [fit_error(8, n, seed=2) for n in n_grid]
err_high = [fit_error(128, n, seed=2) for n in n_grid]
print('planning cost latent(d=64): {:.2e}'.format(N * Hh * d_latent))
print('planning cost pixel(64x64x3):{:.2e}'.format(N * Hh * d_pixel))
planning cost latent(d=64): 9.60e+05 planning cost pixel(64x64x3):1.84e+08
fig, (axL, axR) = plt.subplots(1, 2, figsize=(11, 4.4))
axL.loglog(dims, plan_cost, '-o', color='#1f77b4')
axL.axvline(d_latent, color='#2ca02c', ls='--', label=f'compact latent d={d_latent}')
axL.axvline(d_pixel, color='#d62728', ls='--', label=f'64x64x3 pixels d={d_pixel}')
axL.set_xlabel('Dimension D of the model space (log)')
axL.set_ylabel('Planning cost ~ N*H*D (log)')
axL.set_title('Planning cost scales with dimension')
axL.legend()
axR.plot(n_grid, err_low, '-o', color='#2ca02c', label='low-dim model (d=8)')
axR.plot(n_grid, err_high, '-o', color='#d62728', label='high-dim model (d=128)')
axR.set_xlabel('Number of training transitions')
axR.set_ylabel('Test one-step error (per dim)')
axR.set_title('Higher-dim models need more data to fit')
axR.legend()
fig.suptitle('Latent dynamics cut both planning compute and model sample-complexity', y=1.02)
plt.tight_layout()
plt.show()
What you see¶
Left: planning cost grows linearly in the dimension of the space the model lives in. A $64\times64\times3$ frame sits three orders of magnitude to the right of a $d=64$ latent — evaluating the same $N=1000$ candidate plans over horizon $15$ is roughly a thousand times cheaper in the latent.
Right: the high-dimensional model (red) needs many more training transitions to reach the same one-step test error as the low-dimensional one (green). Fitting dynamics in pixel space is not just expensive to roll out — it is also data-hungry to learn, which can erase the very sample-efficiency that motivated going model-based in the first place.
Takeaway. A latent world model attacks the model-based trade-off from both sides at once: it shrinks planning compute and the data needed to fit the model, while the encoder discards pixel detail that never mattered for the reward. That is what the note means by latent being the sweet spot — not that the trade-off disappears, but that the operating point moves to where model-based RL becomes worthwhile.
Summary / Key Takeaways¶
- Sample efficiency is real (Exp 1). On an identical task with an identical update rule, Dyna-Q (model-based) reached a good policy in far fewer real steps than Q-learning, by replaying imagined transitions from a learned model.
- It is paid for in compute (Exp 2). More planning per real step keeps cutting real-sample cost but raises total value updates — the model-free/model-based choice is "spend data" vs "spend computation".
- Imperfect models compound (Exp 3). A constant per-step error is amplified by the dynamics' Lipschitz constant; with $L>1$ rollout error grows geometrically, which is why long open-loop model rollouts are unsafe and practical systems use short horizons, value bootstraps, and replanning.
- Latent is the sweet spot (Exp 4). Both planning compute and model sample-complexity scale with the dimension of the model's space; a compact latent slashes both versus pixels, moving the trade-off to where model-based RL pays off — the premise for the rest of Tier 7.