Reinforcement Learning Foundations

34.1 MDPs, returns, and values

A Markov decision process has states $s$ , actions $a$ , transition probabilities, rewards, and a discount $0 \le \gamma \le 1$ . The Markov assumption says the next state and reward depend on the past only through the current $(s,a)$ . A policy $\pi(a\mid s)$ chooses actions. A trajectory’s discounted return from time $t$ is

G_t = \sum_{k=0}^{\infty} \gamma^k r_{t+k}.\tag{34.1}

The discount is not just a mathematical trick. It makes far-future rewards count less, and when $\gamma<1$ it keeps infinite-horizon sums finite. In episodic problems the sum also stops at termination. LLM post-training often has short episodes: prompt in, response out, reward at the end; the same notation still applies, with many intermediate rewards equal to zero.

The value of a policy is the expected return after starting from $s$ : $v_\pi(s)=\E_\pi[G_t\mid s_t=s]$ . Split off the first reward and use the Markov property:

v_\pi(s)=\E_\pi[r_t + \gamma v_\pi(s_{t+1}) \mid s_t=s].\tag{34.2}

For a fixed policy this is a linear system, $\vv=\vr+\gamma\mP\vv$ . The code solves it for a three-state chain where state 0 moves to state 1, state 1 receives reward 1 and terminates, and the terminal state has value 0. With $\gamma=0.9$ , the values are exactly $[0.9,1,0]$ , and the test asserts the Bellman residual.

An action-value $q_\pi(s,a)$ is the same idea after forcing the first action. Policy improvement chooses actions with larger $q$ , but estimating $q$ directly is expensive for a large language model because the action space is the vocabulary or the set of whole responses. Policy gradients avoid enumerating all actions by using samples from the current policy.

Listing 34.1 Returns and Bellman policy evaluation

def discounted_returns(rewards, gamma):
    """Return G_t = r_t + gamma r_{t+1} + ... for one trajectory."""
    rewards = np.asarray(rewards, dtype=np.float64)
    returns = np.zeros_like(rewards)
    running = 0.0
    for t in range(len(rewards) - 1, -1, -1):
        running = rewards[t] + gamma * running
        returns[t] = running
    return returns


def evaluate_policy(transition, reward, gamma):
    """Solve v = r + gamma P v for a fixed policy's transition matrix."""
    transition = np.asarray(transition, dtype=np.float64)
    reward = np.asarray(reward, dtype=np.float64)
    system = np.eye(transition.shape[0]) - gamma * transition
    return np.linalg.solve(system, reward)


def tiny_chain(gamma):
    """Three-state chain: 0 -> 1 -> terminal, reward 1 on state 1."""
    transition = np.array([[0, 1, 0], [0, 0, 1], [0, 0, 1]], dtype=np.float64)
    reward = np.array([0, 1, 0], dtype=np.float64)
    return evaluate_policy(transition, reward, gamma)

34.2 The policy-gradient theorem

Let $J(\vtheta)=\E_{\tau\sim\pi_\vtheta}[G_0]$ be expected return. The score-function identity from Section 6.4.1 gives

\nabla_\vtheta J = \E_\tau\big[G_0\nabla_\vtheta\log p_\vtheta(\tau)\big].\tag{34.3}

The environment dynamics do not depend on $\vtheta$ , so the trajectory log-probability contributes only action log-probabilities:

\nabla_\vtheta J = \E_\pi\Big[\sum_t G_t\nabla_\vtheta\log\pi_\vtheta(a_t\mid s_t)\Big].\tag{34.4}

Replacing $G_0$ by $G_t$ is the causality step: rewards before action $a_t$ do not depend on that action, so their expected score term is zero. In a one-state bandit, the theorem says to raise the logit of sampled actions in proportion to reward. The exact categorical gradient is $\pi(a)(q(a)-\E_\pi[q])$ , which the test checks by finite differences; the sampled REINFORCE loop learns to put more than 94% probability on the best arm.

This is a theorem about an expectation, not about a single rollout. One sampled action can be lucky or unlucky, so the estimator is noisy even when it is unbiased. The update becomes useful by averaging many samples, using a baseline, or both. The bandit example keeps rewards deterministic so the only randomness is action sampling; that isolates the policy-gradient estimator itself.

Listing 34.2 REINFORCE on a categorical bandit

def reinforce_bandit(action_values, steps=200, batch_size=64, lr=0.2, seed=0):
    """Sampled REINFORCE updates for a one-state bandit."""
    rng = np.random.default_rng(seed)
    logits = np.zeros(len(action_values), dtype=np.float64)
    for _ in range(steps):
        probabilities = softmax(logits)
        actions = rng.choice(len(action_values), size=batch_size, p=probabilities)
        rewards = action_values[actions]
        baseline = np.mean(rewards)
        grad = np.zeros_like(logits)
        for action, reward in zip(actions, rewards):
            grad_logp = -probabilities.copy()
            grad_logp[action] += 1.0
            grad += (reward - baseline) * grad_logp
        logits += lr * grad / batch_size
    return logits, softmax(logits)

34.3 Baselines, advantages, and off-policy data

Subtracting a baseline that does not depend on the sampled action leaves the policy gradient unchanged:

\E_\pi[b(s_t)\nabla_\vtheta\log\pi_\vtheta(a_t\mid s_t)] = 0.\tag{34.5}

The equality holds because the expectation is $b(s_t)\nabla_\vtheta\sum_a\pi_\vtheta(a\mid s_t)=b(s_t)\nabla_\vtheta 1$ . A good baseline reduces variance. The usual choice is a value estimate, giving an advantage $A_t=G_t-V(s_t)$ : positive means the sampled action did better than expected from that state, negative means it did worse.

The baseline must not depend on which action was sampled at that state. If it did, it could add a systematic push toward or away from that action. A value function is safe because it predicts the average return before seeing the sampled action. In code, many implementations also normalize a batch of advantages to mean zero and unit scale; that changes optimization dynamics but not the sign of which samples were better than their peers.

Logged data often came from a behavior policy $b$ , not the target policy $\pi$ . Importance sampling rewrites one expectation as another:

\E_{a\sim\pi}[f(a)] = \E_{a\sim b}\Big[\frac{\pi(a)}{b(a)}f(a)\Big].\tag{34.6}

For trajectories, the ratio is a product over time, which can have high variance. The chapter code shows the one-step version used by bandits and by per-token corrections.

The formula also states its own failure mode. If $b(a)=0$ for an action that $\pi$ might take, the ratio is undefined and the logged data cannot tell us what would have happened. If $b(a)$ is merely tiny, a few samples receive huge weights. That is why off-policy RL methods usually add clipping, trust regions, or replay rules instead of relying on raw products of ratios over long generations.

Listing 34.3 Importance sampling for logged actions

def importance_ratios(actions, target_probs, behavior_probs):
    """rho_t = pi(a_t) / b(a_t) for logged bandit actions."""
    actions = np.asarray(actions, dtype=np.int64)
    target_probs = np.asarray(target_probs, dtype=np.float64)
    behavior_probs = np.asarray(behavior_probs, dtype=np.float64)
    return target_probs[actions] / behavior_probs[actions]


def off_policy_value(actions, rewards, target_probs, behavior_probs):
    """Ordinary importance-sampling estimate of a target policy value."""
    ratios = importance_ratios(actions, target_probs, behavior_probs)
    return float(np.mean(ratios * rewards))

34.4 Generalized advantage estimation

A learned value function gives a one-step temporal-difference error

\delta_t = r_t + \gamma V(s_{t+1}) - V(s_t).\tag{34.7}

The $k$ -step advantage bootstraps after $k$ rewards:

A_t^{(k)} = \sum_{l=0}^{k-1}\gamma^l r_{t+l} + \gamma^k V(s_{t+k}) - V(s_t).\tag{34.8}

Expanding the TD errors shows a telescoping identity: $A_t^{(k)}=\sum_{l=0}^{k-1}\gamma^l\delta_{t+l}$ . Generalized advantage estimation mixes all $k$ -step advantages with geometric weights [schulman2015highdimensional]:

\hat{A}^{\mathrm{GAE}}_t = \sum_{l=0}^{\infty} (\gamma\lambda)^l \delta_{t+l}.\tag{34.9}

Thus $\lambda=0$ is one-step TD and $\lambda=1$ approaches the Monte Carlo advantage. The backward recursion follows by separating the first term from the sum:

\hat{A}_t = \delta_t + \gamma\lambda\hat{A}_{t+1}.\tag{34.10}

The tests compare this recursion against the explicit weighted sum for several $\gamma$ and $\lambda$ values.

$\lambda$ is a bias-variance knob. Small $\lambda$ trusts the value function and uses short, low-variance estimates; large $\lambda$ trusts sampled returns and uses longer, higher-variance estimates. If the value function is poor, a very small $\lambda$ can be biased; if rewards are noisy, a very large $\lambda$ can make updates unstable. The recursive implementation is the one used in practice because it is linear in trajectory length and works naturally from the end of a rollout buffer backward.

Listing 34.4 Generalized advantage estimation

def gae_recursive(rewards, values, gamma, lam):
    """Generalized advantage estimates by the backward recursion."""
    deltas = td_errors(rewards, values, gamma)
    advantages = np.zeros_like(deltas)
    running = 0.0
    for t in range(len(deltas) - 1, -1, -1):
        running = deltas[t] + gamma * lam * running
        advantages[t] = running
    return advantages

In practice

Modern LLM post-training still uses these ingredients: sample from the current policy, score the sample, subtract a baseline or normalize advantages, and push up the log-probabilities of above-average samples [lambert2025reinforcement]. Value functions and GAE are common when rewards arrive over a sequence, while simpler response-level methods often use one final reward as the return. Importance sampling is mathematically exact but can explode on long trajectories, so practical algorithms usually clip or otherwise constrain policy changes. The next chapter adds those constraints through PPO.

Key equations

G_t=\sum_{k=0}^{\infty}\gamma^k r_{t+k}, \qquad v_\pi(s)=\E_\pi[r_t+\gamma v_\pi(s_{t+1})\mid s_t=s]

\nabla_\vtheta J = \E_\pi\Big[\sum_t G_t\nabla_\vtheta\log\pi_\vtheta(a_t\mid s_t)\Big]

\E_\pi[(G_t-b(s_t))\nabla\log\pi(a_t\mid s_t)] = \E_\pi[G_t\nabla\log\pi(a_t\mid s_t)]

\E_\pi[f(a)] = \E_b\left[\frac{\pi(a)}{b(a)}f(a)\right]

\hat{A}^{\mathrm{GAE}}_t=\sum_{l=0}^{\infty}(\gamma\lambda)^l\delta_{t+l}, \qquad \hat{A}_t=\delta_t+\gamma\lambda\hat{A}_{t+1}

34.5 Teach it

The one-sentence version. RL raises the probability of sampled actions that beat expectation and lowers actions that disappoint.

An analogy. A coach cannot show the perfect move for every board position, but can say whether the game went better than expected after a move.

At the board.

Draw a three-state chain and write $v=r+\gamma Pv$ .
Write $\nabla \E[G]=\E[G\nabla\log p]$ and cross out environment terms.
Subtract a baseline and show $\E[\nabla\log\pi]=0$ .
Write three TD errors and show how GAE discounts them by $\gamma\lambda$ .

Misconceptions to address.

"Reward is a label." It is feedback on a sampled action, not the target action itself.
"A baseline changes the optimum." It changes variance, not the expected gradient.
"Off-policy data is free." Importance ratios can make it very noisy.

Check for understanding. In a bandit with reward 1 for action A and 0 for action B, what sign should the policy-gradient update give to the logit of A when A is sampled?

34.6 Exercises

Exercise 34.1 ★ Bellman arithmetic

For the chain 0 → 1 → terminal, with reward 1 in state 1 and $\gamma=0.8$ , compute $v(0)$ , $v(1)$ , and $v(terminal)$ .

Solution in Appendix D

Exercise 34.2 ★★ Score-function policy gradient

Starting from $J(\vtheta)=\E_{\tau\sim\pi_\vtheta}[G_0]$ , derive (34.4) and explain the causality step from $G_0$ to $G_t$ .

Solution in Appendix D

Exercise 34.3 ★★ Baselines and importance sampling

Prove (34.5). Then, for behavior probabilities $(0.8,0.2)$ , target probabilities $(0.25,0.75)$ , and rewards $(1,3)$ , compute the exact one-step importance-sampling value.

Solution in Appendix D

Exercise 34.4 ★★★ GAE recursion

Show that the $k$ -step advantage equals $\sum_{l=0}^{k-1}\gamma^l\delta_{t+l}$ . Then derive (34.10) from (34.9) and implement a test against the explicit sum.

Solution in Appendix D

References

[williams1992] R. J. Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine Learning 8, 229–256, 1992.
[sutton2018] R. S. Sutton and A. G. Barto. Reinforcement Learning: An Introduction, 2nd edition. MIT Press, 2018. http://incompleteideas.net/book/the-book-2nd.html
[lambert2025reinforcement] N. Lambert. Reinforcement Learning from Human Feedback. 2025. arXiv:2504.12501
[schulman2015highdimensional] J. Schulman et al. High-Dimensional Continuous Control Using Generalized Advantage Estimation. 2015. arXiv:1506.02438