02 · Policy Gradients

What you're learning

How to differentiate an expectation you can only sample from, and how that single trick — the score function estimator — turns “maximize expected return” into a gradient you can actually compute. By the end you should be able to derive the Policy Gradient Theorem from scratch, explain exactly why a baseline is unbiased, and see why every modern RLHF/RLVR method is a REINFORCE estimator in disguise.

This lesson assumes lesson 01 — 01-mdps-and-the-rl-objective — is understood: MDPs, the trajectory distribution, the RL objective $J(\theta )$, and the value/advantage functions $V^{\pi },Q^{\pi },A^{\pi }$.

---

1. Learning map

graph TD
    A["RL objective<br/>J(θ)=E_τ[R(τ)]<br/>(lesson 01)"] --> B["Why not value-based?<br/>argmax over actions"]
    B --> C["Directly optimize π_θ<br/>→ need ∇_θ J"]
    C --> D["Score-function trick<br/>∇E[f]=E[f ∇log p]"]
    D --> E["Policy Gradient Theorem<br/>Σ_t ∇log π · Ψ_t"]
    E --> F["REINFORCE<br/>(MC policy gradient)"]
    F --> G["Variance problem"]
    G --> H1["Reward-to-go<br/>(causality)"]
    G --> H2["Baselines b(s)<br/>(unbiased)"]
    H2 --> I["V(s) as baseline<br/>→ Advantage Ψ_t=A(s,a)"]
    H1 --> I
    I --> J["LLM connection<br/>RLHF / RLVR / GRPO"]

    style E fill:#4a4,color:#fff
    style D fill:#44a,color:#fff
    style J fill:#a44,color:#fff

The spine: we want ${\nabla }_{\theta }J$, we can’t differentiate through sampling naively, the score-function trick rescues us, the Policy Gradient Theorem specializes it to MDPs, REINFORCE turns it into an algorithm, and then the entire back half is fighting variance.

---

2. Why policy gradients at all?

In lesson 01 the objective was

$$J(\theta )={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}[R(\tau )],R(\tau )=\sum _{t=0}^T{\gamma }^t{r}_t,$$

where the trajectory distribution factorizes as

$$p_{\theta }(\tau )={\rho }_0(s_0)\prod _{t=0}^T{\pi }_{\theta }(a_t\mid s_t)P(s_{t+1}\mid s_t,a_t).$$

Value-based methods (Q-learning¹, DQN²) never parameterize a policy directly. They learn $Q(s,a)$ and induce a policy by $\pi (s)=\arg {\max }_{a}Q(s,a)$. That $\arg \max$ is the problem.

Where value-based methods hit a wall

The greedy $\arg {\max }_{a}Q(s,a)$ requires enumerating actions. This is fine for Atari (18 buttons) but breaks for:

• Continuous actions (robot torques): the $\arg \max$ is itself an optimization problem at every step.
• Combinatorially large discrete spaces — e.g. a language model choosing the next token over a vocabulary of $|\mathcal{V}|\approx 128\text{k}$³, or worse, choosing a whole sequence. You cannot enumerate sequences.

Policy gradients sidestep the $\arg \max$ entirely: parameterize ${\pi }_{\theta }$ (e.g. a softmax over logits — exactly what an LLM head is) and do gradient ascent on $J$ directly.

Key insight

Policy-based methods optimize the thing you actually care about — the policy — rather than a proxy (values) from which a policy must be extracted. They naturally produce stochastic policies, handle continuous and huge discrete action spaces, and give smooth policy improvement. The price you pay is variance, which is what the rest of this lesson is about.

The catch: $J(\theta )$ is an expectation over trajectories that themselves depend on $\theta$. You can’t just push ${\nabla }_{\theta }$ inside a sampling operation. That’s the wall we break next.

---

3. The score-function / log-derivative trick

This is the whole game. Suppose we want the gradient of an expectation of some function $f(x)$ under a distribution $p_{\theta }(x)$ that depends on the parameters:

$${\nabla }_{\theta }{\mathbb{E}}_{x\sim p_{\theta }}[f(x)]={\nabla }_{\theta }\int p_{\theta }(x)f(x)dx=\int {\nabla }_{\theta }{p}_{\theta }(x)f(x)dx.$$

(We assume $f$ does not depend on $\theta$; regularity conditions let us swap $\nabla$ and $\int$.) The integrand ${\nabla }_{\theta }{p}_{\theta }$ is not an expectation — we can’t sample it. The trick is the identity

$${\nabla }_{\theta }{p}_{\theta }(x)=p_{\theta }(x)\frac{{\nabla }_{\theta }{p}_{\theta }(x)}{p_{\theta }(x)}=p_{\theta }(x){\nabla }_{\theta }\log p_{\theta }(x),$$

which just rearranges $\nabla \log p=\nabla p/p$. Substituting:

$$\boxed{{\nabla }_{\theta }{\mathbb{E}}_{x\sim p_{\theta }}[f(x)]={\mathbb{E}}_{x\sim p_{\theta }}[f(x){\nabla }_{\theta }\log p_{\theta }(x)]}$$

We turned a gradient-of-an-expectation into an expectation of a gradient — which we can estimate by Monte Carlo: sample $x^{(i)}\sim p_{\theta }$, average $f(x^{(i)}){\nabla }_{\theta }\log p_{\theta }(x^{(i)})$.

The term ${\nabla }_{\theta }\log p_{\theta }(x)$ is called the score function. This estimator is also called the likelihood-ratio estimator.⁴

Why this is remarkable

We can compute the gradient of $\mathbb{E}[f]$ without differentiating $f$ at all — $f$ can be non-differentiable, discontinuous, even a black-box reward from a human or a unit-test harness. All the gradient flows through $\log p_{\theta }$. This is exactly why policy gradients work with reward signals that have no gradient (a reward model, a pass/fail verifier).

Contrast with the reparameterization trick

The other way to differentiate an expectation is the pathwise / reparameterization estimator ($x=g_{\theta }(ϵ)$, push $\nabla$ through $g$), used in VAEs⁵. That requires $f$ to be differentiable and $p_{\theta }$ to be reparameterizable. The score-function trick needs neither — it only needs to evaluate $f$ and to have a differentiable density. That generality is why RL uses it, and the cost is higher variance.

Quiz: In ${\nabla }_{\theta }\mathbb{E}[f]=\mathbb{E}[f{\nabla }_{\theta }\log p_{\theta }]$, we assumed $f$ has no explicit $\theta$-dependence. In the RL objective, does the reward $R(\tau )$ depend on $\theta$?

Answer

No — the reward is a property of the environment: it’s a function of the states and actions in the trajectory, not of the policy parameters. The dependence of $J(\theta )$ on $\theta$ enters only through the trajectory distribution $p_{\theta }(\tau )$, i.e. through which trajectories are likely. This is exactly the setting where the score-function trick applies cleanly, with $f=R(\tau )$ and $p=p_{\theta }(\tau )$. (Subtlety: if you add an entropy bonus or a KL-to-reference term that is computed from ${\pi }_{\theta }$, that piece does depend on $\theta$ and needs its own gradient term — but the environment reward does not.)

---

4. The Policy Gradient Theorem

Now specialize $x=\tau$, $f=R(\tau )$, $p=p_{\theta }(\tau )$. Apply the trick:

$${\nabla }_{\theta }J(\theta )={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}[R(\tau ){\nabla }_{\theta }\log p_{\theta }(\tau )].$$

The magic happens when we expand $\log p_{\theta }(\tau )$ using the factorization from §2:

$$\log p_{\theta }(\tau )=\log {\rho }_0(s_0)+\sum _{t=0}^T[\log {\pi }_{\theta }(a_t\mid s_t)+\log P(s_{t+1}\mid s_t,a_t)].$$

Take ${\nabla }_{\theta }$. The initial-state term $\log {\rho }_0(s_0)$ and every transition term $\log P(s_{t+1}\mid s_t,a_t)$ are independent of $\theta$ — they’re properties of the environment. They vanish under the gradient:

$${\nabla }_{\theta }\log p_{\theta }(\tau )=\sum _{t=0}^T{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t).$$

The dynamics disappear

This is the single most important consequence. The unknown, possibly non-differentiable environment transition $P$ drops out of the gradient. We never need a model of the world — this is why policy gradients are model-free. All that survives is the score of our own policy.

Substituting back gives the raw form of the Policy Gradient Theorem⁶:

$${\nabla }_{\theta }J(\theta )={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}\left[\left(\sum _{t=0}^T{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)\right)\left(\sum _{t^{\prime }=0}^T{\gamma }^{t^{\prime }}{r}_{t^{\prime }}\right)\right].$$

The general form and the choice of ${\Psi }_t$

The theorem is usually written more flexibly. The gradient can be expressed as

$$\boxed{{\nabla }_{\theta }J(\theta )={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}\left[\sum _{t=0}^T{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t){\Psi }_t\right]}$$

where ${\Psi }_t$ is a weight on the score of action $a_t$ — telling the update how much to push up (or down) the log-probability of that action. Several choices of ${\Psi }_t$ all yield the same expected gradient but wildly different variance.⁷ In roughly increasing order of quality:

${\Psi }_t$	Name	Comment
$R(\tau )={\sum }_{t^{\prime }=0}^T{\gamma }^{t^{\prime }}{r}_{t^{\prime }}$	total return	valid but highest variance; every action credited with the whole trajectory
${\sum }_{t^{\prime }=t}^T{\gamma }^{t^{\prime }-t}{r}_{t^{\prime }}$	reward-to-go	drops rewards before $t$ (they can’t be caused by $a_t$)
$\left({\sum }_{t^{\prime }\ge t}{\gamma }^{t^{\prime }-t}{r}_{t^{\prime }}\right)-b(s_t)$	baselined reward-to-go	subtract a state-dependent baseline (§6)
$Q^{\pi }(s_t,a_t)$	action-value	expected reward-to-go
$A^{\pi }(s_t,a_t)=Q^{\pi }-V^{\pi }$	advantage	lowest variance; “how much better than average is $a_t$?”

The one-sentence intuition

Every choice of ${\Psi }_t$ says the same thing: increase the log-probability of actions that led to better-than-expected outcomes, decrease it for worse-than-expected ones. The differences are entirely about how tightly you estimate “better than expected” — which controls variance, not bias.

The last row is the destination of this lesson: the advantage $A^{\pi }$ is the ideal weight. Lesson 03 (03-actor-critic-and-gae) is about how to estimate it well.

Quiz: Why is it legitimate to replace the full return $R(\tau )$ with the reward-to-go ${\sum }_{t^{\prime }\ge t}{\gamma }^{t^{\prime }-t}{r}_{t^{\prime }}$ for the score at time $t$?

Answer

Causality. Action $a_t$ can only influence rewards at times $t^{\prime }\ge t$; rewards collected before $t$ are statistically independent of $a_t$ given the past. Formally, for $t^{\prime }<t$, $\mathbb{E}[{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)r_{t^{\prime }}]=0$ — the same zero-expectation argument as the baseline proof in §6, applied to a past reward that acts as a constant w.r.t. $a_t$. So dropping those terms changes nothing in expectation (unbiased) but removes their noise contribution (lower variance). Crediting an action for rewards it could not possibly have caused is pure noise.

---

5. REINFORCE — the Monte Carlo policy gradient

Williams (1992) gives the simplest concrete algorithm.⁴ Take ${\Psi }_t=$ reward-to-go, estimate the expectation with $N$ sampled trajectories:

$$\hat{g}=\frac{1}{N}\sum _{i=1}^N\sum _{t=0}^T{\nabla }_{\theta }\log {\pi }_{\theta }(a_t^{(i)}\mid s_t^{(i)})\left(\sum _{t^{\prime }=t}^T{\gamma }^{t^{\prime }-t}{r}_{t^{\prime }}^{(i)}\right),$$

then ascend: $\theta \leftarrow \theta +\alpha \hat{g}$.

graph LR
    A["Run policy π_θ<br/>collect N trajectories"] --> B["Compute returns /<br/>rewards-to-go G_t"]
    B --> C["Ĝ = mean of<br/>Σ_t ∇log π(a_t|s_t)·G_t"]
    C --> D["θ ← θ + α·Ĝ"]
    D --> A
    style C fill:#4a4,color:#fff

The full loop:

1. Sample $N$ trajectories by running the current ${\pi }_{\theta }$ in the environment.
2. For each, compute rewards-to-go $G_t={\sum }_{t^{\prime }\ge t}{\gamma }^{t^{\prime }-t}{r}_{t^{\prime }}$.
3. Form the estimator $\hat{g}$ above (in practice: build the surrogate loss $\frac{1}{N}{\sum }_i{\sum }_t\log {\pi }_{\theta }(a_t^{(i)}\mid s_t^{(i)})G_t^{(i)}$ and let autodiff produce $\hat{g}$).
4. Take a gradient-ascent step; repeat.

REINFORCE is on-policy and single-use

Step 1 must use the current ${\pi }_{\theta }$. Once you update $\theta$, the trajectories are stale — they were drawn from the old policy, so the expectation they estimate no longer matches. You must throw them away and resample. This is the root of REINFORCE’s sample inefficiency, and the reason importance-sampling / trust-region methods (PPO, later lessons) exist: to squeeze multiple update steps out of one batch of data.

The surrogate-loss view

You never hand-code $\nabla \log \pi$. You minimize $-\frac{1}{N}\sum \log {\pi }_{\theta }(a_t\mid s_t){\Psi }_t$ with ${\Psi }_t$ treated as a constant (detached / stop_gradient). This is weighted maximum likelihood: it’s ordinary supervised cross-entropy on the sampled actions, but each sample is weighted by how good its outcome was. Keep this framing — it’s exactly how RLHF code looks.

---

6. The variance problem and baselines

REINFORCE is unbiased but notoriously high-variance. Returns can be large in magnitude and the same magnitude gets multiplied onto the score of every action; a single lucky/unlucky trajectory swings the estimate hard. High variance means noisy gradients, tiny usable step sizes, and slow learning.

Baselines: subtract a reference without adding bias

The central trick: subtract a baseline $b(s_t)$ — any function of the state (not of the action) — from the weight:

$${\nabla }_{\theta }J(\theta )=\mathbb{E}\left[\sum _t{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)({\Psi }_t-b(s_t))\right].$$

This is only useful if it doesn’t change the gradient we’re estimating. It doesn’t — and here’s the proof, which is worth internalizing because the same identity underlies causality (§4) and GRPO (§8).

The workhorse identity: expected score is zero

For any distribution ${\pi }_{\theta }(\cdot \mid s)$,
${\mathbb{E}}_{a\sim {\pi }_{\theta }(\cdot \mid s)}[{\nabla }_{\theta }\log {\pi }_{\theta }(a\mid s)]=0.$
Proof.
${\mathbb{E}}_a[{\nabla }_{\theta }\log {\pi }_{\theta }(a\mid s)]={\sum }_a{\pi }_{\theta }(a\mid s)\frac{{\nabla }_{\theta }{\pi }_{\theta }(a\mid s)}{{\pi }_{\theta }(a\mid s)}={\sum }_a{\nabla }_{\theta }{\pi }_{\theta }(a\mid s)={\nabla }_{\theta }{\sum }_a{\pi }_{\theta }(a\mid s)={\nabla }_{\theta }1=0.$
The score of any properly normalized density has mean zero, because the density integrates to a constant.

Now the baseline term vanishes in expectation. Condition on $s_t$ and pull $b(s_t)$ out (it’s constant w.r.t. the action):

$${\mathbb{E}}_{a_t}[{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)b(s_t)]=b(s_t)\underset{=0}{\underbrace{{\mathbb{E}}_{a_t}[{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)]}}=0.$$

So subtracting any $b(s_t)$ leaves ${\nabla }_{\theta }J$ unbiased. The only thing it changes is variance.

The baseline must not depend on the action

If $b$ depended on $a_t$, the pull-out step fails — $b(s_t,a_t)$ can’t come outside the expectation over $a_t$, and $\mathbb{E}[b(s,a)\nabla \log \pi ]\ne 0$ in general. You’d bias the gradient. Baselines are strictly state-dependent. (This is also why the advantage $A(s,a)=Q(s,a)-V(s)$ keeps the action-dependent $Q$ as the weight and uses the state-only $V$ as the baseline — not the other way around.)

Why it reduces variance — and the value function as the near-optimal baseline

Subtracting a constant doesn’t change the mean but does change the second moment. The variance of the per-term estimator is governed by $\mathbb{E}[(\nabla \log \pi {)}^2({\Psi }_t-b{)}^2]$. Choosing $b(s_t)\approx \mathbb{E}[{\Psi }_t\mid s_t]$ centers the weight around zero, shrinking $({\Psi }_t-b{)}^2$ and hence the variance.

The natural choice is the state-value function:

$$b(s_t)=V^{\pi }(s_t)={\mathbb{E}}_{a\sim \pi }[Q^{\pi }(s_t,a)].$$

With ${\Psi }_t=Q^{\pi }(s_t,a_t)$ and $b=V^{\pi }(s_t)$, the weight becomes exactly the advantage:

$${\Psi }_t-b(s_t)=Q^{\pi }(s_t,a_t)-V^{\pi }(s_t)=A^{\pi }(s_t,a_t).$$

Advantage is the punchline

$A^{\pi }(s,a)$ answers: “was $a$ better or worse than the policy’s average action in state $s$?” Its sign is the natural learning signal — push up above-average actions, push down below-average ones — and centering at $V$ makes the magnitude small, minimizing variance. Using $V$ as the baseline isn’t provably variance-optimal (the true minimum-variance baseline is a score-weighted version of $\mathbb{E}[\Psi ]$), but it’s cheap, interpretable, and nearly optimal in practice. This is why the next lesson is “actor-critic”: you learn $V_{\varphi }\approx V^{\pi }$ so you can form $A$.

Quiz (derivation check): Prove that subtracting a state-dependent baseline $b(s_t)$ leaves the policy gradient unbiased. Where exactly would the argument break if $b$ depended on the action?

Answer

We need ${\mathbb{E}}_{\tau }[{\sum }_t{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)b(s_t)]=0$. Take a single term and use the tower rule, conditioning on everything up to and including $s_t$:
$\mathbb{E}[{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)b(s_t)]={\mathbb{E}}_{s_t}[b(s_t){\mathbb{E}}_{a_t\sim {\pi }_{\theta }(\cdot \mid s_t)}[{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)]].$
The inner expectation is the expected score, which is $0$ by the identity in the tip above. So the whole term is ${\mathbb{E}}_{s_t}[b(s_t)\cdot 0]=0$, for every $t$. Bias unchanged. ∎

Where it breaks with $b(s_t,a_t)$: the pull-out $b(s_t){\mathbb{E}}_{a_t}[\cdots ]$ requires $b$ to be constant over the inner expectation. If $b$ depends on $a_t$, it stays inside: you’d need ${\mathbb{E}}_{a_t}[b(s_t,a_t)\nabla \log {\pi }_{\theta }]$, which equals ${\sum }_{a}b(s,a)\nabla {\pi }_{\theta }(a\mid s)$ — generally nonzero. You’d introduce bias.

Quiz: You have two candidate baselines: (a) $b=0$, (b) $b=\frac{1}{T}{\sum }_t{r}_t$ (the batch-average reward, a constant). Both are unbiased. Which is better and why? Could a per-state $V(s)$ beat both?

Answer

(b) beats (a): even a single global constant equal to the mean return recenters ${\Psi }_t$ around zero, cutting variance substantially — this is the cheapest useful baseline and is essentially what “reward whitening / mean-subtraction” does in practice. But a per-state $b=V^{\pi }(s)$ beats both, because returns genuinely differ by state: a good state has high expected return regardless of the action taken, and a constant baseline can’t subtract that off state-by-state. $V^{\pi }(s)$ removes the state-value component of the return, leaving only the action-specific advantage — much tighter. (Note: GRPO’s group-relative baseline in §8 is a modern reincarnation of a per-prompt constant baseline that avoids learning $V$ at all.)

---

7. Practical issues, at a glance

Even with baselines, vanilla policy gradients carry three chronic problems — keep these in mind, they motivate everything downstream.

• High variance. The score-function estimator is inherently noisier than pathwise gradients. Baselines, reward-to-go, and advantage estimation (GAE, lesson 03) all fight this.
• Sample inefficiency. On-policy: data is discarded after one update (§5). PPO’s clipped importance ratio lets you reuse a batch for several epochs — the practical fix.⁸
• Credit assignment. A single scalar return must be apportioned across a long trajectory. Reward-to-go and $\gamma$-discounting help, but with sparse, delayed rewards (win/loss at the end of a game; a correct final answer after 500 tokens) the signal per action is faint. This is the hard problem, and it is exactly the LLM setting.

Discount factor is doing double duty

$\gamma <1$ is often introduced as “we care less about the future,” but in policy gradients it also acts as a variance-reduction / credit-horizon knob: it down-weights far-future rewards that are only weakly caused by the current action. Lowering $\gamma$ reduces variance but adds bias. This bias–variance dial is formalized by GAE’s $\lambda$ in lesson 03.

---

8. The LLM connection — RLHF, RLVR, and GRPO

This is where the whole apparatus becomes the backbone of post-training.⁹ Map the RL vocabulary onto language modeling:

RL concept	LLM instantiation
policy ${\pi }_{\theta }$	the language model (its softmax over the vocabulary)
state $s_t$	the prompt + tokens generated so far (the context)
action $a_t$	the next token $y_t\in \mathcal{V}$
$\log {\pi }_{\theta }(a_t\mid s_t)$	the token log-probability $\log p_{\theta }(y_t\mid x,y_{<t})$
trajectory $\tau$	prompt $x$ + full generated response $y_{1:T}$
reward $R(\tau )$	scalar from a reward model (RLHF) or a verifier (RLVR)

The autoregressive factorization $\log {\pi }_{\theta }(y_{1:T}\mid x)={\sum }_t\log p_{\theta }(y_t\mid x,y_{<t})$ is literally the trajectory-log-prob sum from §4. So the policy gradient for an LLM is

$${\nabla }_{\theta }J(\theta )={\mathbb{E}}_{x,y\sim {\pi }_{\theta }}\left[\sum _{t=1}^T{\nabla }_{\theta }\log p_{\theta }(y_t\mid x,y_{<t}){\Psi }_t\right].$$

This is just REINFORCE on token sequences. Generating a completion is rolling out a trajectory; the token head is the action distribution; huge vocabulary is exactly the “large action space” that killed value-based $\arg \max$ in §2 — which is why policy-gradient methods, not Q-learning, dominate LLM post-training.

Sequence reward → per-token credit

In RLHF/RLVR the reward is usually sequence-level: one scalar for the whole response (e.g. reward-model score, or $1$ if the answer passes the verifier else $0$). Through the return, that single scalar becomes the weight ${\Psi }_t$ on every token’s score. With no intermediate reward, the reward-to-go is the same terminal scalar for all $t$ — so the model uplifts (or suppresses) the log-prob of all tokens in a good (bad) response uniformly. That’s the credit-assignment problem of §7 in its starkest form: the model must sort out which tokens actually mattered purely from statistics across many samples.

Common misconception

“RLHF uses PPO, which is totally different from REINFORCE.” No — PPO is a variance-reduced, sample-reused, trust-region wrapper around the same policy-gradient estimator. Strip away the clipped importance ratio and the learned value baseline, and PPO’s core is the advantage-weighted score sum you derived in §4–§6. Recent methods (RLOO¹⁰, GRPO¹¹) go the other direction — back toward plain REINFORCE — arguing the PPO machinery is often unnecessary for LLMs.

GRPO teaser (lesson 09). A learned value baseline $V_{\varphi }(s)$ is expensive for LLMs (a second large network, and per-token value targets are hard). GRPO (Group Relative Policy Optimization) drops the learned critic entirely.¹¹ For a prompt $x$, it samples a group of $G$ completions $\{y^{(1)},\dots ,y^{(G)}\}$, scores each with reward $R^{(i)}$, and uses the group statistics as the baseline:

$$A^{(i)}=\frac{R^{(i)}-\mathrm{mean}(\{R^{(j)}{\}}_{j=1}^G)}{\mathrm{std}(\{R^{(j)}{\}}_{j=1}^G)}.$$

The group mean is a baseline in the exact sense of §6 — a quantity that depends on the state (the prompt) but not on the specific sampled action — so it’s unbiased and needs no critic network. This is the per-prompt-constant baseline from the §6 quiz, made concrete at scale. We derive GRPO fully in lesson 09.

Quiz: In LLM RL with a sequence-level 0/1 verifier reward and no baseline, what does the gradient do to a completion that gets reward 0, versus one with reward 1? Why is a group-relative baseline important here?

Answer

With ${\Psi }_t=R\in \{0,1\}$ and no baseline: a reward-1 completion contributes $+\nabla \log \pi$ on all its tokens (push their probabilities up); a reward-0 completion contributes $0\cdot \nabla \log \pi =0$ — no learning signal at all from failures. The update only ever reinforces successes and never explicitly suppresses failures, and its magnitude scales with the raw success rate (a batch that’s mostly correct barely moves; mostly-wrong barely moves). Subtracting the group-mean baseline fixes both: completions above the group’s average reward get positive advantage (reinforced), those below get negative advantage (suppressed) — so failures now push probabilities down, and the signal is centered regardless of the absolute success rate. That recentring is precisely the variance-reduction-without-bias argument of §6.

---

9. Practice problems

Practice problem 1: Derive the minimum-variance constant baseline

Problem: Consider the single-parameter, single-step estimator $\hat{g}(b)=(R-b){\nabla }_{\theta }\log {\pi }_{\theta }(a)$ with scalar $b$. You proved any $b$ is unbiased. Find the $b^{\star }$ that minimizes $\mathrm{Var}(\hat{g})$. Interpret it.

Worked solution:
Since $\hat{g}$ is unbiased, minimizing variance means minimizing the second moment $\mathbb{E}[{\hat{g}}^2]$. Let $s:={\nabla }_{\theta }\log {\pi }_{\theta }(a)$ (a vector; take a scalar component or its squared-norm form). Then
$\mathbb{E}[{\hat{g}}^2]=\mathbb{E}[(R-b{)}^2{s}^2]=\mathbb{E}[R^2{s}^2]-2b\mathbb{E}[R{s}^2]+b^2\mathbb{E}[s^2].$
Differentiate w.r.t. $b$ and set to zero:
$\frac{d}{db}\mathbb{E}[{\hat{g}}^2]=-2\mathbb{E}[R{s}^2]+2b\mathbb{E}[s^2]=0\Rightarrow \boxed{b^{\star }=\frac{\mathbb{E}[R{s}^2]}{\mathbb{E}[s^2]}}.$
Interpretation: the optimal baseline is a score-magnitude-weighted average of the return — not simply $\mathbb{E}[R]$. It’s $\mathbb{E}[R]$ reweighted by $\Vert \nabla \log \pi {\Vert }^2$. In practice people use $V(s)\approx \mathbb{E}[R\mid s]$ because it’s interpretable, gives the advantage, and captures most of the benefit; the $s^2$-weighting correction is usually second-order and rarely worth the complexity. (This is the “optimal-ish” hedge from §6.)

Practice problem 2: Why the naive full-return estimator wastes signal

Problem: Two-step episode ($t=0,1$), $\gamma =1$. Rewards $r_0,r_1$. Write the full-return policy gradient and the reward-to-go policy gradient. Show they have the same expectation, and identify the specific term that reward-to-go removes and why it’s pure noise.

Worked solution:
Full return $R=r_0+r_1$ applied to both scores:
${\hat{g}}_{\text{full}}=\nabla \log \pi (a_0\mid s_0)(r_0+r_1)+\nabla \log \pi (a_1\mid s_1)(r_0+r_1).$
Reward-to-go: action $a_0$ gets $r_0+r_1$, action $a_1$ gets only $r_1$:
${\hat{g}}_{\text{rtg}}=\nabla \log \pi (a_0\mid s_0)(r_0+r_1)+\nabla \log \pi (a_1\mid s_1)(r_1).$
The difference is the term $\nabla \log \pi (a_1\mid s_1)r_0$. In expectation it’s zero: $r_0$ is determined by $s_0,a_0$ (before $a_1$ is even chosen), so conditioning on $s_1$ and taking ${\mathbb{E}}_{a_1}$,
$\mathbb{E}[\nabla \log \pi (a_1\mid s_1)r_0]=\mathbb{E}\left[r_0\underset{=0}{\underbrace{{\mathbb{E}}_{a_1}[\nabla \log \pi (a_1\mid s_1)]}}\right]=0.$
So both estimators are unbiased — but that dropped term has nonzero variance and contributes only noise (it credits $a_1$ for a reward that occurred before $a_1$ was taken, which is causally impossible). Removing it strictly lowers variance. This is the §4 causality quiz made concrete.

---

10. Reading order

Work through these in order:

1. Foundational text: Sutton & Barto, Reinforcement Learning: An Introduction (2nd ed.), Chapter 13 — Policy Gradient Methods. The canonical derivation of the Policy Gradient Theorem and REINFORCE with baseline. Read §13.1–13.4 first.
2. Best intuition-builder: OpenAI Spinning Up — “Intro to Policy Optimization” (spinningup.openai.com). Rederives the gradient, the reward-to-go trick, and baselines with runnable code (vpg). The clearest single resource for the estimator-implementation mapping in §5.
3. Primary source: Williams (1992), “Simple statistical gradient-following algorithms for connectionist reinforcement learning” — the original REINFORCE paper. Read for the score-function estimator and the baseline discussion in its original framing.
4. Forward pointer (read after lesson 03): Schulman et al. (2015), “High-Dimensional Continuous Control Using Generalized Advantage Estimation” (arXiv:1506.02438) — GAE, the bias–variance-controlled advantage estimator that makes the ${\Psi }_t=A^{\pi }$ choice practical. This is the bridge to actor-critic.

---

11. What’s next

Two threads, in order:

1. 03-actor-critic-and-gae — you now know the advantage $A^{\pi }$ is the ideal weight ${\Psi }_t$, but you can’t compute it exactly. Actor-critic learns a value network $V_{\varphi }$ as the baseline; GAE ($\lambda$) tunes the bias–variance tradeoff in estimating $A$. This is the direct sequel.
2. Back-reference: 01-mdps-and-the-rl-objective — if the trajectory factorization, $V^{\pi }/Q^{\pi }/A^{\pi }$, or the objective $J(\theta )$ felt shaky in §2–§4, revisit it; everything here rests on it.

Further out: PPO (trust-region / clipped surrogate for sample reuse) and lesson 09 on GRPO (group-relative baseline, critic-free RL for LLMs) both build directly on the estimator derived here.

---

References

---

Topic hub: rl-for-llms | Filed: 2026-09-02

Footnotes

1. Watkins, C. J. C. H., & Dayan, P. (1992). “Q-learning.” Machine Learning, 8(3–4), 279–292. https://doi.org/10.1007/BF00992698 ↩

2. Mnih, V., Kavukcuoglu, K., Silver, D., Rusu, A. A., Veness, J., Bellemare, M. G., et al. (2015). “Human-Level Control Through Deep Reinforcement Learning” (DQN). Nature, 518, 529–533. https://doi.org/10.1038/nature14236 ↩

3. Grattafiori, A., et al. (Meta AI) (2024). “The Llama 3 Herd of Models.” arXiv:2407.21783. Llama 3 uses a tokenizer with a vocabulary of 128,256 tokens. https://arxiv.org/abs/2407.21783 ↩

4. Williams, R. J. (1992). “Simple Statistical Gradient-Following Algorithms for Connectionist Reinforcement Learning.” Machine Learning, 8(3–4), 229–256. Introduces REINFORCE and the score-function / likelihood-ratio estimator with baselines. https://doi.org/10.1007/BF00992696 ↩ ↩²

5. Kingma, D. P., & Welling, M. (2013). “Auto-Encoding Variational Bayes.” arXiv:1312.6114. Introduces the reparameterization (pathwise) gradient estimator used in VAEs. https://arxiv.org/abs/1312.6114 ↩

6. Sutton, R. S., McAllester, D., Singh, S., & Mansour, Y. (2000). “Policy Gradient Methods for Reinforcement Learning with Function Approximation.” Advances in Neural Information Processing Systems 12 (NIPS 1999), 1057–1063. The canonical statement and proof of the Policy Gradient Theorem. https://papers.nips.cc/paper/1713-policy-gradient-methods-for-reinforcement-learning-with-function-approximation ↩

7. Schulman, J., Moritz, P., Levine, S., Jordan, M., & Abbeel, P. (2015). “High-Dimensional Continuous Control Using Generalized Advantage Estimation” (GAE). arXiv:1506.02438. Frames the choice of weight ${\Psi }_t$ (return, reward-to-go, $Q$, advantage) and the $\lambda$ bias–variance tradeoff. https://arxiv.org/abs/1506.02438 ↩

8. Schulman, J., Wolski, F., Dhariwal, P., Radford, A., & Klimov, O. (2017). “Proximal Policy Optimization Algorithms.” arXiv:1707.06347. Clipped surrogate objective enabling multiple update epochs per batch. https://arxiv.org/abs/1707.06347 ↩

9. Ouyang, L., Wu, J., Jiang, X., Almeida, D., Wainwright, C. L., Mishkin, P., Zhang, C., Agarwal, S., Slama, K., Ray, A., et al. (2022). “Training Language Models to Follow Instructions with Human Feedback” (InstructGPT). arXiv:2203.02155. The canonical RLHF-with-PPO post-training recipe. https://arxiv.org/abs/2203.02155 ↩

10. Ahmadian, A., Cremer, C., Gallé, M., Fadaee, M., Kreutzer, J., Pietquin, O., Üstün, A., & Hooker, S. (2024). “Back to Basics: Revisiting REINFORCE Style Optimization for Learning from Human Feedback in LLMs” (RLOO). arXiv:2402.14740. Argues PPO’s machinery is often unnecessary for RLHF and that REINFORCE/RLOO variants match or exceed it. https://arxiv.org/abs/2402.14740 ↩

11. Shao, Z., Wang, P., Zhu, Q., Xu, R., Song, J., Bi, X., Zhang, H., Zhang, M., Li, Y. K., Wu, Y., & Guo, D. (2024). “DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models.” arXiv:2402.03300. Introduces Group Relative Policy Optimization (GRPO), which replaces the learned critic with a group-relative (per-prompt) baseline. https://arxiv.org/abs/2402.03300 ↩ ↩²