Value learning: Q-learning from scratch

Put the states back and the question changes: not which option pays off now, but which actions lead somewhere good later. This section builds the classic answer in full, as a table of what each action in each state is worth, an update rule that nudges those numbers towards what actually happened, and a runnable agent on a small environment. It then covers what changes when the table is replaced by a neural network, and why the naive version of that diverges.

What Q means#

Q(s, a) is the total future reward you expect if you take action a in state s and act well afterwards. If you knew Q exactly, acting well would be trivial: in each state, take the action with the highest Q. Q-learning estimates it from experience.

The difference between Q and the plain value V is worth being precise about, because it is what makes Q the useful one. V(s) says how good a state is, which tells you nothing about what to do unless you also know where each action leads. Q(s, a) says how good each action is from that state, so the policy falls out as an argmax over four numbers and no model of the environment is needed. That is the whole reason value methods are written in terms of Q. It is also the reason they need a small, countable set of actions: you cannot take an argmax over a steering angle.

The update#

Q[s, a] += alpha * ( r + gamma * max(Q[s_next]) - Q[s, a] )
                     \______________________/   \_____/
                       what we just learned     what we believed

alpha   step size: how far to move towards the new estimate
gamma   discount: how much future reward counts (0.99 is common)
r + gamma * max(Q[s_next])   the target: reward now plus the best value from where we landed

The difference in brackets is the temporal difference error: how surprised we were. The update nudges the old belief towards the new evidence. Reward information flows backwards one step at a time, from the goal to the states that lead to it. That is credit assignment in practice.

An environment small enough to check#

Write the environment rather than importing one. Thirteen squares you can decide in means the whole table fits on the screen, and it means you can compute the right answer separately and hold the agent to it, which no benchmark environment lets you do. This is a 4-by-4 grid with a goal worth 1, a hole worth -1, a wall, and a cost of 0.04 per step so that dawdling is penalised. The floor is slippery: an action goes where you intended 80 per cent of the time and sideways otherwise, which is what stops the problem being a maze puzzle.

import random

GRID = ["S...",
        ".#.H",
        "....",
        "...G"]                                  # S start, # wall, H hole, G goal
MOVES = [(-1, 0), (0, 1), (1, 0), (0, -1)]       # up right down left
ARROWS = "^>v<"
GAMMA, STEP_COST, SLIP = 0.95, -0.04, 0.1

def at(s):                                       # the character on square s
    r, c = divmod(s, 4)
    return GRID[r][c]

def is_end(s):
    return at(s) in "GH"

def arrive(s):                                   # reward for arriving on square s
    return {"G": 1.0, "H": -1.0}.get(at(s), STEP_COST)

def slide(s, a):                                 # where action a lands you
    r, c = divmod(s, 4)
    dr, dc = MOVES[a]
    nr, nc = r + dr, c + dc
    if not (0 <= nr < 4 and 0 <= nc < 4) or GRID[nr][nc] == "#":
        return s                                 # walls and edges: you stay put
    return nr * 4 + nc

def outcomes(s, a):                              # the floor is slippery
    p = {}
    for act, prob in ((a, 0.8), ((a + 1) % 4, 0.1), ((a - 1) % 4, 0.1)):
        ns = slide(s, act)
        p[ns] = p.get(ns, 0.0) + prob
    return list(p.items())

def step(s, a, rng):                             # one sampled transition
    u = rng.random()
    if u < SLIP:       a = (a + 1) % 4
    elif u < 2 * SLIP: a = (a - 1) % 4
    ns = slide(s, a)
    return ns, arrive(ns), is_end(ns)

squares = [s for s in range(16) if at(s) != "#"]
inner = [s for s in squares if not is_end(s)]
print(len(inner), "squares to decide in,", len(inner) * 4, "state-action pairs in the table")
# 13 squares to decide in, 52 state-action pairs in the table

Because the transition probabilities are written down, the exact answer is available without any learning. Value iteration applies the Bellman assignment to every square until nothing moves, and what comes out is the best any policy can do. Keep it: it is the ruler that every later claim gets measured against.

V = [0.0] * 16
for sweep in range(1, 1001):
    nxt = list(V)
    for s in inner:
        nxt[s] = max(sum(p * (arrive(ns) + GAMMA * V[ns]) for ns, p in outcomes(s, a))
                     for a in range(4))
    delta = max(abs(nxt[s] - V[s]) for s in inner)
    V = nxt
    if delta < 1e-9:
        break

def greedy_policy(value):
    return {s: max(range(4), key=lambda a: sum(p * (arrive(ns) + GAMMA * value[ns])
                                               for ns, p in outcomes(s, a)))
            for s in inner}

def draw(policy):
    for r in range(4):
        print(" ".join(ARROWS[policy[r * 4 + c]] if r * 4 + c in policy else at(r * 4 + c)
                       for c in range(4)))

best = greedy_policy(V)
print("converged after", sweep, "sweeps")
for r in range(4):
    print(" ".join(f"{V[r * 4 + c]:6.3f}" if at(r * 4 + c) != "#" else "  wall" for c in range(4)))
draw(best)
# converged after 36 sweeps
#  0.481  0.402  0.453  0.231
#  0.576   wall  0.570  0.000
#  0.666  0.765  0.858  0.965
#  0.737  0.847  0.965  0.000
# v < v <
# v # v H
# > > v v
# > > > G

Read the policy and the environment is explained back to you. The top-right corner steers left, away from the hole below it, even though that is the long way round: with a 10 per cent chance of sliding sideways, standing next to a -1 is not worth 0.04 a step. The start square is worth 0.481, which is the number to remember, because everything the learner does from here is judged against it.

The whole thing, runnable#

Now throw the transition probabilities away. The agent gets only the sampler: it acts, it sees where it landed and what it got. Epsilon starts at 1, which is pure random flailing, and decays as the table fills.

def q_learn(seed, episodes, alpha=0.1, init=0.0, mask_terminal=True):
    rng = random.Random(seed)
    Q = [[init] * 4 for _ in range(16)]
    visits = [0] * 16
    eps, env_steps = 1.0, 0
    for ep in range(episodes):
        s = 0
        for t in range(100):
            a = rng.randrange(4) if rng.random() < eps else max(range(4), key=lambda x: Q[s][x])
            ns, r, done = step(s, a, rng)
            target = r + (0.0 if done and mask_terminal else GAMMA * max(Q[ns]))
            Q[s][a] += alpha * (target - Q[s][a])
            visits[s] += 1
            env_steps += 1
            s = ns
            if done:
                break
        eps = max(0.05, eps * 0.999)                     # explore less as the table fills
    return Q, visits, env_steps

Q, visits, env_steps = q_learn(seed=0, episodes=5000)
learned = {s: max(range(4), key=lambda a: Q[s][a]) for s in inner}
print(env_steps, "environment steps; value of the start square:",
      round(max(Q[0]), 3), "against", round(V[0], 3), "from value iteration")
draw(learned)
print("squares where the two policies disagree:", [s for s in inner if learned[s] != best[s]])
# 53882 environment steps; value of the start square: 0.474 against 0.481 from value iteration
# v < < <
# v # v H
# v > v v
# > > > G
# squares where the two policies disagree: [2, 8]

Two squares out of thirteen came out different, and the reason is not that the agent is half-trained. Count where it has been.

def true_value(policy):
    """Exact expected return of a policy, by sweeping its own Bellman equation."""
    Vp = [0.0] * 16
    for _ in range(2000):
        nxt = list(Vp)
        for s in inner:
            nxt[s] = sum(p * (arrive(ns) + GAMMA * Vp[ns]) for ns, p in outcomes(s, policy[s]))
        if max(abs(nxt[s] - Vp[s]) for s in inner) < 1e-12:
            return nxt
        Vp = nxt
    return Vp

print("visits per square, seed 0:")
for r in range(4):
    print(" ".join(f"{visits[r * 4 + c]:6d}" for c in range(4)))

scores = []
for seed in range(20):
    Qs, _, _ = q_learn(seed, 5000)
    pol = {s: max(range(4), key=lambda a: Qs[s][a]) for s in inner}
    scores.append(true_value(pol)[0])
print(f"optimal start value {V[0]:.4f}")
print(f"20 seeds: {sum(1 for v in scores if abs(v - V[0]) < 1e-9)} reached it exactly, "
      f"worst {min(scores):.4f}, mean {sum(scores) / 20:.4f}")
for s in (8,):
    qs = [sum(p * (arrive(ns) + GAMMA * V[ns]) for ns, p in outcomes(s, a)) for a in range(4)]
    print("true action values on square 8:", " ".join(f"{ARROWS[a]}{qs[a]:+.3f}" for a in range(4)))
# visits per square, seed 0:
#   8861   2430    716    264
#   8743      0    747      0
#   7300   4883   4275   1212
#   4148   4364   5939      0
# optimal start value 0.4815
# 20 seeds: 9 reached it exactly, worst 0.4670, mean 0.4751
# true action values on square 8: ^+0.534 >+0.666 v+0.656 <+0.591

The start square was visited 8,861 times and the square in the top-right corner 264 times, a factor of 34, and the two disagreements are in the thin part of that map. This is the third property of reinforcement learning made visible: the agent's own policy decides where its data comes from, so it learns the squares it walks through and stays vague about the ones it avoids. It avoids them because they are bad, which is usually fine and occasionally the reason it never finds something better.

The second disagreement is a different thing entirely, and it is why you should measure policies rather than compare them square by square. On square 8 the exact values of going right and going down are 0.666 and 0.656 — a gap of 0.01 that no finite amount of sampling will reliably resolve, and which costs almost nothing to get wrong. Across 20 seeds, 9 recovered the optimal policy exactly and the worst was worth 0.4670 against the optimum of 0.4815, a loss of 3 per cent. So more than half the seeds produce a policy that is technically wrong and practically fine. Reporting the fraction of exactly-correct seeds would have made the algorithm look far worse than it is.

That bug is worth running rather than believing, because it has a condition attached that nobody mentions. The q_learn above takes a mask_terminal flag; turn it off and the target bootstraps from the terminal square's row of the table, which should be nonsense. Whether it is depends on what that row was initialised to.

for init in (0.0, 1.0):
    for mask in (True, False):
        vals = []
        for seed in range(10):
            Qs, _, _ = q_learn(seed, 3000, init=init, mask_terminal=mask)
            vals.append(Qs[3][2])                  # square (0,3), action down, into the hole
        print(f"init {init}, mask_terminal={mask}: Q(into the hole) = {sum(vals) / 10:+.3f}")
# init 0.0, mask_terminal=True: Q(into the hole) = -0.817
# init 0.0, mask_terminal=False: Q(into the hole) = -0.817
# init 1.0, mask_terminal=True: Q(into the hole) = -0.747
# init 1.0, mask_terminal=False: Q(into the hole) = +0.173

With zeros everywhere the missing mask adds gamma times zero, so the two columns agree to three decimal places and the bug hides completely — which is how it survives review and reaches the run where it matters. Initialise optimistically at 1.0, a standard trick for encouraging exploration, and stepping into a hole is valued at +0.173 instead of -0.747: a -1 reward followed by 0.95 times an imaginary future. The general lesson is the one to carry: a bug whose effect is multiplied by an initialisation is a bug that only appears when something else changes.

The same agent against a standard environment#

Once it works on a grid you wrote, the identical algorithm runs against the environments papers use. Gymnasium's FrozenLake is the same shape of problem, and the only real differences are that the environment owns the state and that reset and step return tuples rather than plain values. This needs gymnasium installed; it is not part of the standard library.

import gymnasium as gym
import numpy as np

env = gym.make("FrozenLake-v1", is_slippery=False)       # 16 states, 4 actions
Q = np.zeros((env.observation_space.n, env.action_space.n))
alpha, gamma, eps = 0.1, 0.99, 1.0

for episode in range(5000):
    s, _ = env.reset()
    done = False
    while not done:
        a = env.action_space.sample() if np.random.rand() < eps else int(Q[s].argmax())
        s_next, r, terminated, truncated, _ = env.step(a)
        target = r + gamma * Q[s_next].max() * (not terminated)
        Q[s, a] += alpha * (target - Q[s, a])
        s, done = s_next, terminated or truncated
    eps = max(0.05, eps * 0.999)                         # explore less over time

print(Q.argmax(axis=1).reshape(4, 4))                    # the learned policy

When the table is too big: DQN#

For images or large state spaces, a table is impossible, so a neural network predicts Q values from the state. Training it naively diverges. The DQN work that learned Atari games from pixels made it stable with two changes.

• Experience replay: store transitions in a buffer and train on random samples from it, which breaks the strong correlation between consecutive steps.
• A target network: compute targets with a copy of the network that is updated only occasionally, so the thing you are chasing does not move with every step.

Both fixes address the same thing, which is that the table version had a property the network version quietly loses. Updating one cell of a table changes one cell. Updating a network changes its prediction everywhere, including at the state whose value you are using as the target, so the target moves as you chase it. Add to that a max over noisy estimates, which is biased upwards because the largest of several noisy numbers tends to be one that was overestimated — the overestimation that double DQN exists to correct — and the naive version diverges rather than merely learning slowly.

What to take from this section and what to leave. The update rule, the meaning of each term, the off-policy property and the two DQN fixes are standard interview ground and worth knowing cold. The rest of the value-method literature — duelling architectures, prioritised replay, distributional value estimates, Rainbow's combination of all of them — is a body of work aimed at Atari scores, and reading it will not help you unless you are doing exactly that. The habits transfer better than the algorithms: compute the exact answer on a small case, count your visits, judge a policy by its measured value rather than by matching a reference policy square by square.