🕹 Reinforcement Learning
Reinforcement Learning
Every other learning page in this section shows the model the right answer. Supervised learning has labels; [[cross-validation]] measures how well they were fitted. Reinforcement learning has none of that, and three things change:
| supervised | reinforcement | |
|---|---|---|
| signal | the correct output | a scalar reward, possibly much later |
| data | fixed, given | generated by your own behaviour |
| objective | match the label | maximise long-run return |
The second row is the one that makes RL genuinely different. Your policy decides which states you visit, so it decides which data you get to learn from. Change the policy and the data distribution changes underneath you — there is no fixed training set to cross-validate against.
This page builds the machinery from the ground up on a gridworld small enough to print, and ends at the policy-gradient estimator that [[rlhf-grpo-constitutional-ai]] uses. Every value is computed by hand and checked against an exact solution.
Display helpers ready: banner, step, check, expand, show_matrix, show_matmul, show_elementwise.
==============================================================================
A GRIDWORLD SMALL ENOUGH TO PRINT
==============================================================================
The board
==============================================================================
reward on entering each cell (0 = plain, costs -0.04 to step)
+--------+--------+--------+--------+
| +0.000 | +0.000 | +0.000 | +1 END|
+--------+--------+--------+--------+
| +0.000 | #### | +0.000 | -1 END|
+--------+--------+--------+--------+
| +0.000 | +0.000 | +0.000 | +0.000 |
+--------+--------+--------+--------+
states: 11 (one wall removed)
terminals: (0, 3) = +1, (1, 3) = -1
step reward -0.04, discount gamma = 0.9
actions move as intended with p = 0.8, and veer to each
side with p = 0.1. Bumping a wall leaves you in place.
One transition, written out
==============================================================================
from (2, 0), action 'up':
p = 0.8 -> (1, 0)
p = 0.1 -> (2, 0)
p = 0.1 -> (2, 1)
probabilities sum to 1.0
1. Return, and why rewards are discounted
The agent does not maximise the next reward, it maximises the return — the whole discounted future:
$$G_t = R_{t+1} + \gamma R_{t+2} + \gamma^2 R_{t+3} + \cdots = \sum_{k=0}^{\infty} \gamma^k R_{t+k+1}$$
The discount $\gamma \in [0,1)$ does three jobs at once, and it is worth being clear which is which:
- It keeps the sum finite. On a task that never ends, undiscounted return diverges. With $\gamma < 1$ and bounded reward $|R| \le R_{\max}$, the return is bounded by $R_{\max}/(1-\gamma)$.
- It expresses preference for sooner. A reward $k$ steps away is worth $\gamma^k$ of the same reward now.
- It sets an effective horizon. Roughly $1/(1-\gamma)$ steps — beyond that, contributions are negligible. This is the number to think in.
$\gamma$ is not a minor knob. It changes which policy is optimal.
==============================================================================
DISCOUNTING, BY HAND
==============================================================================
A single trajectory
==============================================================================
rewards: [-0.04, -0.04, -0.04, -0.04, 1.0] (four steps, then reach the goal)
gamma=1.0: -0.04 + 1.0^1x-0.04 + 1.0^2x-0.04 + 1.0^3x-0.04 + 1.0^4x+1.00
= +0.840000
hand check +0.840000 MATCH: True
gamma=0.9: -0.04 + 0.9^1x-0.04 + 0.9^2x-0.04 + 0.9^3x-0.04 + 0.9^4x+1.00
= +0.518540
hand check +0.518540 MATCH: True
gamma=0.5: -0.04 + 0.5^1x-0.04 + 0.5^2x-0.04 + 0.5^3x-0.04 + 0.5^4x+1.00
= -0.012500
hand check -0.012500 MATCH: True
Effective horizon and the bound
==============================================================================
gamma 1/(1-gamma) gamma^10 gamma^50 max |return|
------------------------------------------------------------
0.000 1.0 0.0000 0.000000 1.0
0.500 2.0 0.0010 0.000000 2.0
0.900 10.0 0.3487 0.005154 10.0
0.990 100.0 0.9044 0.605006 100.0
0.999 1000.0 0.9900 0.951206 1000.0
------------------------------------------------------------
At gamma=0.9 anything more than ~10 steps away barely registers.
At gamma=0.999 the agent plans over ~1000 steps -- and every value
estimate becomes correspondingly harder to learn.
Gamma changes which policy is optimal
==============================================================================
Two routes to a reward:
A: 2 steps, then +1 -> G = -0.04 + g(-0.04) + g^2(1)
B: 6 steps, then +2 -> G = -0.04(1+g+..+g^5) + g^6(2)
gamma route A route B prefers
----------------------------------------
0.50 0.1900 -0.0475 A
0.80 0.5680 0.3767 A
0.90 0.7340 0.8755 B
0.95 0.8245 1.2583 B
0.99 0.9005 1.6489 B
----------------------------------------
Impatient agents take the near reward; patient ones wait for the
bigger one. Gamma is a statement about what you want, not a
convergence trick.
2. Value functions and the Bellman equation
A policy $\pi(a \mid s)$ says what to do. Its state-value function is the expected return from following it:
$$V^\pi(s) = \mathbb{E}_\pi\big[\,G_t \mid S_t = s\,\big]$$
The Bellman equation is just that definition, unrolled one step:
$$V^\pi(s) = \sum_a \pi(a|s) \sum_{s'} P(s'|s,a)\Big[\,R(s,a,s') + \gamma V^\pi(s')\,\Big]$$
The value of where I am equals the immediate reward plus the discounted value of where I land. That self-referential structure is the whole of dynamic programming — one equation per state, all coupled.
The action-value $Q^\pi(s,a)$ is the same thing after committing to one action first:
$$Q^\pi(s,a) = \sum_{s'} P(s'|s,a)\big[R + \gamma V^\pi(s')\big], \qquad V^\pi(s) = \sum_a \pi(a|s)\, Q^\pi(s,a)$$
$Q$ is the more useful object in practice: given $Q$, acting well needs no model at all — take $\arg\max_a Q(s,a)$. Given only $V$, you would need to know where each action leads.
==============================================================================
POLICY EVALUATION: SOLVING BELLMAN BY ITERATION
==============================================================================
Iterate the backup until nothing changes
==============================================================================
sweep max change V(2,0) V(0,0)
----------------------------------------------
1 8.08e-01 -0.0400 -0.0400
2 5.63e-01 -0.0760 -0.0760
3 3.99e-01 -0.1084 +0.3229
4 1.05e-01 -0.2092 +0.4281
5 7.14e-02 -0.2806 +0.4545
10 2.44e-02 -0.4712 +0.4580
20 3.19e-03 -0.5659 +0.4507
40 5.41e-05 -0.5795 +0.4489
100 1.86e-10 -0.5797 +0.4489
----------------------------------------------
converged after 147 sweeps
V for the 'always go right' policy
+--------+--------+--------+--------+
| +0.449 | +0.646 | +0.791 | +1 END|
+--------+--------+--------+--------+
| -0.185 | #### | -0.799 | -1 END|
+--------+--------+--------+--------+
| -0.580 | -0.654 | -0.689 | -0.716 |
+--------+--------+--------+--------+
Check against the exact linear solution
==============================================================================
Bellman for a FIXED policy is linear: V = r + gamma P V
so (I - gamma P) V = r can be solved directly.
state iterative exact solve diff
----------------------------------------------
(0, 0) 0.448911 0.448911 5.38e-15
(0, 1) 0.646042 0.646042 5.55e-16
(0, 2) 0.791326 0.791326 4.44e-16
(1, 0) -0.184899 -0.184899 3.96e-14
(1, 2) -0.798818 -0.798818 3.00e-15
(2, 0) -0.579708 -0.579708 3.42e-14
----------------------------------------------
max|diff| over all 9 states = 3.96e-14 MATCH: True
3. Finding the best policy: value iteration
Policy evaluation scores a policy you already have. To find the best one, replace the average over actions with a maximum:
$$V^*(s) = \max_a \sum_{s'} P(s'|s,a)\big[\,R + \gamma V^*(s')\,\big]$$
This is the Bellman optimality equation, and iterating it is value iteration. It converges because the backup is a $\gamma$-contraction: each sweep shrinks the worst-case error by at least a factor of $\gamma$, so
$$\|V_{k+1} - V^*\|_\infty \le \gamma \|V_k - V^*\|_\infty$$
That guarantee is unusually strong for an iterative method — convergence is geometric from any starting point, and the rate is a parameter you chose.
==============================================================================
VALUE ITERATION
==============================================================================
Iterate the max-backup
==============================================================================
converged after 37 sweeps
sweep max change ratio to previous
----------------------------------------
1 7.920e-01 nan
2 5.630e-01 0.7109
3 3.989e-01 0.7085
6 1.371e-01 0.6205
11 3.670e-03 0.4394
21 5.518e-07 0.4076
----------------------------------------
observed contraction ratio 0.4250, gamma = 0.9
the theorem gives an UPPER bound, ratio <= gamma: True
Converging faster than the bound is normal -- gamma is the
worst case over all MDPs, and this one is easier than the worst.
V* (optimal values)
+--------+--------+--------+--------+
| +0.610 | +0.766 | +0.928 | +1 END|
+--------+--------+--------+--------+
| +0.487 | #### | +0.585 | -1 END|
+--------+--------+--------+--------+
| +0.374 | +0.327 | +0.428 | +0.189 |
+--------+--------+--------+--------+
Read the optimal policy off the values
==============================================================================
optimal action in each cell
+--------+--------+--------+--------+
| RIGHT | RIGHT | RIGHT | +1 |
+--------+--------+--------+--------+
| UP | #### | UP | -1 |
+--------+--------+--------+--------+
| UP | RIGHT | UP | LEFT |
+--------+--------+--------+--------+
One state's Q values, in full
==============================================================================
state (1, 2) (directly left of the -1 pit):
Q((1, 2), up ) = +0.5849 [0.8->(0, 2), 0.1->(1, 2), 0.1->(1, 3)] <- best
Q((1, 2), down ) = +0.2245 [0.8->(2, 2), 0.1->(1, 2), 0.1->(1, 3)]
Q((1, 2), left ) = +0.5032 [0.8->(1, 2), 0.1->(0, 2), 0.1->(2, 2)]
Q((1, 2), right) = -0.6860 [0.8->(1, 3), 0.1->(0, 2), 0.1->(2, 2)]
V*((1, 2)) = max_a Q = +0.584934
from the table +0.584934 MATCH: True
Note the agent avoids 'right' here even though it is the direct
route past the pit -- the 0.1 slip chance makes it too dangerous.
4. Learning without a model: Q-learning
Everything so far needed $P(s'|s,a)$ — a full map of the world. Real agents do not have one. Temporal-difference learning replaces the expectation with a single sampled transition:
$$Q(s,a) \leftarrow Q(s,a) + \alpha\Big[\underbrace{r + \gamma \max_{a'} Q(s',a')}_{\text{TD target}} - Q(s,a)\Big]$$
The bracket is the TD error — the gap between what you predicted and what one step of experience suggests. You nudge toward it by $\alpha$.
Two things make this remarkable:
- No model. It only needs $(s, a, r, s')$ tuples, which you get by acting.
- Bootstrapping. The target contains $Q$ itself — you learn a guess from a guess. That should not obviously work, and the fact that it converges is the central theoretical result of the field.
Q-learning is off-policy: the $\max_{a'}$ evaluates the greedy policy no matter how the action was actually chosen. So you can explore randomly and still learn the optimal policy. The next section shows what that costs.
==============================================================================
Q-LEARNING FROM EXPERIENCE ALONE
==============================================================================
The first five updates, written out
==============================================================================
step 1: in (2, 0), took 'up' -> landed (1, 0), reward -0.04
target = -0.04 + 0.9 x +0.0000 = -0.0400
TD error = -0.0400 - +0.0000 = -0.0400
Q((2, 0),up) <- +0.0000 + 0.2 x -0.0400 = -0.0080
step 2: in (1, 0), took 'up' -> landed (0, 0), reward -0.04
target = -0.04 + 0.9 x +0.0000 = -0.0400
TD error = -0.0400 - +0.0000 = -0.0400
Q((1, 0),up) <- +0.0000 + 0.2 x -0.0400 = -0.0080
step 3: in (0, 0), took 'up' -> landed (0, 0), reward -0.04
target = -0.04 + 0.9 x +0.0000 = -0.0400
TD error = -0.0400 - +0.0000 = -0.0400
Q((0, 0),up) <- +0.0000 + 0.2 x -0.0400 = -0.0080
step 4: in (0, 0), took 'down' -> landed (0, 0), reward -0.04
target = -0.04 + 0.9 x +0.0000 = -0.0400
TD error = -0.0400 - +0.0000 = -0.0400
Q((0, 0),down) <- +0.0000 + 0.2 x -0.0400 = -0.0080
step 5: in (0, 0), took 'left' -> landed (0, 0), reward -0.04
target = -0.04 + 0.9 x +0.0000 = -0.0400
TD error = -0.0400 - +0.0000 = -0.0400
Q((0, 0),left) <- +0.0000 + 0.2 x -0.0400 = -0.0080
Coverage and step size decide whether it converges at all
==============================================================================
Q-learning converges to Q* only if every state-action pair is
visited infinitely often, and the step size satisfies the
Robbins-Monro conditions: sum(alpha) = infinity, sum(alpha^2) < infinity.
Both are easy to violate by accident. Three runs, same algorithm:
setup max|V-V*| optimal min visits
--------------------------------------------------------------------------
always start at (2,0), alpha = 0.1 0.3498 6/9 2,793
exploring starts, alpha = 0.1 0.1223 7/9 8,308
exploring starts, alpha decaying 0.0260 9/9 6,542
--------------------------------------------------------------------------
Row 1: starting always at (2,0) leaves far corners barely visited,
so their Q values never settle.
Row 2: exploring starts fix the coverage.
Row 3: a constant alpha keeps injecting noise forever -- it converges
to a BALL around Q*, not to Q*. Decaying alpha closes it.
max_a Q(s,a) after 20,000 episodes, no model used
+--------+--------+--------+--------+
| +0.613 | +0.768 | +0.927 | +1 END|
+--------+--------+--------+--------+
| +0.491 | #### | +0.568 | -1 END|
+--------+--------+--------+--------+
| +0.378 | +0.329 | +0.421 | +0.215 |
+--------+--------+--------+--------+
Compare with the V* computed from the model in section 3
==============================================================================
state V* (model) Q-learning diff
----------------------------------------------
(0, 0) 0.6105 0.6133 0.0028
(0, 1) 0.7662 0.7685 0.0023
(0, 2) 0.9282 0.9274 0.0007
(1, 0) 0.4872 0.4907 0.0035
(1, 2) 0.5849 0.5682 0.0168
(2, 0) 0.3739 0.3775 0.0037
(2, 1) 0.3266 0.3288 0.0022
(2, 2) 0.4275 0.4213 0.0063
(2, 3) 0.1888 0.2148 0.0260
----------------------------------------------
max|diff| = 0.0260
optimal action recovered in 9/9 states
every optimal action recovered: True
The agent never saw the transition probabilities. It reconstructed
the value of every state from sampled experience alone.
5. On-policy vs off-policy: SARSA and the cliff
Change one symbol in the update and the behaviour changes completely:
$$\text{Q-learning (off-policy)}: \quad r + \gamma \max_{a'} Q(s',a')$$ $$\text{SARSA (on-policy)}: \quad\quad\ \ r + \gamma\, Q(s', a')\ \text{ where } a' \text{ is the action actually taken next}$$
Q-learning asks "what if I acted greedily from here?" SARSA asks "what if I kept behaving the way I actually behave — exploration and all?"
On the classic cliff walk this produces two different policies. The shortest path runs along the cliff edge; one exploratory step off it is a large penalty.
- Q-learning learns the optimal path — right along the edge — because its target ignores the fact that it explores.
- SARSA learns a safer path further from the edge, because its target includes the cost of its own occasional random step.
The subtlety worth keeping: Q-learning learns the better policy while earning the worse return during training. If your agent is learning in the real world, where falling off cliffs is expensive, that distinction matters.
==============================================================================
CLIFF WALKING: SARSA vs Q-LEARNING
==============================================================================
Average return per episode during training
==============================================================================
method last 500 episodes
----------------------------------
Q-learning -51.05
SARSA -24.14
----------------------------------
SARSA earns MORE while learning, despite learning a longer path.
Q-learning keeps walking the edge and keeps falling off.
The routes they learned (greedy, no exploration)
==============================================================================
Q-learning: path length 13, rows visited [2, 3]
SARSA: path length 17, rows visited [0, 1, 2, 3]
grid: S start, G goal, C cliff, q = Q-learning route, s = SARSA route
. s s s s s s s s s s s
s s . . . . . . . . . s
* q q q q q q q q q q *
S C C C C C C C C C C G
Q-learning hugs the cliff (row 2, the optimal path).
SARSA keeps a row of clearance, because its target accounts for
the 10% chance that its next move is random.
6. Policy gradients — the bridge to RLHF
Everything above learns values and derives a policy from them. For a language model that is impractical: the action space is the whole vocabulary at every token. Policy-gradient methods skip values and optimise the policy directly.
The objective is expected return $J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta}[R(\tau)]$. The difficulty is that $\theta$ appears in the distribution being averaged over, not in the thing being averaged. The score-function identity resolves it:
$$\nabla_\theta J = \mathbb{E}_{\tau \sim \pi_\theta}\big[\,R(\tau)\,\nabla_\theta \log \pi_\theta(\tau)\,\big]$$
which is an expectation you can estimate by sampling — that is REINFORCE. In words: increase the log-probability of whatever produced high reward.
The estimator is unbiased but extremely noisy, and the standard fix is a baseline:
$$\nabla_\theta J = \mathbb{E}\big[\,(R(\tau) - b)\,\nabla_\theta \log \pi_\theta(\tau)\,\big]$$
Any $b$ not depending on the action leaves the estimator unbiased, because $\mathbb{E}[\nabla_\theta \log \pi_\theta] = 0$ — but it can cut the variance enormously. This is exactly what GRPO's group mean is: sample $G$ completions, subtract their average reward, and you have a baseline that needed no value network. Both claims — unbiasedness and variance reduction — are verified below.
==============================================================================
REINFORCE, AND WHY BASELINES MATTER
==============================================================================
A two-armed problem with a softmax policy
==============================================================================
rewards: arm0 = 1.0, arm1 = 3.0
policy at theta=0.0: [0.5 0.5]
J(theta) = E[R] = 0.5000x1.0 + 0.5000x3.0 = 2.000000
The exact gradient, analytically
==============================================================================
finite difference dJ/dtheta = 0.50000000
score-function form E[R dlog pi] = 0.50000000
MATCH: True
Monte Carlo estimates: is the baseline still unbiased?
==============================================================================
32 samples per estimate, 4000 independent estimates
baseline b mean estimate bias std dev var ratio
------------------------------------------------------------------
0 (none) 0.498922 -1.08e-03 0.175545 1.0000
1.0 0.498016 -1.98e-03 0.087756 0.2499
E[R] (mean) 0.500000 2.63e-10 0.000000 0.0000
3.0 0.501336 1.34e-03 0.088167 0.2522
------------------------------------------------------------------
exact gradient = 0.500000
Every baseline gives the same mean -- all unbiased, as the identity
E[dlog pi] = 0 promises. Check that directly:
sum_a pi(a) dlog pi(a) = 0.00e+00 is zero: True
But the VARIANCE differs a lot -- a 4x reduction from any sensible
baseline, and the mean reward does better still.
Why the mean baseline gives EXACTLY zero variance here
==============================================================================
That row is not a bug, and it is not the general case. With
b = E[R] = 2.0, work out the per-sample estimate for each arm:
arm 0: (R - b) x dlog pi = (1.0 - 2.0) x -0.5 = +0.5000
arm 1: (R - b) x dlog pi = (3.0 - 2.0) x +0.5 = +0.5000
Both arms give the same number, so every single sample returns the
exact gradient and the estimator has no variance at all. That is an
algebraic coincidence of this symmetric two-arm setup (p = 0.5).
In general the optimal baseline REDUCES variance, it does not
eliminate it -- and it is not quite E[R] either, but a
gradient-magnitude-weighted average of returns. E[R] is the cheap
approximation everyone actually uses.
And that approximation is precisely GRPO's group baseline: sample a
group of completions, subtract the group mean reward, no value
network needed anywhere.
Where this lands in the LLM setting
==============================================================================
concept here -> in [[rlhf-grpo-constitutional-ai]]
--------------------------------------------------------------------
trajectory tau -> a sampled completion
reward R(tau) -> verifier: is the answer right
baseline b -> mean reward of the sampled group
(R - b) -> the advantage A_i
dlog pi(tau) -> gradient of token log-probs
step size control -> the clipped ratio + KL penalty
--------------------------------------------------------------------
GRPO is REINFORCE with a group-mean baseline and a trust region.
Everything on this page is underneath it.
7. What to take away
RL is different because the policy chooses the data. There is no fixed training set. Improving the policy changes the distribution you learn from, which is why nothing from [[cross-validation]] transfers directly.
The discount is a modelling choice, not a technicality. It sets an effective horizon of about $1/(1-\gamma)$ steps and can change which policy is optimal.
Bellman is one idea stated twice. The value of a state is the immediate reward plus the discounted value of the next. With an average over actions it evaluates a policy; with a maximum it finds the best one. Both converge geometrically at rate $\gamma$.
TD learning removes the need for a model by replacing an expectation with one sample, and bootstraps a guess from a guess. That it converges at all is the field's central result.
Off-policy learns the better policy; on-policy earns the better return while learning. The cliff makes this concrete: Q-learning finds the optimal edge path and keeps falling off it during training; SARSA keeps clearance and scores better throughout.
Policy gradients are what scale to language models, and REINFORCE with a baseline is the whole idea. Any action-independent baseline keeps the estimator unbiased while cutting variance — and GRPO's group mean is exactly that baseline.
Verified in this notebook
- Discounted returns computed term by term and matched against the closed form
- Iterative policy evaluation matched to the exact linear solve $(I - \gamma P)V = r$ to $10^{-6}$ across all states
- Value iteration's observed contraction ratio measured and confirmed to satisfy the $\le \gamma$ bound (it converges faster, as it should — $\gamma$ is the worst case over all MDPs)
- $V^*(s) = \max_a Q^*(s,a)$ checked at a state, with all four action values shown
- Q-learning from sampled experience alone reconstructing $V^*$ to within 0.026 and recovering the optimal action in all nine states, with no access to the transition model
- The two conditions that decide whether it converges at all, each shown breaking it: poor state coverage from a fixed start (6/9 actions correct) and a constant step size that leaves a permanent noise floor (7/9). Exploring starts plus a decaying $lpha$ gives 9/9
- SARSA and Q-learning learning genuinely different routes on the cliff, with SARSA's higher training return quantified
- The score-function gradient matched to a finite-difference gradient
- Four baselines shown to be unbiased ($\mathbb{E}[\nabla \log \pi] = 0$ verified directly) with substantially different variances
Where this connects
[[rlhf-grpo-constitutional-ai]] is this machinery applied to language models — GRPO is REINFORCE with a group-mean baseline and a trust region. [[self-improving-llms]] is what happens when the reward comes from the model itself. [[optimisers-sgd-to-adam]] is the update rule underneath the policy gradient.
References
- Reinforcement Learning: An Introduction (2nd ed.)
- Learning from Delayed Rewards (Q-learning)
- Simple Statistical Gradient-Following Algorithms for Connectionist Reinforcement Learning (REINFORCE)
- Policy Gradient Methods for Reinforcement Learning with Function Approximation