Everything RLHF, DPO, and GRPO are built on, from scratch. MDPs and Bellman equations, policy gradients and their variance, baselines and critics, generalized advantage estimation, trust regions and PPO, and the language model seen as a policy.

Suppose you want a language model to write programs that pass a hidden test suite. Supervised fine-tuning (Chapter 14) would need a correct program for every prompt, written by someone, so the model can copy it token by token. But for most interesting prompts nobody has that program. What you do have is the test runner. Give it any program and it says pass or fail.

That is a different kind of signal. It does not tell the model what to write. It tells the model how well it did, after the fact, for the one thing it tried. Learning from that kind of signal is reinforcement learning (RL), and it is how reasoning models learn to solve hard problems (§18.3) and how chat models learn to be helpful from human preferences (Chapter 16).

This chapter builds RL from nothing, for a reader who has never seen it, and ends with the exact objective that RLHF, DPO, and GRPO optimize. This first section sets up the vocabulary (agents, environments, states, actions, rewards, returns, policies, and value functions) and then studies the simplest problem where the central difficulty of RL already shows up, the choice between exploiting what you know and exploring to learn more.

Learning from a score instead of an answer

In supervised learning every input comes with the right output. The loss tells the model exactly which way to move, and the training examples do not depend on the model. Reinforcement learning breaks all three of those comforts.

• The feedback is evaluative, not instructive. A reward says how good the chosen action was. It says nothing about the actions that were not chosen, so the only way to learn about them is to try them.
• The consequences are delayed. A chess move can win or lose the game forty moves later, and a single wrong line of code fails the whole test suite. Working out which of many earlier decisions deserves the credit or the blame is the credit assignment problem.
• The data depends on the learner. The agent only sees the outcomes of the actions it takes, so a change in behavior changes the training distribution. A bad policy collects uninformative data, which can keep it bad.

The goal is also different in kind. A supervised model can at best match its teacher. An RL agent is judged only by reward, so it can discover behavior no teacher showed it, which is what AlphaGo's famous move 37 was and what a reasoning model's self-correction is. It can also discover behavior nobody wanted, if the reward is wrong, and we will meet that danger throughout Chapters 16 to 19.

Agents, environments, and the interaction loop

RL splits the world in two. The agent is the learner and decision maker. The environment is everything else. They interact in discrete steps $t = 0, 1, 2, \dots$. At each step the agent observes the state $s_t$, chooses an action $a_t$, and the environment answers with a scalar reward $r_t$ and a new state $s_{t+1}$. The record of one run,

is a trajectory (also called an episode or a rollout). We write the reward for step $t$ as $r_t$, received after taking $a_t$ in $s_t$. This is the convention of Spinning Up, CS285, and the PPO and GAE papers. Sutton and Barto call the same number $R_{t+1}$, so their formulas are ours with every reward index shifted by one.

Here is the environment we will use for most of this chapter, a gridworld from Russell and Norvig's textbook. The agent starts in the bottom-left corner and can try to move up, right, down, or left. Reaching the top-right cell pays $+1$ and ends the episode. Falling into the cell below it pays $-1$ and also ends it. One cell is a wall. Moves are slippery, so the agent goes where it intended with probability 0.8 and slips to one of the two perpendicular directions with probability 0.1 each. Bumping into the wall or the edge leaves it where it was. Every step costs $0.04$, so dawdling is mildly punished.

      col 0   col 1   col 2   col 3
    +-------+-------+-------+-------+
r0  |       |       |       |  +1   |
    +-------+-------+-------+-------+
r1  |       | wall  |       |  -1   |
    +-------+-------+-------+-------+
r2  | start |       |       |       |
    +-------+-------+-------+-------+

Because the world is small, we can write it down completely as two arrays. The transition array $P$ has shape $(|\mathcal{S}|, |\mathcal{A}|, |\mathcal{S}|)$ and holds $P(s' \mid s, a)$, the probability of landing in $s'$ after taking $a$ in $s$. The reward array has shape $(|\mathcal{S}|, |\mathcal{A}|)$ and holds the expected reward $r(s, a)$. The two terminal cells become absorbing states that loop to themselves with zero reward, which lets us treat an episode that has ended as one that goes on forever with nothing happening.

From the start cell, "up" reaches the cell above with probability 0.8, slips right with probability 0.1, and slips left into the edge, staying put, with probability 0.1. Everything an agent could ever learn about this world is in those two arrays. The catch, for most of RL, is that the agent is not given them.

Markov decision processes

The mathematical object behind the loop is a Markov decision process (MDP), a tuple $(\mathcal{S}, \mathcal{A}, P, r, p_0, \gamma)$.

• $\mathcal{S}$ is the set of states and $\mathcal{A}$ the set of actions (11 states and 4 actions in the gridworld).
• $P(s' \mid s, a)$ is the transition distribution.
• $r(s, a)$ is the expected reward for taking $a$ in $s$. (Some texts let the reward depend on $s'$ too. Averaging over $s'$ turns that into our $r(s, a)$, which is all the expected-value formulas need.)
• $p_0(s)$ is the distribution of the first state (always the start cell here).
• $\gamma \in [0, 1]$ is the discount factor, which we meet in a moment.

The word Markov is the important assumption. The next state and reward depend on the history only through the current state and action,

In plain words, the state is a complete summary of the past, as far as the future is concerned. The cell you are in tells you everything about where you can go next. How you got there does not matter.

Whether a problem is Markov depends on what you call the state. A single frame of a video game is not Markov, because it does not show which way the ball is moving, and the classic fix is to stack the last few frames. The most extreme fix is to make the state the entire history so far, which makes any process Markov by construction. That is exactly what a language model does. Its state is the prompt plus every token generated so far, and we will see in §15.7 that this makes text generation a perfectly Markov, and unusually simple, MDP. When the agent truly cannot see the state (a poker player cannot see the other hands), the problem is a partially observable MDP, and the agent has to act on a belief built from its observations. We will not need that here.

Tasks come in two kinds. Episodic tasks end, in a terminal state (the goal, the pit, the end of a chess game, the end-of-sequence token). Continuing tasks go on forever (a thermostat, a recommender). Absorbing states let one set of formulas cover both.

Where does the reward come from? In the gridworld, from the designer. In RLHF, from a reward model trained on human comparisons (§16.3). In reasoning RL, from a program that checks the final answer (§18.3). The working assumption of the field, which Sutton and Barto call the reward hypothesis, is that any goal can be expressed as maximizing the expected sum of a scalar reward. It is a modeling choice, not a theorem, and a careless reward gets maximized in ways nobody intended, which is called specification gaming (§19.1).

Returns and discounting

The agent's goal is not to get a large reward now but a large total reward over time. The return from step $t$ is the discounted sum of the rewards that follow,

and for an episode that ends at step $T$ (the last reward is $r_{T-1}$) the sum stops there. In plain words, a reward $k$ steps in the future is worth $\gamma^k$ of the same reward now.

The definition hides a recursion that the rest of the chapter runs on. Pull out the first term and factor $\gamma$ from the rest,

In plain words, the return now is this step's reward plus the discounted return from the next step. It means all the returns of an episode can be computed in one backward pass, starting from $G_T = 0$.

Worked example. Starting in the cell just above the start, the agent moves up once and then right three times, with no slips, reaching the goal on its fourth step. The first three steps cost $0.04$ each and the last pays $1 - 0.04 = 0.96$. With $\gamma = 0.9$ the rewards are $(-0.04, -0.04, -0.04, 0.96)$, and backwards,

The same walk with $\gamma = 1$ has return $0.84$. Discounting makes the distant $+1$ worth less, and the longer the route, the less it is worth.

Why discount at all? There are four good reasons, and it helps to keep all four in mind.

• Mathematics. If every reward satisfies $|r_t| \le r_{\max}$ and $\gamma < 1$, the return is bounded, $|G_t| \le \sum_k \gamma^k r_{\max} = r_{\max}/(1 - \gamma)$, even for a task that never ends. Every convergence proof in §15.2 uses this.
• An effective horizon. The weights $\gamma^k$ sum to $1/(1-\gamma)$, and they fall to $1/e$ after about $1/(1-\gamma)$ steps (because $\gamma^k = e^{k\ln\gamma} \approx e^{-k(1-\gamma)}$ when $\gamma$ is near 1). So $\gamma = 0.9$ looks about 10 steps ahead and $\gamma = 0.99$ about 100.
• A chance of ending. Suppose the episode could be cut off after each step with probability $1-\gamma$, independently. The reward at step $k$ is collected only if the episode survives $k$ cut-offs, which has probability $\gamma^k$, so the expected undiscounted total is exactly $\sum_k \gamma^k r_{t+k}$. A discounted return is an undiscounted return in a world that might end at any moment.
• Variance. Distant rewards are the noisiest part of a return. Down-weighting them lowers the variance of every estimate we will make, at the price of bias, and in §15.5 we will treat $\gamma$ as a bias-variance knob.

Discounting is not mandatory. In an episodic task where every episode ends in bounded time, $\gamma = 1$ is fine, and it is the usual choice for language models, whose responses are capped at a maximum length (§15.7).

The matrix form is useful when you want the returns for a whole batch at once without a Python loop, though it costs $O(T^2)$ memory, so for long episodes the backward loop (or a library routine such as scipy.signal.lfilter) is the better tool.

A rollout buffer in practice holds many episodes back to back, with a flag marking where each one ended, and the most common bug in RL code is letting the return of one episode leak into the previous one. The next exercise is the version of this you will write in every RL codebase.

import numpy as np

def discounted_returns(rewards, dones, gamma):
    raise NotImplementedError

Policies

The agent's behavior is a policy, a rule for choosing actions. A deterministic policy is a function $a = \pi(s)$. A stochastic policy is a conditional distribution $\pi(a \mid s)$, the probability of choosing $a$ in state $s$, with $\sum_a \pi(a \mid s) = 1$ for every $s$. A deterministic policy is the special case that puts all its probability on one action.

When the policy has learnable parameters $\theta$ we write $\pi_\theta$. For a finite set of actions, the standard choice is a softmax over action scores (logits) produced by a network,

where $\mathbf{z}_\theta(s) \in \R^{|\mathcal{A}|}$ is a vector of logits. That is exactly a language model's next-token distribution, with the context as the state and the vocabulary as the action set. For continuous actions (joint torques of a robot) the usual choice is a Gaussian whose mean is a network output, but we will not need continuous actions in this chapter.

Stochastic policies are not only a modeling convenience. They explore, because they sometimes try actions that currently look worse. In games against an opponent (rock, paper, scissors) any deterministic policy can be exploited. When two different states look the same to the agent, the best policy can be genuinely random, and we will meet a tiny example in §15.3. And for a language model, sampling is how it produces varied text.

Value functions

How good is a state? It depends on what the agent will do from there, so values are always defined for a policy. The state-value function of $\pi$ is the expected return when starting in $s$ and following $\pi$,

and the action-value function is the expected return when starting in $s$, taking action $a$ first, and following $\pi$ afterward,

The subscript on $\mathbb{E}_\pi$ means that every action is drawn from $\pi$ and every transition from $P$. In plain words, $V^\pi(s)$ is how much total reward you expect from here if you play $\pi$, and $Q^\pi(s, a)$ is the same if you first make the move $a$ and then play $\pi$.

The two are linked in both directions. Averaging over the first action gives

and splitting off the first reward (using $G_t = r_t + \gamma G_{t+1}$) gives

We will derive the second one carefully in §15.2, where it becomes the Bellman equation. Their difference has a name we will use constantly, the advantage $A^\pi(s, a) = Q^\pi(s, a) - V^\pi(s)$, which says how much better action $a$ is than the policy's average behavior in $s$.

Worked example, by hand. Take an MDP with two states. In state $X$ you can stay, which pays $1$ and keeps you in $X$, or go, which pays $0$ and moves you to $Y$. State $Y$ has a single action that pays $2$ and keeps you in $Y$ forever. Consider the policy that always stays. Its value is $V(X) = 1 + \gamma + \gamma^2 + \dots = 1/(1-\gamma)$. If you go instead, you collect $0$ now and then $2/(1-\gamma)$ from $Y$ onward, discounted by one step, so $Q(X, \textit{go}) = 2\gamma/(1-\gamma)$. Going is better exactly when $2\gamma > 1$, that is when $\gamma > 0.5$. With $\gamma = 0.9$, $V(X) = 10$ under staying, while going is worth $18$. With $\gamma = 0.4$, staying is worth $1.67$ and going only $1.33$. The discount factor does not only scale values. It can change which behavior is best, because it decides how much the agent cares about a payoff that takes time to reach.

When the environment is unknown, the definition itself suggests an estimator. Run the policy many times from $s$ and average the returns. This is the Monte Carlo method. Here is the uniformly random policy in the gridworld, estimated from 2,000 episodes, next to the exact value, which we compute by solving a linear system that §15.2 will derive.

The Monte Carlo estimate lands within about one standard error of the exact value $-0.41$, but it took 2,000 complete episodes to get two decimal places, because each return is a noisy sum of many random steps. The random walker usually blunders into the pit or wanders around paying step costs. That noise is the central practical problem of RL, and most of this chapter is about taming it.

Optimal policies and optimal values

One policy is at least as good as another if its value is at least as large in every state. The optimal value functions are the best achievable,

It is not obvious that a single policy can be best in every state at once, since the maximum is taken separately for each $s$. This is a theorem. For a finite MDP with $\gamma < 1$ there is a deterministic, stationary (time-independent) policy $\pi^\star$ with $V^{\pi^\star}(s) = V^\star(s)$ for every state simultaneously (Puterman 1994; we prove it in §15.2). Once you know $Q^\star$, such a policy is easy to read off,

and acting this way needs no model of the environment, only the table of $Q^\star$. Knowing $V^\star$ alone is not enough without the model, since turning state values into a choice of action needs $P$ to look one step ahead. This is why model-free methods like Q-learning (§15.2) learn $Q$ rather than $V$.

Brute force is hopeless. The gridworld has 9 non-terminal states and 4 actions, so it has $4^9 = 262{,}144$ deterministic policies, and a language model with a 100,000-token vocabulary has more policies than anything can enumerate. The Bellman equations of §15.2 are what make the search feasible.

Exploration versus exploitation

To find good actions an agent must try actions whose value it does not know. To collect reward it must take the actions it already believes are good. Every RL algorithm strikes some balance between exploration and exploitation, and the cleanest place to study the trade-off is a problem with no states at all.

A $K$-armed bandit (named after the slot machine, the one-armed bandit) is an MDP with a single state and episodes of length one. There are $K$ actions, called arms. Pulling arm $a$ pays a random reward with an unknown mean $\mu_a$, and you get $n$ pulls in total. Let $\mu^\star = \max_a \mu_a$ be the best mean and $\Delta_a = \mu^\star - \mu_a$ the gap of arm $a$. The standard measure of how well an algorithm does is its regret, the reward lost compared with always pulling the best arm,

where $N_a(n)$ is the number of times arm $a$ was pulled in $n$ rounds. The second form follows by grouping the pulls by arm, and it says regret is the expected number of pulls of each bad arm, weighted by how bad that arm is.

Every bandit algorithm keeps an estimate $\hat\mu_a$ of each mean, usually the average reward seen from that arm. Recomputing the average from scratch is wasteful. After the $N$-th reward $r$ from an arm,

which is the running mean rearranged ($\frac{(N-1)\hat\mu + r}{N} = \hat\mu + \frac{r - \hat\mu}{N}$). In plain words, new estimate equals old estimate plus a step size times the error. That shape, estimate plus step times (target minus estimate), is the update behind TD learning, Q-learning, and value-function regression later in this chapter, so it is worth recognizing.

Greedy fails, and epsilon-greedy fixes it

The greedy algorithm always pulls the arm with the highest estimate. It can lock onto a bad arm forever. Take two Bernoulli arms with means $0.4$ and $0.6$, pull each once, then go greedy. If the first pull of the good arm pays $0$ and the bad arm pays $1$, the bad arm's estimate can never fall to $0$ (it has at least one success), the good arm's estimate stays at exactly $0$, and the good arm is never tried again. That happens with probability $0.4 \times 0.4 = 0.16$ (the bad arm pays and the good arm does not), so about one greedy run in six earns $0.4$ per pull forever, and ties at zero, broken toward the bad arm, can make it worse.

The simplest fix is $\varepsilon$-greedy. With probability $\varepsilon$ pull an arm uniformly at random, and otherwise act greedily. Every arm keeps being sampled, so every estimate converges to its true mean by the law of large numbers. The price is that exploration never stops. A fraction $\varepsilon(K-1)/K$ of all pulls goes to arms other than the best one, even after the best arm is known, so the regret grows linearly in $n$. Decaying $\varepsilon$ over time (for example $\varepsilon_t \propto 1/t$) can fix this, but the right schedule depends on the unknown gaps.

import numpy as np

def epsilon_greedy(means, n_steps, eps, seed):
    raise NotImplementedError

Optimism in the face of uncertainty, UCB

$\varepsilon$-greedy explores blindly. It is as likely to re-try an arm it has pulled a thousand times as one it has pulled twice. A smarter rule explores where it is uncertain. The idea is to act as if each arm were as good as it could plausibly be, given the data, and pull the arm whose plausible best is highest. Either that arm really is good, and you profit, or it is not, and pulling it shrinks its uncertainty, so you stop being fooled.

To make "plausibly" precise, suppose rewards lie in $[0, 1]$. Hoeffding's inequality (a theorem) says that the average $\hat\mu_a$ of $N_a$ independent rewards rarely falls far below the true mean,

Now choose $u$ so that this failure probability shrinks as time goes on, say to $t^{-4}$ at round $t$. Solving $e^{-2N_a u^2} = t^{-4}$ gives $u = \sqrt{2\ln t / N_a}$. So with high probability $\mu_a$ is below the upper confidence bound $\hat\mu_a + \sqrt{2\ln t/N_a}$, and the UCB1 algorithm of Auer, Cesa-Bianchi, and Fischer (2002) pulls each arm once and then always pulls

with $c = \sqrt 2$ in the original analysis. In plain words, score each arm by its average plus an exploration bonus that is large for arms pulled rarely and grows slowly with time, and pull the top scorer.

The bonus behaves exactly as it should. An arm that is never pulled has its bonus grow like $\sqrt{\ln t}$, so it is eventually tried again. An arm pulled often has a small bonus, so it is chosen only if its average is genuinely high. This is a theorem (Auer et al. 2002, Theorem 1). For rewards in $[0,1]$, UCB1's expected regret after $n$ pulls is at most

which grows only logarithmically in $n$ (here $\pi$ is the number 3.14159, not a policy). Lai and Robbins (1985) proved that no reasonable algorithm can do better than logarithmic regret, so UCB1 is optimal up to constant factors. Compare that with the linear regret of fixed $\varepsilon$-greedy.

import numpy as np

def ucb1(means, n_steps, c, seed):
    raise NotImplementedError

Racing the strategies

Now compare them on the same problem, a 10-armed Bernoulli bandit whose means are drawn at random, averaged over 200 independent runs of 3,000 pulls each. The code runs all 200 bandits at once as rows of arrays, which is the same trick RL libraries use to run many environments in parallel.

Greedy is the worst by far, because many of its runs lock onto a mediocre arm early and never leave, and its regret grows at more than 0.4 per pull. $\varepsilon = 0.1$ learns quickly, then keeps paying for its random pulls, so its regret keeps climbing in a straight line (about 0.045 per pull from 1,000 to 3,000). $\varepsilon = 0.01$ explores too little at first. UCB1 with the textbook constant $c = \sqrt 2$ is a surprise, worse than $\varepsilon = 0.1$ over this horizon. Its constant comes from a worst-case proof, and with ten arms it spends many pulls confirming what it already knows. Its curve bends, though, as the theory says it must. UCB1 lost about 0.097 per pull between pulls 300 and 1,000 and only 0.046 per pull after that, while $\varepsilon$-greedy loses a steady 0.045 to 0.05 per pull forever. UCB1's per-pull loss keeps shrinking, because its total regret grows only like $\ln n$, while $\varepsilon$-greedy's can never drop below $\varepsilon$ times the average gap over all arms, so with enough pulls UCB1 overtakes it. With a smaller bonus, $c = 0.5$, UCB wins at every horizon, and its regret is nearly flat by the end. The lesson generalizes to everything in this chapter. The theory tells you the shape of the curve. The constants still have to be tuned.

The third agent in the arena, the gradient bandit, does not estimate values at all. It keeps a preference (a logit) for each arm, samples arms from the softmax of those preferences, and after each pull it raises the logit of the pulled arm (and lowers the others in proportion to their probabilities) when the reward beats the average reward so far, and does the opposite when it falls short. That is a policy gradient method in miniature, and §15.3 derives its update from first principles. The comparison of the three agents is a preview of the whole chapter. Value-based methods (greedy, $\varepsilon$-greedy, UCB, and Q-learning in §15.2) learn how good actions are and act on the estimates. Policy-based methods learn the action probabilities directly.

Two more ideas are worth knowing by name. Optimistic initial values (start every $\hat\mu_a$ high, so every arm disappoints and gets tried) are a crude form of UCB's optimism. Thompson sampling (Thompson 1933) keeps a Bayesian posterior over each arm's mean, for Bernoulli arms the Beta posterior of §0.4, draws one sample from each posterior, and pulls the arm with the largest sample. It explores in proportion to the probability that each arm is the best, and it has the same logarithmic regret as UCB (Agrawal and Goyal 2012). In practice it often beats UCB, and it is easy to extend.

From bandits to language models

A contextual bandit adds a state that is drawn fresh each round and is not affected by the action. A news site sees a user (the context), shows an article (the action), and observes a click (the reward). Prompting a language model and scoring its whole response with a reward model has exactly this shape, if we treat the entire response as one action, and InstructGPT's paper describes its RL environment in just those terms (Ouyang et al. 2022). Treating each token as an action instead turns the same problem into a many-step MDP. Both views are correct, and §15.7 shows when each one is the useful one.

In PyTorch

In real code a policy is a module that maps states to logits, and torch.distributions.Categorical handles sampling, log-probabilities, and entropy. Every policy gradient method in this chapter needs exactly those three operations.

Read more

• Sutton & Barto (2018). Reinforcement Learning: An Introduction, 2nd edition. Chapters 2 and 3 cover bandits and finite MDPs. Free online, and the standard reference for everything in this chapter.
• OpenAI Spinning Up. Part 1: Key Concepts in RL. A compact statement of the vocabulary and notation we use, with the same reward indexing.
• Auer, Cesa-Bianchi & Fischer (2002). Finite-time Analysis of the Multiarmed Bandit Problem. The UCB1 algorithm and its logarithmic regret bound.
• Lattimore & Szepesvári (2020). Bandit Algorithms. A rigorous, readable book on bandits, including the Lai-Robbins lower bound and Thompson sampling.
• Russo et al. (2018). A Tutorial on Thompson Sampling. The Bayesian alternative to UCB, explained with examples.
• Levine. CS 285, Deep Reinforcement Learning. Berkeley's course, whose notation for trajectories and policy gradients we follow from §15.3 on.

Watch next

▶

▶

▶