Markdown

Policy gradient foundations for LLM RL¶

Scope: the shared mathematical core that PPO, GRPO, REINFORCE and RLOO, and every GRPO variant instantiate differently. This page covers the score-function (log-derivative) estimator, the value, action-value, and advantage functions, why a baseline reduces variance without introducing bias, reward-to-go and credit assignment, and the two ways an LLM completion gets modelled as a trajectory. Read it before the optimizer pages if the difference between a token-level and a sequence-level objective is not already obvious; the post-training system map places all of it in one pipeline.

The numpy block is executed and asserted in this page. It enumerates a four-action bandit exactly, so the identities are proven rather than sampled, and the one Monte Carlo claim is labelled as such.

What it is¶

Nearly every RL algorithm used to train an LLM is a policy gradient algorithm: it estimates the gradient of expected return with respect to the policy parameters, then takes a gradient ascent step. All of them share one structure, a per-action log-probability gradient multiplied by a scalar score:

grad J(theta) = E[ sum_t  psi_t * grad_theta log pi_theta(a_t | s_t) ]

The algorithms differ only in what they put in psi_t and how they constrain the resulting step:

`psi_t`	Algorithm	Cost
Trajectory return `R(tau)`	Vanilla policy gradient	Highest variance, no extra model
Return minus a batch-average baseline	REINFORCE	One rollout per prompt
Return minus a leave-one-out group baseline	RLOO	K rollouts per prompt
Group-normalized reward	GRPO	G rollouts per prompt
Learned advantage (GAE over a critic)	PPO	A second trained model

The sign convention is the whole intuition: psi_t > 0 pushes the policy to make action a_t more likely, psi_t < 0 pushes it to make a_t less likely, and psi_t near zero produces almost no update. Everything else in RL research for LLMs is an argument about how to compute that scalar with less variance, less bias, or less compute.

Value, action-value, and advantage¶

Three expectations recur, and confusing them is the most common source of a wrong advantage implementation:

Value V(s): expected return starting from state s and acting under the current policy.
Action-value Q(s, a): expected return starting from s, taking action a, then acting under the policy.
Advantage A(s, a) = Q(s, a) - V(s): how much better action a is than the policy's own average from that state.

The sampled return G_t observed after taking a_t is a single-sample Monte Carlo estimate of Q(s_t, a_t). That gives two routes to an advantage estimate, and the split defines the whole algorithm family: subtract a cheap action-independent baseline (REINFORCE, RLOO, GRPO), or learn V with a critic and subtract its prediction (PPO). A critic is an LLM backbone with a scalar head, architecturally near-identical to a reward model, with one decisive difference: a reward model scores a finished completion and stays frozen, while a critic predicts expected future return at every token position and must be retrained continuously because V is defined against the current policy.

Two ways to model a completion¶

The same completion can be modelled as a Markov decision process, where each token is an action and the state is the prompt plus the tokens emitted so far, or as a bandit, where the entire completion is one action that receives one outcome reward. Both are used in production. REINFORCE and RLOO are usually written as bandits; PPO is usually written as an MDP; GRPO mixes them, computing a sequence-level advantage and then broadcasting it to a token-level objective.

The MDP transition function is deterministic for a plain LLM completion: the next state is the current state with one token concatenated. That is why transition dynamics, which dominate classical RL, barely appear in LLM post-training until tool calls and environments enter the loop (agentic RL), where the environment genuinely does return states the policy did not generate.

Outcome versus process rewards¶

An outcome reward scores the whole completion; a process reward scores intermediate steps. Outcome rewards dominate because they are what a verifier naturally produces (RLVR), but they are sparse: one scalar has to explain a 4,000-token trace. Process rewards give finer credit assignment at the cost of needing a step-level signal. The executed block below shows why this distinction is not cosmetic: under an outcome-only reward with no discounting, the token-level and sequence-level objectives are algebraically the same expression, and they stop being the same the moment a single intermediate reward appears.

Why use it¶

It tells you which knob to turn. Entropy collapse, length bias, and zero-gradient batches are all psi_t problems. Knowing that DAPO, Dr. GRPO, and GSPO each modify one term in this one expression turns the variant zoo into a short list of orthogonal choices.
It bounds what a fix can do. A baseline cannot change the expected gradient, only its variance. Any proposal claiming a baseline improves the direction of the update is either changing the objective or is wrong.
It explains the memory bill. psi_t from a learned critic means a second trainable model, which is roughly 16 GB per billion parameters in half precision with AdamW, against roughly 2 GB per billion for an inference-only model. The choice of estimator sets the cluster shape before any code is written.
It makes the token-versus-sequence debate concrete. Whether an importance ratio, a loss normalization, or an advantage lives at the token or the sequence level is the single most common cause of a silently biased LLM RL run.

When to use it (and when not)¶

Read this first when an RL run is unstable and the reward is not obviously broken. Most instability is variance in the gradient estimator, not a reward bug.
Read this first when porting an objective between libraries. TRL, verl, and slime differ in loss aggregation and advantage normalization defaults, and those are exactly the terms defined here.
Skip to GRPO if the goal is to launch a standard verifiable-reward run on a supported stack; the defaults are sane and this page is theory you can return to.
Skip to DPO if there is no rollout budget at all. DPO is not a policy gradient method and none of this applies to it directly.

Architecture¶

flowchart TB
  OBJ["RL objective: maximise E[return]"] --> SF["Score-function estimator<br/>grad log pi times psi_t"]
  SF --> PSI{"What is psi_t?"}
  PSI -->|"return"| VPG["Vanilla policy gradient"]
  PSI -->|"return minus baseline"| BASE["REINFORCE / RLOO"]
  PSI -->|"group-normalized reward"| GRP["GRPO"]
  PSI -->|"critic advantage (GAE)"| AC["PPO actor-critic"]
  BASE --> CON{"Constrain the step?"}
  GRP --> CON
  AC --> CON
  CON -->|"KL penalty"| KL["KL regularization"]
  CON -->|"clipped importance ratio"| CLIP["PPO / GRPO clipping"]
  CON -->|"hard KL constraint"| TR["TRPO trust region"]

The left branch is a variance decision and the right branch is a stability decision. They are independent: any psi_t composes with any constraint, which is why the recent literature is mostly a cross product of the two.

How to use it¶

The two identities that make the whole family work are the score function and the EGLP lemma. Both are provable by exact enumeration on a small policy, with no sampling noise to hide behind, and the block below does that before checking the one claim that genuinely needs Monte Carlo.

# pg_foundations.py -- the two identities every policy-gradient algorithm rests on:
# the score function (log-derivative trick) and the EGLP lemma that makes baselines free.
# Exact enumeration over a 4-action bandit, so nothing here is a sampling artefact. numpy only.
import numpy as np

theta = np.array([0.4, -0.2, 1.1, 0.0])          # policy logits
R = np.array([1.0, 0.0, 0.0, 1.0])               # reward per action (outcome reward)

def pi(th):
    e = np.exp(th - th.max())
    return e / e.sum()

def score(a, th):                                 # grad_theta log pi(a) for softmax
    g = -pi(th)
    g[a] += 1.0
    return g

def exact_pg(th, psi):                            # sum_a pi(a) * psi(a) * grad log pi(a)
    p = pi(th)
    return sum(p[a] * psi[a] * score(a, th) for a in range(len(th)))

p = pi(theta)
g_true = exact_pg(theta, R)

# (1) SCORE FUNCTION: the estimator matches the analytic gradient of E[R] to machine
#     precision. Finite differences on the true objective are the independent reference.
def J(th):
    return float(pi(th) @ R)
fd = np.array([(J(theta + 1e-6 * np.eye(4)[i]) - J(theta - 1e-6 * np.eye(4)[i])) / 2e-6
               for i in range(4)])
assert np.allclose(g_true, fd, atol=1e-8), (g_true, fd)

# (2) EGLP LEMMA: E[b * grad log pi] == 0 for ANY action-independent b, so subtracting a
#     baseline leaves the gradient untouched. Tested with a deliberately absurd baseline.
for b in (0.0, float(p @ R), -1e4):
    assert np.allclose(exact_pg(theta, R - b), g_true, atol=1e-12)
    assert np.allclose(exact_pg(theta, np.full(4, b)), 0.0, atol=1e-12)

# (3) ADVERSARIAL: an ACTION-DEPENDENT baseline breaks the lemma. b(a) = R(a) is the
#     limiting case and zeroes the gradient entirely, which is the wrong answer, not a
#     smaller one. This is why RLOO excludes a completion's own reward from its baseline.
g_bad = exact_pg(theta, R - R)
assert np.allclose(g_bad, 0.0) and not np.allclose(g_true, 0.0)

# (4) VARIANCE: the baseline is free in expectation but not in variance. Monte Carlo the
#     single-sample estimator both ways and compare the trace of its covariance.
rng = np.random.default_rng(0)
acts = rng.choice(4, size=200_000, p=p)
est0 = np.array([R[a] * score(a, theta) for a in acts])
est1 = np.array([(R[a] - p @ R) * score(a, theta) for a in acts])
v0, v1 = est0.var(axis=0).sum(), est1.var(axis=0).sum()
assert v1 < v0
assert np.allclose(est0.mean(axis=0), g_true, atol=5e-3)   # both still unbiased
assert np.allclose(est1.mean(axis=0), g_true, atol=5e-3)

# (5) MDP vs BANDIT: with an outcome-only reward and gamma = 1, the token-level reward-to-go
#     is the SAME scalar at every token, so the token-level (MDP) loss is algebraically the
#     sequence-level (bandit) loss. The two formulations are not two algorithms here.
logps = np.array([-0.31, -1.20, -0.05, -2.40])        # per-token log-probs of one completion
outcome = 0.7
per_token_r = np.array([0.0, 0.0, 0.0, outcome])
rtg = np.cumsum(per_token_r[::-1])[::-1]              # reward-to-go
assert np.allclose(rtg, outcome)
assert np.isclose(-(logps * rtg).sum(), -(logps.sum() * outcome))

# (6) BOUNDARY: add one process reward mid-sequence and the equivalence breaks, which is
#     exactly the credit assignment a token-level formulation buys.
per_token_r2 = np.array([0.0, 0.3, 0.0, outcome])
rtg2 = np.cumsum(per_token_r2[::-1])[::-1]
assert not np.isclose(-(logps * rtg2).sum(), -(logps.sum() * per_token_r2.sum()))

print("exact policy gradient:", np.round(g_true, 4))
print("variance no baseline / with baseline:", round(float(v0), 4), "/", round(float(v1), 4))
print("reward-to-go, outcome only:", rtg, "| with a process reward:", rtg2)

Executed output:

exact policy gradient: [ 0.143  -0.0512 -0.1877  0.0959]
variance no baseline / with baseline: 0.2918 / 0.1134
reward-to-go, outcome only: [0.7 0.7 0.7 0.7] | with a process reward: [1.  1.  0.7 0.7]

Check (4) is the practical headline: the baseline cut the trace of the estimator covariance from 0.2918 to 0.1134, a factor of 2.6, while leaving the expected gradient where it was. Check (3) is the trap the arithmetic of RLOO and GRPO turns on, and it is why footnote-level details about whether a sample appears in its own baseline are not pedantry.

How to develop with it¶

Framework code never computes a gradient expression directly. It builds a surrogate loss whose autodiff gradient equals the intended policy gradient, and that indirection is where implementations go wrong.

# Reference template (needs torch). The surrogate whose backward pass IS the vanilla
# policy gradient. Nothing here is executed on this page; the maths it implements is
# the sequence-level case validated above.
import torch

logits = policy(input_ids)[:, :-1, :]                       # (B, L-1, V)
sampled = input_ids[:, 1:]                                  # (B, L-1)
per_token_logps = torch.log_softmax(logits, dim=-1).gather(
    dim=-1, index=sampled.unsqueeze(-1)).squeeze(-1)        # (B, L-1)

seq_logps = (per_token_logps * completion_mask).sum(dim=-1)  # (B,)
loss = -(seq_logps * advantages.detach()).mean()             # negative: torch descends
loss.backward()

Three rules follow from the derivation and each maps to a real defect class:

Detach the scalar. advantages.detach() is not an optimization. The derivation treats psi_t as a constant; leaving a gradient path through it optimizes a different objective. This bites hardest when a KL term is folded into the reward, because the policy log-probs appear on both sides.
Mask before summing, and know what the denominator is. The prompt and padding tokens must not contribute. Whether the sum is divided by that sequence's own length, by the batch token count, or by a fixed constant is precisely the Dr. GRPO and DAPO length-bias argument, not a formatting choice.
Shift once. The logit at position t predicts the token at t+1. An off-by-one here trains the model on the wrong action and still produces a plausible loss curve.

How to maintain it¶

Instrument the estimator, not just the reward. The quantities that reveal a broken policy gradient before the reward curve does:

Gradient norm and its variance across microbatches. Rising variance at flat reward is the signature of a failing baseline (degenerate groups, or a critic that has stopped tracking the policy).
Fraction of zero-advantage samples. In group methods this is frac_reward_zero_std; every such sample is compute spent on a guaranteed zero gradient.
Entropy of the next-token distribution. The policy gradient increases the probability of reinforced actions without bound unless something opposes it; monotone entropy decay is the run telling you exploration has ended.
Advantage mean and standard deviation per batch. A mean drifting away from zero means the baseline is biased, which by the EGLP lemma means it has become action-dependent somewhere.

How to run it in production¶

The estimator choice sets the cluster shape, so decide it before provisioning:

Critic-free (REINFORCE, RLOO, GRPO) trains one model. Budget the rollout fleet as a serving deployment, because generation dominates wall-clock (async and disaggregated RL, rollout fleet sizing).
Actor-critic (PPO) trains two. At roughly 16 GB per billion parameters for a trained model against roughly 2 GB per billion for an inference-only copy, adding a critic costs far more than adding a frozen reference or reward model. PPO covers the four-model arithmetic.
Either way, the rollout engine and the trainer disagree on token probabilities. That disagreement is an importance-sampling problem, not a numerical bug, and it is corrected rather than eliminated (truncated importance sampling, async RL systems).
Discounting is usually off. LLM post-training normally runs gamma = 1 with a finite horizon because the reward arrives at the completion level. The exceptions are real: KL penalties are commonly expressed as a per-token reward, and GAE uses gamma and lambda when estimating the advantage.

Failure modes¶

Gradient flows through the advantage. Forgetting .detach() silently changes the objective. Symptom: loss falls while reward does not move.
Action-dependent baseline. A batch mean that includes the sample it baselines is biased, per check (3). Symptom: advantages with a persistent non-zero mean.
Off-by-one in the causal shift. Trains on the wrong action, produces a smooth and meaningless loss curve.
Prompt tokens in the loss. Unmasked prompt log-probs turn the objective into a mix of policy gradient and unintended language modelling.
Sparse outcome reward on very long traces. One scalar over thousands of tokens is a credit-assignment problem, not a tuning problem; process rewards or a group baseline are the answers, not a larger learning rate.
Treating variance as bias. Reaching for a KL penalty or a smaller clip when the real problem is a noisy return estimate suppresses learning instead of stabilising it.

References¶

Sutton and Barto, Reinforcement Learning: An Introduction (2nd ed.): http://incompleteideas.net/book/the-book-2nd.html
Williams, Simple statistical gradient-following algorithms for connectionist reinforcement learning (REINFORCE, 1992): https://link.springer.com/article/10.1007/BF00992696
Schulman et al., High-Dimensional Continuous Control Using Generalized Advantage Estimation: https://arxiv.org/abs/1506.02438
OpenAI Spinning Up, Intro to Policy Optimization (score function, EGLP lemma, reward-to-go): https://spinningup.openai.com/en/latest/spinningup/rl_intro3.html
Lambert, The RLHF Book, Policy Gradient Algorithms: https://rlhfbook.com/c/11-policy-gradients.html
Ahmadian et al., Back to Basics: Revisiting REINFORCE Style Optimization for Learning from Human Feedback in LLMs: https://arxiv.org/abs/2402.14740
Cameron R. Wolfe, Reinforcement Learning for LLMs: The Complete Guide: https://cameronrwolfe.substack.com/p/llm-rl
Cameron R. Wolfe, Policy Gradients: The Foundation of RLHF: https://cameronrwolfe.substack.com/p/policy-gradients-the-foundation-of