Implement Tabular Dyna-Q: Learn From Real and Simulated Transitions

A Q-learning agent that has already found the goal still has to walk back there to learn from it again. Dyna-Q refuses that tax: it caches what happened and replays it internally, so one real step can feed many updates.

You already have a working tabular Q-learning agent. It updates from real transitions, and it has no memory of the environment beyond the Q-table itself. That is the whole point of model-free learning, and it is also the whole cost: every lesson must be paid for with a real interaction.

Dyna-Q keeps the Q-learning update exactly as it is and adds two things on top. First, a model table that caches the observed outcome of each state-action pair. Second, a planning loop that samples from that cache and applies the same update again, using simulated transitions instead of real ones. The algorithm is not a different learner. It is the same learner with a cheaper source of practice.

This tutorial assumes a deterministic environment with discrete states and actions, and a tabular Q-table. The cache-the-last-observation model is only valid when the same (state, action) always produces the same outcome. We will build the learner, compare it against pure Q-learning, and then break it on purpose to see where the assumption fails.

What You Are Adding to Q-Learning

Recap in two sentences: Q-learning updates Q[s, a] from one real transition using r + gamma * max(Q[s_next]) as the target, and it stores nothing about the environment. Dyna-Q adds a model table keyed by (state, action) that caches the observed (next_state, reward, done), plus a planning loop that samples cached pairs and applies the identical update rule.

The direct RL update and the planning update use the same line of code. Only the source of the transition differs. That symmetry is the reason Dyna-Q is easy to implement once you have Q-learning working.

Three knobs control the behavior: alpha (learning rate), gamma (discount), and k (planning steps per real step). Everything else is bookkeeping.

Set Up the Environment and the Q-Table

You need Python 3 and NumPy. Use a small deterministic gridworld, or the same toy environment from the Q-learning tutorial. The environment needs three things: a reset() that returns a start state, a step(state, action) that returns (next_state, reward, done), and a fixed set of discrete actions.

import numpy as np

n_states, n_actions = 16, 4
Q = np.zeros((n_states, n_actions))
model = {}  # (state, action) -> (next_state, reward, done)

def epsilon_greedy(state, epsilon):
    if np.random.random() < epsilon:
        return np.random.randint(n_actions)
    return int(np.argmax(Q[state]))

def q_update(state, action, next_state, reward, done, alpha, gamma):
    target = reward
    if not done:
        target += gamma * np.max(Q[next_state])
    Q[state, action] += alpha * (target - Q[state, action])

The done flag is not decoration. When a transition ends an episode, bootstrapping from the next state's Q-values teaches the agent that reward continues beyond termination. Suppressing the continuation value on terminal transitions keeps the target honest.

Run a training loop with k = 0 first. Print the episode return or roll out the greedy policy and confirm it reaches the goal. That is your success criterion before any planning is added: a baseline that learns from real experience alone.

Store the Model From Real Transitions

After each real step, write the observed outcome into the model:

model[(state, action)] = (next_state, reward, done)

Overwriting is the entire model-learning rule in the deterministic case. Print model after a few episodes and you will see it only contains pairs the agent has actually tried. Unvisited pairs are absent, not zero. That distinction matters because planning can only start from pairs the model has seen, so early planning is thin and grows with coverage.

Warning: In a stochastic environment, overwriting silently keeps only the last sample and the model becomes a noisy lie. A count-based or averaged model is the fix, but it is out of scope here. Keep the environment deterministic until you have the basic loop working.

Add the Planning Loop

After the real update, run k planning steps. Each step samples a (state, action) pair uniformly from the keys already in the model, retrieves the cached outcome, and applies the same Q-learning update.

def plan(k, alpha, gamma):
    if not model:
        return
    keys = list(model.keys())
    for _ in range(k):
        s, a = keys[np.random.randint(len(keys))]
        s_next, r, done = model[(s, a)]
        q_update(s, a, s_next, r, done, alpha, gamma)

The guard matters. Sampling from an empty key list is the first crash most readers hit, and it happens on the very first real step before the model has any entries.

Wrap the loop in a function so k can be passed in without editing the training code. The interpretation is worth stating plainly: planning updates are cheap, but they reinforce whatever the model currently believes, including its errors.

Compare k = 0 Against k > 0

Run the same seed and the same episode budget twice: once with k = 0 (pure Q-learning) and once with k = 5 or k = 10. Plot episodes-to-threshold or cumulative reward per episode for both runs.

With planning, the curve rises earlier and the agent reaches good greedy behavior in fewer real steps. The gap is largest in the first episodes, when real data is scarce and each cached transition is worth replaying many times.

Two honest caveats. The advantage shrinks as the environment gets larger, because uniform sampling from the model wastes planning updates on irrelevant pairs. And k is a compute-for-data tradeoff: more planning steps cost CPU, not environment interactions. If real steps are cheap and CPU is expensive, Dyna-Q is the wrong trade.

Inspect Model Error, Not Just Reward

Reward curves tell you whether learning worked. They do not tell you whether the model is telling the truth. Turn the model table into a diagnostic object.

Compare each cached (next_state, reward, done) against the true environment transition for the same (state, action). Report two numbers:

In a deterministic environment, accuracy should reach 1.0 for covered pairs. If it does not, the bug is in how the model is written, not in the learning rate.

This is why planning amplifies model error. A wrong cached entry gets replayed k times per real step, so it can dominate the Q-table. When Dyna-Q underperforms plain Q-learning, check model accuracy before tuning alpha or epsilon.

Break It on Purpose: A Changing Environment

Change one transition in the environment mid-training, or make one action stochastic, and rerun with the same k. Watch what happens.

The model keeps serving the stale cached transition until that state-action pair is encountered again in the real environment. Real updates push the Q-values toward the new dynamics; planning updates pull them back toward the old cached entry. The agent fights itself, and more planning makes it worse.

The standard remedies are bonus rewards for long-untried pairs (the Dyna-Q+ heuristic) or prioritized replay of surprising transitions. We are not implementing them here. The point is to see the boundary: this simple cache-the-last-observation model stops being adequate the moment the world stops being deterministic. The algorithm is not broken. The assumption is.

Where to Take This Next

Use Dyna-Q when real interactions are expensive and the environment is close to deterministic. Do not expect it to rescue a bad model.

The first extension worth trying is prioritized sweeping: instead of sampling uniformly, focus planning on pairs whose values just changed. The second is swapping the cached model for a count-based one and rerunning the stochastic variant to watch the accuracy metric move.

Keep the model-accuracy check in the loop. It is the cheapest early warning that planning has started lying to you, and it costs less than a hyperparameter sweep.