Model-Based vs Model-Free

Model-based and model-free reinforcement learning apply the same equation; they differ only in where the next state comes from.

The Bellman optimality equation needs, for each state and action, a next state and a reward. Value iteration reads them from a known model. Q-learning reads them from the real environment, one sample at a time. Dyna-Q reads them from both. This post derives all three from one equation, codes them on one maze, and measures the cost of each in the only currency that matters when the world is expensive: real environment steps.

The equation underneath#

For a deterministic MDP with discount $\gamma$, the optimal action value satisfies

where $f(s, a)$ is the next state. The right-hand side is a backup: it computes a better estimate of $Q(s, a)$ from estimates one step ahead. Every algorithm below is a schedule of backups. The only question is how each obtains the pair $(f(s, a), r(s, a))$:

The maze#

The test bed is a 6x9 maze modelled on the one in Sutton and Barto's treatment of Dyna: start at row 2 on the left, goal in the top-right corner, a few wall segments in between, reward 1 only on reaching the goal, $\gamma = 0.95$. The shortest path is 14 steps.

import numpy as np
 
ROWS, COLS = 6, 9
WALLS = {(1, 2), (2, 2), (3, 2), (4, 5), (0, 7), (1, 7), (2, 7)}
START, GOAL = (2, 0), (0, 8)
MOVES = [(-1, 0), (1, 0), (0, -1), (0, 1)]
GAMMA = 0.95
S = [(r, c) for r in range(ROWS) for c in range(COLS) if (r, c) not in WALLS]
IDX = {s: i for i, s in enumerate(S)}
 
def env_step(s, a):
    r, c = s[0] + MOVES[a][0], s[1] + MOVES[a][1]
    if not (0 <= r < ROWS and 0 <= c < COLS) or (r, c) in WALLS:
        r, c = s
    s2 = (r, c)
    return s2, float(s2 == GOAL), s2 == GOAL

The sparse reward makes the comparison sharp. Nothing is learned until the goal is reached once, and after that, value has to propagate backward one backup at a time across 14 cells.

Value iteration: the model is given#

If env_step may be called freely, as an oracle rather than as experience, then a full sweep over states applies the Bellman backup everywhere at once:

def value_iteration(tol=1e-6):
    V = np.zeros(len(S)); sweeps = 0
    while True:
        sweeps += 1; delta = 0.0
        for s in S:
            if s == GOAL: continue
            best = max(r + (0 if d else GAMMA * V[IDX[s2]])
                       for s2, r, d in (env_step(s, a) for a in range(4)))
            delta = max(delta, abs(best - V[IDX[s]])); V[IDX[s]] = best
        if delta < tol: return V, sweeps

It converges in 13 sweeps, with $V(\text{start}) = 0.5133 = 0.95^{13}$: the value of a reward received on step 14. Real environment steps used: zero. That number is not a free lunch. It is the price of having been handed $f$. The cost has moved from interaction into knowledge, and in most real problems that knowledge does not exist.

The limitation is the assumption itself. Value iteration requires that every $(s, a)$ can be queried, which is exactly what a robot, a game or a clinic does not allow.

Q-learning: no model at all#

Q-learning drops the assumption. It replaces the exact backup with a sampled, incremental one, using whatever transition the agent just experienced:

The expectation over next states is replaced by a single sample, and the full update by a step of size $\alpha$. Under standard conditions on exploration and step sizes it converges to $Q^*$ without ever representing $f$. The price is that each real step funds exactly one backup, and the backup touches exactly one $(s, a)$ entry. When the agent first reaches the goal, only the final cell learns anything. The next episode pushes value one more cell back, and so on.

Dyna-Q: remember, then rehearse#

Dyna-Q keeps the Q-learning update and adds two lines. After each real transition it stores (s, a) -> (s', r) in a table (a lookup-table world model), and then performs $n$ additional backups on transitions sampled from that table. With n_plan = 0 the function below is plain Q-learning.

def greedy_len(Q, limit=200):
    s = START
    for t in range(limit):
        s, _, done = env_step(s, int(np.argmax(Q[IDX[s]])))
        if done: return t + 1
    return None
 
def dyna_q(n_plan, episodes, rng, alpha=0.1, eps=0.1):
    Q = np.zeros((len(S), 4)); model = {}; steps_per_ep = []; total = 0
    first_opt = None
    for ep in range(episodes):
        s, done, t = START, False, 0
        while not done:
            i = IDX[s]
            a = int(rng.integers(4)) if rng.random() < eps else \
                int(rng.choice(np.flatnonzero(Q[i] == Q[i].max())))
            s2, r, done = env_step(s, a)                       # one REAL step
            Q[i, a] += alpha * (r + (0 if done else GAMMA * Q[IDX[s2]].max()) - Q[i, a])
            model[(s, a)] = (s2, r, done)                      # remember it
            keys = list(model)
            for _ in range(n_plan):                            # n IMAGINED steps
                ps, pa = keys[rng.integers(len(keys))]
                ps2, pr, pd = model[(ps, pa)]
                j = IDX[ps]
                Q[j, pa] += alpha * (pr + (0 if pd else GAMMA * Q[IDX[ps2]].max()) - Q[j, pa])
            s = s2; t += 1
        total += t; steps_per_ep.append(t)
        if first_opt is None and (greedy_len(Q) or 999) <= 16:
            first_opt = total
    return np.array(steps_per_ep), first_opt

The imagined backups are the same equation as the real ones. They differ only in where (s', r) came from. That is the entire conceptual content of "model-based": a model is a way to run more backups per unit of real experience.

The measurement#

Each configuration runs for 50 episodes with $\epsilon = 0.1$, $\alpha = 0.1$, averaged over 30 seeds. "Steps to near-optimal" is the median number of real environment steps until the greedy policy reaches the goal in at most 16 steps (optimal is 14).

V, sweeps = value_iteration()
print("VI sweeps:", sweeps, "V(start)=%.4f" % V[IDX[START]], "real steps: 0")
for n in [0, 5, 50]:
    curves, firsts = [], []
    for seed in range(30):
        c, f = dyna_q(n, 50, np.random.default_rng(seed))
        curves.append(c); firsts.append(f if f is not None else np.nan)
    C = np.mean(curves, 0)
    print(f"n={n:2d}", [round(C[k]) for k in (0, 1, 4, 9)], round(C[49], 1),
          round(np.mean([c.sum() for c in curves])),        # real steps over 50 episodes
          np.nanmedian(firsts), np.isnan(firsts).sum())     # steps to near-optimal, failures

Columns "ep k" are mean steps taken in episode $k$.

Three things are visible.

• Episode 1 is the same for everyone. Before the goal has been found, there is no reward to propagate, and imagined backups of zero are still zero. A model cannot plan toward something it has never seen. Exploration is not replaced by planning.
• Episode 2 is where the model pays. Once the goal is in the table, fifty rehearsed backups per real step drag value across the whole visited region in a single episode. Q-learning needs dozens of episodes to do the same, one cell at a time.
• The returns diminish. Going from 5 to 50 planning steps cuts episode 2 from 166 to 40 steps but barely changes the total. After value has propagated, extra imagined backups refine an already-correct policy. A few seeds in each of the $n = 0$ and $n = 50$ runs ended with a greedy path just over 16 steps; with $\epsilon$-greedy behaviour and a fixed step size, that residue is expected.

The numbers translate to a factor of about 2.6 fewer real steps over 50 episodes, and five to eight times fewer to reach a near-optimal greedy policy. That ratio is the case for world models in its smallest form.

Which backups to imagine#

Uniform sampling from the table wastes most imagined backups on pairs whose values are not changing. Prioritized sweeping orders them instead: after a real step, push the pair onto a queue keyed by the size of its TD error, pop the most surprising pair, back it up, and then queue every remembered predecessor that leads into it. Value flows backward from the goal along the paths that actually lead there.

import heapq
 
def td(Q, model, s, a):
    s2, r, d = model[(s, a)]
    return r + (0 if d else GAMMA * Q[IDX[s2]].max()) - Q[IDX[s], a]
 
def sweep(Q, model, preds, pq, n_plan, alpha, theta=1e-4):
    for _ in range(n_plan):
        if not pq: break
        _, ps, pa = heapq.heappop(pq)                      # most surprising pair first
        Q[IDX[ps], pa] += alpha * td(Q, model, ps, pa)
        for qs, qa in preds.get(ps, ()):                   # everyone who leads into ps
            p = abs(td(Q, model, qs, qa))
            if p > theta: heapq.heappush(pq, (-p, qs, qa))

Measured on the same maze with $n = 5$, $\alpha = 0.5$ and 30 seeds, prioritized sweeping cut episode 2 from 166 steps to 58, but the total over 50 episodes did not improve (1864 against 1829 for uniform Dyna-Q). With $\alpha = 1$ it got worse: once a pair's error is consumed it leaves the queue, and the queue can run dry while the policy is still poor. Ordering matters most in large state spaces where uniform rehearsal is hopeless; in a 47-cell maze, brute repetition is already cheap.

Where the free lunch ends#

Dyna-Q works this well because its model is perfect: the maze is deterministic and the table stores observed transitions verbatim. Every imagined backup is a true one. Three things break that.

1. Stochastic dynamics. Storing the last observed s' is wrong when the next state is random. The table must store counts or a distribution, and imagined backups become sampled rather than exact.
2. A changing world. If a wall moves, the table keeps rehearsing the old maze. Sutton and Barto's Dyna-Q+ adds an exploration bonus that grows with the time since a pair was last tried, so stale entries get re-tested in reality.
3. Function approximation. With a learned network instead of a table, imagined transitions can be wrong in ways the real data never shows. Backups through those errors are confidently wrong, and the planner exploits them. The usual remedies are short imagined rollouts branched from real states, and ensembles that measure their own disagreement.

The model-free side has its own answer to sample inefficiency, replay buffers, which are a non-parametric model in disguise: replaying stored transitions is Dyna with a model that can only reproduce what it saw. The line between the two families is not "model or no model". It is how far the model is allowed to generalise beyond the data.

A checklist#

• Report real environment steps, not wall-clock time or gradient updates, when comparing sample efficiency.
• Report planning compute alongside it. Dyna-Q with $n = 50$ does 51 backups per real step; that is the trade, not a free win.
• Separate the first-success phase from the propagation phase. Models speed up the second; only exploration fixes the first.
• Test with the world changed mid-run. A model that cannot notice it is wrong is worse than no model.
• When moving from tables to networks, keep imagined horizons short and branch them from real states, as in Dreamer, and measure how fast imagined rollouts diverge before trusting them, as in evaluating a world model.
• The next step up, planning online with the model at every decision rather than distilling it into $Q$, is model-predictive control.