Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Model-Based Reinforcement Learning › Code Implementation: Dyna Against Model-Free Q-Learning

Michael BrenndoerferJune 12, 20267 min read

World Models HandbookAll chapters

Part 2 of 12Inference, Decisions, and Control

Chapter 12Model-Based Reinforcement Learning › Code Implementation: Dyna Against Model-Free Q-Learning

Reading section 7 of 11 · Code Implementation: Dyna Against Model-Free Q-Learning · 7 minSection 7 of 11 · 7 min

Section 7 of 11 · 7 min read

Model-Based Reinforcement Learning › Code Implementation: Dyna Against Model-Free Q-LearningOne section · Sections in this chapter7 / 11 ↓

Sections in this chapter11

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

We now build the worked example in full. The plan: implement the gridworld and a tabular Dyna agent, run three variants (no planning, 10 planning steps, 30 planning steps) across many random seeds, and plot episode return against real environment interactions, which is the resource we care about. Plotting against episodes instead would flatter nothing and show nothing: an episode is a variable-length bundle of interactions, and the whole point of model-based methods is their exchange rate between computation and reality.

First, the environment. It is deterministic, with bounded grid movement and a terminal goal state. The clipping at the walls means attempted moves off the grid result in staying in place, which gives the agent a small but real amount of transition structure to learn about the boundary.

In[13]:

Code

from collections import defaultdict

GRID_SIZE = 6
START = (0, 0)
GOAL = (GRID_SIZE - 1, GRID_SIZE - 1)
ACTIONS = [(-1, 0), (1, 0), (0, -1), (0, 1)]
STEP_REWARD = -1.0
GOAL_REWARD = 10.0


def env_step(state, action_idx):
    dr, dc = ACTIONS[action_idx]
    r = int(np.clip(state[0] + dr, 0, GRID_SIZE - 1))
    c = int(np.clip(state[1] + dc, 0, GRID_SIZE - 1))
    next_state = (r, c)
    done = next_state == GOAL
    reward = GOAL_REWARD if done else STEP_REWARD
    return next_state, reward, done

Next, the agent. Q maps each state to four action values, defaulting to zeros for unvisited states. The model dictionary maps visited $(s, a)$ pairs to outcomes; because the environment is deterministic, one observation per pair makes the model exact. The plan_steps parameter is Dyna's dial: set it to zero and this is plain Q-learning.

In[14]:

Code

def dyna_episode(
    Q,
    model,
    behavior_rng,
    planning_rng,
    alpha=0.5,
    gamma=0.95,
    eps=0.1,
    plan_steps=0,
):
    """Run one episode: epsilon-greedy control, tabular Q-learning,
    and plan_steps imagined backups per real step."""
    state = START
    total_reward, env_steps = 0.0, 0
    while state != GOAL and env_steps < 250:
        if behavior_rng.random() < eps:
            action_idx = int(behavior_rng.integers(len(ACTIONS)))
        else:
            greedy = np.flatnonzero(Q[state] == np.max(Q[state]))
            action_idx = int(behavior_rng.choice(greedy))
        next_state, reward, done = env_step(state, action_idx)

        target = reward if done else reward + gamma * np.max(Q[next_state])
        Q[state][action_idx] += alpha * (target - Q[state][action_idx])
        model[(state, action_idx)] = (next_state, reward, done)

        keys = list(model.keys())
        for _ in range(plan_steps):
            ps, pa = keys[planning_rng.integers(len(keys))]
            pnext, preward, pdone = model[(ps, pa)]
            ptarget = preward if pdone else preward + gamma * np.max(Q[pnext])
            Q[ps][pa] += alpha * (ptarget - Q[ps][pa])

        state = next_state
        total_reward += reward
        env_steps += 1
    return total_reward, env_steps

Planning pairs are sampled uniformly from visited pairs. Separate random-number streams keep exploratory actions paired across planning budgets until their learned values cause the policies to diverge; planning draws no longer advance the behavior stream. Random tie-breaking also avoids a fixed preference for action 0.

Now the experiment. We run each variant for 15 seeds and 120 episodes, recording the return of every episode and the cumulative count of real environment steps, then average across seeds. This cell is the expensive one; everything after it only rearranges these arrays.

In[15]:

Code

def run_experiment(plan_steps, n_seeds=15, n_episodes=120):
    """Return per-seed cumulative real steps and episode returns."""
    all_returns = np.zeros((n_seeds, n_episodes))
    all_cum_steps = np.zeros((n_seeds, n_episodes))
    for seed in range(n_seeds):
        behavior_rng = np.random.default_rng(seed)
        planning_rng = np.random.default_rng(10_000 + seed)
        Q = defaultdict(lambda: np.zeros(len(ACTIONS)))
        model = {}
        cum_steps = 0
        for ep in range(n_episodes):
            ret, steps = dyna_episode(
                Q, model, behavior_rng, planning_rng, plan_steps=plan_steps
            )
            cum_steps += steps
            all_returns[seed, ep] = ret
            all_cum_steps[seed, ep] = cum_steps
    return all_cum_steps, all_returns


steps_q, returns_q = run_experiment(plan_steps=0)
steps_d10, returns_d10 = run_experiment(plan_steps=10)
steps_d30, returns_d30 = run_experiment(plan_steps=30)

Before plotting, let's compute a crisp summary: for each variant, how many real environment interactions were needed before the smoothed return first crossed a fixed threshold. This converts the learning curves into a single interpretable number per method, one that can be compared and quoted, rather than eyeballed off a chart.

In[16]:

Code

def steps_to_threshold_per_seed(steps, returns, threshold=-2.0, window=10):
    """Per-seed real steps to the first qualifying moving-average return."""
    crossings = []
    kernel = np.ones(window) / window
    for seed_steps, seed_returns in zip(steps, returns):
        smooth = np.convolve(seed_returns, kernel, mode="valid")
        crossed = np.flatnonzero(smooth >= threshold)
        crossings.append(
            np.nan if len(crossed) == 0 else seed_steps[crossed[0] + window - 1]
        )
    return np.asarray(crossings)

Out[17]:

Console

Q-learning: 15/15 seeds reached -2; mean  846.8 +/-  21.9 real steps
   Dyna-10: 15/15 seeds reached -2; mean  273.1 +/-  14.5 real steps
   Dyna-30: 15/15 seeds reached -2; mean  237.0 +/-   9.6 real steps

The threshold is computed per seed before aggregation, so a crossing in one run cannot hide a failure in another. The plot below interpolates each seed's smoothed returns onto one common real-interaction grid before taking the mean and standard error. It does not extrapolate a seed before its first complete 10-episode window; a method's curve begins only where all 15 seeds contribute observations.

Out[18]:

Visualization

Learning curves with standard-error bands showing both Dyna variants improving earlier than Q-learning on a shared interaction axis. — Mean 10-episode-smoothed return on a common real-interaction grid for Q-learning and two Dyna variants across 15 seeds; shaded bands show one standard error. Each curve begins only once all 15 seeds for that method have completed a smoothing window. In this deterministic visited-pair model, both planning variants reach a -2 return threshold in about one-third of Q-learning's real interactions.

In this exact deterministic tabular setting, the two planning variants reach the displayed threshold in roughly one-third of Q-learning's real interactions. The difference between 10 and 30 planning updates is smaller than the difference between either Dyna variant and no planning. This is one experiment, not a monotonic law: with an approximate model, more planning can reinforce bias, while with an exact model and a compute-free objective, no theorem here requires a finite optimal planning budget.

You now have a complete Dyna-style MBRL system: model learning, imagined backups, and explicit accounting of real interactions. Other model-based methods may instead plan directly or differentiate through a model without this value-learning loop.