8.5 Reinforcement learning

In reinforcement learning nobody gives the right answers. An agent tries actions, receives rewards, and learns which actions pay off, balancing the urge to use what it already knows against the need to try things that might be better.

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.axes import Axes
from scipy import stats

from statly import style

style.use()
rng = np.random.default_rng(1)

8.5.1 Learning from rewards

Supervised learning (Part 6) learns from examples labelled with the right answer. Reinforcement learning learns from consequences. An agent observes the state of its environment, chooses an action, receives a reward and a new state, and repeats. Its goal is a policy, a rule for choosing actions, that collects as much reward as possible over time. Game-playing programs, robot control and the final tuning stage of language models (Chapter 8.3) all work this way.

fig, ax = plt.subplots(figsize=(7, 2.6))
for x, name, color in ((0, "agent\n(policy)", style.BLUE), (4, "environment", style.GREEN)):
    ax.add_patch(plt.Rectangle((x - 0.9, -0.5), 1.8, 1, color=color, alpha=0.35))
    ax.text(x, 0, name, ha="center", va="center", fontsize=10)
ax.annotate(
    "",
    xy=(3.1, 0.25),
    xytext=(0.9, 0.25),
    arrowprops=dict(arrowstyle="->", color=style.INK, lw=1.5),
)
ax.text(2, 0.42, "action", ha="center", fontsize=9)
ax.annotate(
    "",
    xy=(0.9, -0.25),
    xytext=(3.1, -0.25),
    arrowprops=dict(arrowstyle="->", color=style.INK, lw=1.5),
)
ax.text(2, -0.6, "reward and new state", ha="center", fontsize=9)
ax.set(xlim=(-1.2, 5.2), ylim=(-0.9, 0.8))
ax.set_axis_off()
plt.show()

Figure 8.5.1: The reinforcement learning loop. The agent acts on the environment; the environment returns a reward and a new state; the agent updates its policy and acts again.

8.5.2 Explore or exploit: the multi-armed bandit

The simplest case has no states, just repeated choices. Three versions of a web page have unknown click rates; each visitor sees one version and clicks or not. Always showing the version that has done best so far (exploiting) can lock in an early lucky winner; showing versions at random (exploring) wastes visitors on poor versions forever. Good strategies explore less as they learn more.

• ε-greedy: show the best version so far, except with a small probability ε choose at random.
• Thompson sampling: keep a Bayesian posterior for each click rate (Chapter 3.12), draw one value from each posterior, and show the version whose draw is highest. Uncertain versions sometimes draw high and get tried; versions that are clearly worse almost never do.

RATES = np.array([0.04, 0.05, 0.06])
VISITORS, RUNS = 10_000, 200


def run_bandit(strategy: str, seed: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Cumulative regret, and click and show counts per version, for one simulated run."""
    local = np.random.default_rng(seed)
    clicks, shows = np.zeros(3), np.zeros(3)
    regret = np.empty(VISITORS)
    lost = 0.0
    for t in range(VISITORS):
        if strategy == "random":
            arm = local.integers(3)
        elif strategy == "epsilon":
            arm = (
                local.integers(3)
                if local.uniform() < 0.1 or shows.min() == 0
                else int(np.argmax(clicks / shows))
            )
        else:
            arm = int(np.argmax(local.beta(1 + clicks, 1 + shows - clicks)))
        clicks[arm] += local.uniform() < RATES[arm]
        shows[arm] += 1
        lost += RATES.max() - RATES[arm]
        regret[t] = lost
    return regret, clicks, shows


STRATEGIES = {"random": style.GREY, "epsilon": style.GOLD, "thompson": style.BLUE}
NAMES = {"random": "random choice", "epsilon": "ε-greedy", "thompson": "Thompson sampling"}
fig, axes = style.subplots(2, height=3.2, width=10.5)
for strategy, color in STRATEGIES.items():
    regrets = np.mean([run_bandit(strategy, s)[0] for s in range(RUNS)], axis=0)
    axes[0].plot(np.arange(1, VISITORS + 1), regrets, color=color, lw=2, label=NAMES[strategy])
axes[0].legend(fontsize=8)
axes[0].set(xlabel="visitors", ylabel="clicks lost (regret)")
_, clicks, shows = run_bandit("thompson", 0)
grid = np.linspace(0.02, 0.09, 400)
for k, color in enumerate((style.CORAL, style.GOLD, style.GREEN)):
    posterior = stats.beta(1 + clicks[k], 1 + shows[k] - clicks[k])
    axes[1].plot(
        grid,
        posterior.pdf(grid),
        color=color,
        lw=2,
        label=f"true {RATES[k]:.0%}, shown {int(shows[k]):,}",
    )
axes[1].xaxis.set_major_formatter(style.PERCENT)
axes[1].legend(fontsize=8)
axes[1].set(xlabel="click rate", ylabel="how plausible", yticks=[])
fig.tight_layout()
plt.show()

Figure 8.5.2: Three page versions with true click rates of 4%, 5% and 6%, shown to 10,000 visitors, averaged over 200 simulated runs. Left: regret, the clicks lost compared with always showing the best version. Random choice loses clicks steadily; ε-greedy (ε = 0.1) learns but keeps exploring at a fixed rate; Thompson sampling’s regret flattens as it settles on the best version. Right: one run of Thompson sampling, the posterior for each click rate after 10,000 visitors; the best version has been shown most and has the narrowest posterior.

Bandits are an alternative to the fixed-size A/B test of Chapter 5.1: they lose fewer clicks while learning, at the price of a less clean comparison at the end.

8.5.3 States, values and the Bellman equation

With states, an action’s worth depends on where it leads, and rewards may come only much later. The value of taking action a in state s, written Q(s, a), is the expected total of future rewards, with rewards further in the future discounted by a factor γ < 1 per step. The values satisfy the Bellman equation: the value of an action is its immediate reward plus the discounted value of the best action in the state it leads to,

$$Q(s, a) = E\big[r + \gamma \max_{a'} Q(s', a')\big]$$

Q-learning estimates these values by experience: after each step, nudge the current estimate towards the reward received plus the discounted best estimate for the new state,

$$Q(s, a) \leftarrow Q(s, a) + \alpha\big(r + \gamma \max_{a'} Q(s', a') - Q(s, a)\big)$$

and act mostly greedily on the current estimates, with some exploration.

ROWS, COLS = 5, 6
WALLS = {(1, 1), (2, 1), (3, 3), (1, 4)}
GOAL, PIT, START = (0, 5), (2, 4), (4, 0)
MOVES = [(-1, 0), (1, 0), (0, -1), (0, 1)]  # up, down, left, right
ARROWS = ["↑", "↓", "←", "→"]


def step(state: tuple[int, int], action: int) -> tuple[tuple[int, int], float, bool]:
    """Move if possible; return the new state, the reward and whether the episode ended."""
    r, c = state[0] + MOVES[action][0], state[1] + MOVES[action][1]
    if not (0 <= r < ROWS and 0 <= c < COLS) or (r, c) in WALLS:
        r, c = state
    if (r, c) == GOAL:
        return (r, c), 1.0, True
    if (r, c) == PIT:
        return (r, c), -1.0, True
    return (r, c), -0.01, False


def q_learning(
    episodes: int = 3000, alpha: float = 0.2, gamma: float = 0.95, epsilon: float = 0.1
) -> tuple[np.ndarray, list[int]]:
    """Tabular Q-learning with ε-greedy exploration."""
    local = np.random.default_rng(0)
    q = np.zeros((ROWS, COLS, 4))
    lengths = []
    for _ in range(episodes):
        state, done, steps = START, False, 0
        while not done and steps < 200:
            action = (
                local.integers(4) if local.uniform() < epsilon else int(np.argmax(q[state]))
            )
            new_state, reward, done = step(state, action)
            target = reward + (0 if done else gamma * q[new_state].max())
            q[state][action] += alpha * (target - q[state][action])
            state, steps = new_state, steps + 1
        lengths.append(steps)
    return q, lengths


q, lengths = q_learning()
values = q.max(axis=2)
fig, axes = style.subplots(2, height=3.4, width=10.5)
shown = np.where(
    np.isin(np.arange(ROWS * COLS), [r * COLS + c for r, c in WALLS]).reshape(ROWS, COLS),
    np.nan,
    values,
)
axes[0].imshow(shown, cmap="Blues", vmin=-0.3, vmax=1.8)
for r in range(ROWS):
    for c in range(COLS):
        if (r, c) in WALLS:
            axes[0].add_patch(plt.Rectangle((c - 0.5, r - 0.5), 1, 1, color=style.GREY))
        elif (r, c) == GOAL:
            axes[0].text(c, r, "G", ha="center", va="center", fontsize=12, fontweight="bold")
        elif (r, c) == PIT:
            axes[0].add_patch(
                plt.Rectangle((c - 0.5, r - 0.5), 1, 1, color=style.CORAL, alpha=0.7)
            )
        else:
            label = ARROWS[int(np.argmax(q[r, c]))] + ("\nS" if (r, c) == START else "")
            axes[0].text(c, r, label, ha="center", va="center", fontsize=12, color="black")
axes[0].set(xticks=[], yticks=[], title="Learned values and policy")
axes[0].grid(False)
blocks = np.array(lengths).reshape(-1, 50).mean(axis=1)
axes[1].plot(np.arange(len(blocks)) * 50 + 25, blocks, color=style.BLUE, lw=2)
axes[1].set(xlabel="episode", ylabel="steps to finish", title="Learning")
fig.tight_layout()
plt.show()

Figure 8.5.3: Q-learning in a small grid world. The agent starts at the bottom left (S) and is rewarded +1 for reaching the goal (G); stepping into the pit (red) costs 1 and ends the episode; walls are grey; every step costs 0.01. Left: the learned value of the best action in each cell after 3,000 episodes, and the action it would take (arrows). The values fall with distance from the goal, and the arrows lead around the pit. Right: steps per episode while learning (average over blocks of 50 episodes); early wandering gives way to a nearly shortest route.

8.5.4 From tables to networks

A table of values works for a few dozen states. A board game or a robot has far too many states to list, so deep reinforcement learning replaces the table by a neural network (Chapter 8.1) that takes a state and outputs values or action probabilities, and generalises from states it has seen to similar ones. Policy gradient methods learn the policy directly, nudging the probability of each action in proportion to how much better than expected the following rewards turned out.

In reinforcement learning from human feedback, used to tune language models (Chapter 8.3), the “environment” is a reward model trained to predict which of two answers people prefer, and the policy is the language model itself.

Reinforcement learning is powerful and brittle: rewards that are sparse or poorly specified lead to slow learning or to agents that exploit loopholes in the reward instead of doing the intended task, and results can vary a lot between random seeds.

8.5.5 In one picture

• An agent acts, receives rewards, and learns a policy that maximises future reward.
• Bandits: balance exploration and exploitation; Thompson sampling draws from Bayesian posteriors; judge by regret.
• Values obey the Bellman equation $Q(s,a) = E[r + \gamma \max_{a'} Q(s',a')]$; Q-learning estimates them from experience.
• Deep reinforcement learning replaces tables by networks; RLHF tunes language models with a learned reward.

Next: Part 9, data that come in order, through time and until an event.