How policy-gradient methods became practical for reasoning models
May 2026 · Cédric Caruzzo
TL;DR. This post derives the main line of policy-gradient methods used to train language models: REINFORCE → Actor-Critic → GAE → TRPO → PPO → RLHF → GRPO → Dr.GRPO. Each section introduces a limitation and the method that addresses it.
The chain runs from REINFORCE (1992) to Dr.GRPO (2025). Basic gradients and probability are the only prerequisites.
Start with the mathematical setup.
We start with a prompt, defined as a state $s$, and the model, defined as the policy and denoted $\pi_\theta$. When we pass the prompt to the model, the model takes an action $u$ by sampling from the policy given the current state. The process is autoregressive, so we add a subscript to refer to the current step. The full trajectory ends when the model generates the <EOS> token:
And at each time step, the environment can give us a scalar reward. The total reward over a full trajectory is:
The reward function is the only way we can communicate to the model what we want it to do. It is a function that takes the current state and action, and outputs a number that tells us how good or bad that action was in that state. For example, if we want the model to generate a correct answer to a question, we can give it a reward of +1 for a correct answer and 0 for an incorrect one.
The reward can be sparse (only given at the end of the trajectory), for example when training a LLM to reason, or dense (given at every step), such as teaching an agent to play a game.
The goal is to choose actions $u_t$ from states $s_t$ that maximize expected reward.
Written out:
where $P(\tau;\theta)$ is the probability of trajectory $\tau$ under policy $\pi_\theta$.
therefore:
The Objective Function: we want to find the parameters $\theta$ that maximize the expected reward over all possible trajectories.
So how do we achieve that? Like any other optimization problem: compute the gradient of the utility, and update the parameters in that direction. The plot below illustrates that idea visually.
Compute the gradient step by step:
Step 1. Linearity of differentiation: push the gradient inside the sum. $R(\tau)$ doesn't depend on $\theta$, so it just tags along:
$$\nabla_\theta U(\theta) = \nabla_\theta \sum_\tau P(\tau;\theta) \, R(\tau) = \sum_\tau \nabla_\theta P(\tau;\theta) \, R(\tau)$$Step 2. We have $\nabla_\theta P(\tau;\theta)$ floating around, but no distribution to sample from. Multiply and divide by $P(\tau;\theta)$ to introduce one:
$$= \sum_\tau P(\tau;\theta) \cdot \frac{\nabla_\theta P(\tau;\theta)}{P(\tau;\theta)} \cdot R(\tau)$$Step 3. The fraction $\frac{\nabla_\theta P(\tau;\theta)}{P(\tau;\theta)}$ is exactly $\nabla_\theta \log P(\tau;\theta)$ by the log-derivative identity. Swap it in:
$$= \sum_\tau P(\tau;\theta) \cdot \nabla_\theta \log P(\tau;\theta) \cdot R(\tau)$$Step 4. Now the form $\sum_\tau P(\tau;\theta)(\cdot)$ is just an expectation over trajectories sampled from $\pi_\theta$:
$$\nabla_\theta U(\theta) = \mathbb{E}_{\tau \sim \pi_\theta}\!\left[ \nabla_\theta \log P(\tau;\theta) \cdot R(\tau) \right]$$The log-derivative identity used in Step 3 is just the chain rule applied to the log. You have $\theta \to P(\tau;\theta) \to \log P(\tau;\theta)$, and the chain rule gives:
What we did in Step 3 is use this in the opposite direction: we started with $\nabla_\theta P$ and rewrote it as $P \cdot \nabla_\theta \log P$. Same chain rule, just rearranged. The "normal" chain rule goes $x \to g(\cdot) \to y \to h(\cdot) \to z$ and tells you how to get $\frac{dz}{dx}$ from the parts. Here we're working backwards: we know the derivative of the composed function ($\nabla_\theta \log P$), and we recover the ingredient we need ($\nabla_\theta P$).
The derivation gives us two useful properties.
First, the result is an expectation. Instead of summing over every possible trajectory, we can sample from $\pi_\theta$ and average.
Second, it replaces $P(\tau;\theta)$ with $\log P(\tau;\theta)$. This improves numerical stability because trajectory probability is a product of token-level probabilities:
Each $\pi_\theta(u_t \mid s_t)$ is a probability in $[0, 1]$. You are chaining $H$ of these multiplications together. For a long sequence, that product becomes astronomically small. The kind of number that underflows to zero in floating point before you can do anything useful with it.
Taking the log converts that product into a sum:
Sums of log-probabilities stay in a perfectly manageable range no matter how long the sequence gets. If you have done maximum likelihood estimation before, this is exactly the same principle. We always maximize the log-likelihood precisely because the log turns the product into a sum.
Back to the other half of what the derivation gave us: the expectation. Look at the exact form of the gradient we now want to compute: $\mathbb{E}_{\tau \sim \pi_\theta}[\cdots]$. Computing this expectation exactly means summing over every possible trajectory $\tau$. A trajectory is a sequence of $H$ tokens, and at each step the model can generate any of the $V$ tokens in its vocabulary. The total number of distinct trajectories is:
So we need to approximate. And this is exactly why recovering the expectation form matters: we can now estimate it by Monte Carlo sampling. Instead of summing over all trajectories, we sample $N$ trajectories from $\pi_\theta$ and average. This gives an unbiased estimator of the true gradient:
where $\tau^{(i)} \sim \pi_\theta$ are trajectories rolled out from the current policy.
And the intuition is simple:
Reinforce what worked, punish what didn't.
The model maps a state to a distribution over actions. Decompose the trajectory into state-action pairs:
We can see two buckets. The first is the dynamic model: the probability of transitioning from one state to another given the action taken (the environment, outside our control). The second is the policy: the probability of the model choosing action $u_t$ given the current state.
Distribute the log (product becomes a sum):
$$= \nabla_\theta \left[\,\sum_{t=0}^{H-1} \log P(s_{t+1}^{(i)} \mid s_t^{(i)}, u_t^{(i)}) \;+\; \sum_{t=0}^{H-1} \log \pi_\theta(u_t^{(i)} \mid s_t^{(i)})\right]$$Distribute the gradient inside the sum:
$$= \underbrace{\sum_{t=0}^{H-1} \nabla_\theta \log P(s_{t+1}^{(i)} \mid s_t^{(i)}, u_t^{(i)})}_{\text{Bucket 1: environment}} \;+\; \underbrace{\sum_{t=0}^{H-1} \nabla_\theta \log \pi_\theta(u_t^{(i)} \mid s_t^{(i)})}_{\text{Bucket 2: policy}}$$The environment transition probability is independent of $\theta$, so its gradient is zero:
Substituting this into the estimator gives:
The trajectory reward is:
For each time step $t$, we weight the gradient by the cumulative reward for the whole trajectory. But this breaks a very simple causality principle: an action taken at time $t$ cannot be blamed for rewards collected before $t$. The past is already fixed. Causally, an action can only be blamed for the rewards it provokes in the future, not actions taken in the past.
For step $t$, split the trajectory reward into past and future terms:
It can be shown that in expectation, the past terms contribute nothing to the gradient. They are independent of the action taken at $t$, so they only add variance without adding any useful signal. We can safely drop them. What remains is the reward-to-go:
Plugging this back in, our estimator becomes:
Same idea, less noise. Each gradient step is now weighted only by what actually follows from the action, not by the whole episode history.
Suppose one game-playing agent hides in a corner and collects a modest positive reward, while another tries promising but unfinished strategies and sometimes dies early. Raw return rewards the first agent even if its behavior has little long-term value.
The estimator needs a reference point that distinguishes an expected outcome from one that is unusually good for the current state.
The fix is to subtract a baseline $b(s_t)$ from the reward-to-go. The mathematical key: as long as the baseline only depends on the state and not the specific action taken, subtracting it does not introduce any bias into the gradient in expectation. So we can safely center our estimator:
The natural choice is the expected reward-to-go from state $s_t^{(i)}$: what would a typical trajectory collect from here? This is exactly the definition of the value function:
Computing this expectation exactly would require considering every future trajectory from every state. The branching cost is infeasible at scale.
Instead, a value function parameterized by $\phi$ approximates the expected reward-to-go:
In practice, we simply clone the LLM weights and add a small projection head on top to predict a scalar. This gives us a second model with its own parameters $\phi$, called the critic. The policy being optimized is called the actor. Together they form the Actor-Critic family of algorithms.
Actor $\pi_\theta$: the policy network. Looks at a state and outputs a probability distribution over actions. Updated to maximize the Advantage.
Critic $V_\phi^\pi$: the value network. Looks at a state and outputs a single scalar: the predicted expected reward-to-go. Updated by minimizing prediction error (MSE).
The training objective is intuitive: we want $V_\phi^\pi$ to predict the actual cumulative reward as accurately as possible, so we minimize the squared error between the value estimate and the observed reward-to-go across all sampled trajectories:
With the Critic in place, we can now name the quantity we are actually computing. The difference between the observed reward-to-go and the Critic's prediction is called the Advantage: how much better or worse than average was this specific action?
This directly gives us the reinforce-what-worked signal:
The critic predicts an expectation, not the exact return of each stochastic rollout.
Training it with mean squared error over many sampled returns makes the optimal prediction their conditional mean.
That conditional mean is the baseline we need. Subtracting it lets the actor learn from how a sampled return differs from the expected return for that state.
At initialization, the critic's predictions are noisy.
As long as the baseline depends on state and not on the sampled action, subtracting it does not bias the actor's expected gradient. It changes the variance.
An inaccurate critic can add variance, while a better one usually reduces it and centers the learning signal. This does not guarantee smooth or rapid convergence, but it can make updates much less noisy.
There is an interesting circularity here: the Critic is trained to predict its own future value (bootstrapping). The only thing keeping it grounded is the real reward $r_t$ anchoring each update step. Without it, the network would collapse to predicting 0 everywhere, which is always consistent with itself.
The reward-to-go $\sum_{k=t}^{H-1} r(s_k^{(i)}, u_k^{(i)})$ comes from one sampled trajectory. It is unbiased but has high variance because each rollout follows a different path.
A raw single-sample return has low bias and high variance. Relying on the critic lowers variance but introduces bias from its approximation errors. GAE balances the two.
To set up that balance properly, we first need to introduce one more function.
So far, $V^\pi(s)$ tells us the expected reward-to-go from a state, averaged over all actions the policy might take. But we can condition more precisely: the action-state value function, or Q-function, fixes both the state and the action taken at that step, then averages over everything that follows:
Action-state value function / Q-function
$$Q^\pi(s, u) \;=\; \mathbb{E}\!\left[r_t + r_{t+1} + r_{t+2} + \cdots \;\middle|\; s_t = s,\; u_t = u\right]$$The relationship between $Q^\pi$ and $V^\pi$ is clean: the value of a state is just the Q-value averaged over the policy's action distribution.
Now look at the Q-function more carefully. It is the expectation of the cumulative reward from $t$ onward, attributing uniform credit to every future state-action pair based on the single action taken at time $t$. This is a source of variance: we are assigning equal blame to events that are increasingly less causally connected to $u_t$. Variance is entering the expectation along the temporal axis.
We can reasonably address this by introducing a temporal weighting: actions taken at time $t$ have a stronger causal impact on the near future than on the distant future, so we discount future rewards by a factor $\gamma \in (0, 1]$ raised to the number of steps elapsed. This gives us the discounted Q-function:
Action-state value function / Q-function (discounted)
$$Q^{\pi,\gamma}(s, u) \;=\; \mathbb{E}\!\left[r_t + \gamma r_{t+1} + \gamma^2 r_{t+2} + \cdots \;\middle|\; s_t = s,\; u_t = u\right]$$Substitute $Q^{\pi,\gamma}$ for reward-to-go and $V_\phi^{\pi,\gamma}$ for the critic baseline. Their difference is the advantage.
Back to the Advantage definition. To actually compute $\hat{A}_t(s_t, u_t) = Q^{\pi,\gamma}(s_t, u_t) - V^{\pi,\gamma}(s_t)$, we need an estimate of $Q^{\pi,\gamma}$. The simplest approach: run the trajectory and use the actual discounted cumulative reward from $t$ onward as a single-sample estimate. This is Monte Carlo.
Monte-Carlo
$$\hat{A}_t^{\text{MC}} = -V(s_t) + r_t + \gamma r_{t+1} + \gamma^2 r_{t+2} + \cdots + \gamma^{H-1} r_{H-1}$$Accurate, low bias. But still high variance: a single trajectory is a single random draw through the environment, and the reward-to-go is noisy.
Now look at the terms from $t+1$ onward: $r_{t+1} + \gamma r_{t+2} + \cdots + \gamma^{H-2} r_{H-1}$. Factor $\gamma$ out front:
That parenthesised tail is the discounted return from state $s_{t+1}$ onward, which is exactly what $V(s_{t+1})$ is trained to estimate. Substitute it in:
One real reward step, then the Critic absorbs the rest. We traded the high variance of a full noisy rollout for the low variance of the Critic's smooth average, at the cost of carrying whatever bias the Critic holds. And nothing is special about stopping after one step: we can take $n$ actual reward steps and only then hand off to the Critic for what remains.
Monte-Carlo
$$\hat{A}_t^{\text{MC}} = -V(s_t) + r_t + \gamma r_{t+1} + \gamma^2 r_{t+2} + \cdots + \gamma^{H-1} r_{H-1}$$Temporal-Difference
$$\hat{A}_t^{(1)} = -V(s_t) + r_t + \gamma V(s_{t+1})$$ $$\hat{A}_t^{(2)} = -V(s_t) + r_t + \gamma r_{t+1} + \gamma^2 V(s_{t+2})$$ $$\hat{A}_t^{(3)} = -V(s_t) + r_t + \gamma r_{t+1} + \gamma^2 r_{t+2} + \gamma^3 V(s_{t+3})$$Async Advantage Actor-Critic (A3C)
Each additional real step reduces variance (we trust the environment over the Critic for one more timestep) and also reduces bias (we rely less on the Critic's approximation error). More steps: less bias, more variance. Fewer steps: more bias, less variance. The question becomes: which $n$?
Instead of committing to a single $n$, we can generalise the temporal-difference idea into a single framework: Generalised Advantage Estimation. With one parameter $\lambda$, all $n$-step estimates get an exponentially decaying weight, and the bias-variance tradeoff becomes a continuous knob rather than a discrete choice.
Generalised Advantage Estimation (GAE) $0 < \lambda < 1$
When $\lambda \to 0$, the series collapses to the 1-step TD estimate: fully trust the Critic, low variance, biased. When $\lambda \to 1$, all $n$-step estimates contribute equally and we recover Monte Carlo: zero bias, high variance. At intermediate values, $\lambda$ smoothly interpolates. In practice, $\lambda \approx 0.95$ is the default; it gives enough variance reduction to stabilise training without letting Critic errors dominate.
One might wonder how the Critic training still works in this expression. We are now training the Critic to match the immediate reward plus its own guess from time $t+1$ onward. This is no longer the full sequence of real rewards. The loss becomes:
$\text{sg}[\cdot]$ = stop-gradient (detach from computation graph)
The main difference from before is that instead of waiting for the end of the trajectory, the Critic is updated step by step. It uses the actual observed reward plus its own expectation of all future rewards as the target. This creates a low-variance but high-bias setup, and it makes optimisation trickier: the Critic now relies on its own assumptions about the future, constantly self-correcting as the trajectory rolls out.
Finally, a crucial implementation detail. To keep gradients stable and avoid an infinite computational loop (the Critic trying to differentiate through its own future prediction), the target $V_\phi^\pi(s_{t+1})$ is detached from the computation graph. We freeze it and use only its hard numerical value as the training target for the current step. This is the $\text{sg}[\cdot]$ in the loss above.
The full Actor-Critic training loop is now three clean steps that repeat until convergence.
Initialize policy network $\pi_{\theta_0}$ and value network $V_{\phi_0}$
Collect rollouts $\bigl\{s_t^{(k)},\, u_t^{(k)},\, s_{t+1}^{(k)},\, r_t^{(k)}\bigr\}_{t=0}^{H-1}$ and compute $\hat{Q}^{\pi,\gamma}(s_t^{(k)}, u_t^{(k)})$
Update value network
$$\phi_{i+1} \leftarrow \min_{\phi} \sum_{\{s_t,\, u_t,\, s_{t+1},\, r_t\}} \left\|\hat{Q}^{\pi,\gamma}(s_t^{(k)}, u_t^{(k)}) - V_\phi^{\pi,\gamma}(s_t^{(k)})\right\|_2^2 + \lambda\|\phi - \phi_i\|_2^2$$Update policy network
$$\theta_{i+1} \leftarrow \theta_i + \alpha\,\frac{1}{m}\sum_{k=1}^{m}\sum_{t=0}^{H-1} \nabla_{\theta_i}\!\log\pi_{\theta_i}(u_t^{(k)} \mid s_t^{(k)}) \left(\hat{Q}^{\pi,\gamma}(s_t^{(k)}, u_t^{(k)}) - V_{\phi_i}^{\pi,\gamma}(s_t^{(k)})\right)$$We now have a working framework for updating policy parameters. But it still has a critical dependence on step size. Too small, and training is painfully slow. Too large, and the policy degenerates: a single bad update can push the parameters into a region from which recovery is very hard.
The deeper issue is that policy updates are completely unbounded. The gradient step has no ceiling. If a single rollout happens to return an unusually high or low reward (one lucky or unlucky Monte Carlo draw), the resulting gradient can be enormous, and the update will trash the policy in a single step. There is nothing in the algorithm preventing this.
The framework is also inefficient: collect an episode, make one update, discard the data, and repeat. The gradient estimator explains why.
Replacing the Q minus V term with the Advantage $\hat{A}_t$, and as $m \to \infty$, this sample average converges to an expectation:
The distribution we are summing over is $P_\theta(s_t, u_t)$: the joint probability of visiting state $s_t$ and taking action $u_t$ under the current policy $\theta$. The sampling distribution and the policy being optimised are the same object. This is on-policy learning: you cannot reuse data collected under any other policy without corrupting the gradient estimate.
Every time $\theta$ changes, the old trajectories become samples from the wrong distribution. One update, one episode, discard. This is what makes vanilla Actor-Critic so expensive.
To break the on-policy constraint, we can use a classical statistical trick: Importance Sampling.
The general problem: we want to estimate the expected value of some function $f(x)$ under a distribution $P(x)$, but we want to draw our samples from a completely different distribution $Q(x)$. To correct for the mismatch, we apply an importance weight to each sample: the ratio $\frac{P(x)}{Q(x)}$. If $Q$ rarely visits a region where $P$ is large, those rare samples get a high weight to compensate. The expectation is preserved.
The reason it works is a simple algebraic trick. Start from the definition of expectation as an integral, multiply by the dummy term $\frac{Q(x)}{Q(x)} = 1$, then rearrange so $Q(x)$ sits out front. The resulting integral is exactly an expectation under $Q$.
In practice the integral is replaced by a sample average. The two equivalent estimators side by side:
Sampling directly from P
$$\mathbb{E}_{x \sim P(x)}[f(x)] \approx \frac{1}{n}\sum_{i=1}^{n} f(x^{(i)}), \qquad x^{(i)} \sim P(x)$$Sampling from Q (Importance Sampling)
$$\mathbb{E}_{x \sim P(x)}[f(x)] \approx \frac{1}{n}\sum_{i=1}^{n} \frac{P(x^{(i)})}{Q(x^{(i)})} \, f(x^{(i)}), \qquad x^{(i)} \sim Q(x)$$The visualisation below shows this. $P$ is fixed (blue), $Q$ is the sampling distribution (green, shift it with the slider), and $f(x) = x$ is the linear target (yellow). Samples drawn from $Q$ appear as dots below the axis; dot size is proportional to the importance weight $P(x^{(i)})/Q(x^{(i)})$. Despite all samples coming from the wrong distribution, the Importance-Sampling-weighted average still converges to the true expectation under $P$.
There is one important caveat. When $Q$ is very far from $P$, most samples land in a region where $P$ is tiny, so they receive near-zero weights (small dots). The rare samples that do land near $P$'s peak receive enormous weights (large dots) to compensate. The estimate stays unbiased in theory, but a single outlier sample can dominate the sum and variance explodes. Slide $Q$ to an extreme to see this directly. This is why Importance Sampling only works reliably when $P$ and $Q$ are not too far apart.
We now apply the exact same trick to our on-policy gradient. The sampling distribution is $P_\theta(s_t, u_t)$, the old policy is our new $Q$, and the function being evaluated is $\nabla_\theta \log \pi_\theta \cdot \hat{A}_t$. Multiply by $\frac{P_{\theta_\text{old}}}{P_{\theta_\text{old}}} = 1$, swap, and the expectation shifts to samples from $\pi_\text{old}$.
Now factorize the joint state-action probability: $P_\theta(s_t, u_t) = \pi_\theta(u_t \mid s_t)\, P_\theta(s_t)$, and the same for old. The Importance-Sampling ratio splits into a policy ratio and a state-marginal ratio:
Key assumption (remember this): $P_\theta(s_t) \approx P_{\theta_\text{old}}(s_t)$. If the policy does not change too much, the state visitation distributions stay close and their ratio is approximately 1. Under this assumption, the state marginals cancel. This assumption will come back.
After cancellation, the state marginals vanish and we are left with the surrogate objective:
We can verify that differentiating $U(\theta)$ recovers the correct gradient. Push $\nabla_\theta$ inside the expectation (the sampling distribution $\pi_\text{old}$ does not depend on $\theta$), then differentiate only the $\pi_\theta$ in the numerator:
We just proved something important. We can use this surrogate objective to run gradient-based optimisation on our policy model using data collected from an older version of that same model. The rollout cost is paid once; $\theta$ can then be updated many times against that fixed batch before collecting new data. The on-policy bottleneck is broken.
This requires the old and new policies to stay close, so $\frac{\pi_\theta}{\pi_{\theta_\text{old}}}$ remains near 1. When they diverge, two failure modes appear.
1. Importance-Sampling variance explosion. When the new and old policies differ substantially, the ratio $\frac{\pi_\theta}{\pi_{\theta_\text{old}}}$ can be very large for some state-action pairs and near zero for others. Exactly as we saw with $P$ and $Q$ in the visualisation: a few samples dominate the sum, the gradient estimate becomes extremely noisy, and training destabilises.
2. Advantage function staleness. $\hat{A}_t^{\pi_\text{old}}$ was estimated under the old policy. If the new policy has moved far from $\pi_\text{old}$, the state-action pairs it visits are drawn from a different distribution than the one used to compute the Advantages. The signal telling the policy what is better or worse than average is now based on a world the policy no longer inhabits. The gradient is pointing in the wrong direction.
Both problems have the same root: we need to keep $\theta$ close to $\theta_\text{old}$ during each update. We need a principled way to enforce that constraint.
TRPO turns “do not drift too far” into a constraint. It maximizes the surrogate objective $U(\theta)$ while limiting the difference from the old policy.
The measure of difference between two policies is the KL divergence, averaged over the states the old policy visits:
The full TRPO optimization problem is then:
In parameter space, the KL constraint defines a trust region around $\pi_{\theta_\text{old}}$. TRPO finds the point inside that region that maximizes the surrogate objective.
Solving this exactly requires the local curvature of the KL around the current policy, captured by the Fisher Information Matrix.
For a network with $N$ parameters, that matrix is $N \times N$. For a 7B-parameter LLM, that is roughly $5 \times 10^{19}$ entries, about the number of grains of sand on Earth. There are clever ways around it (Schulman's original paper avoids ever building the matrix explicitly by using conjugate gradient on Fisher-vector products), but even with the tricks, TRPO is slow, fragile, and never really made the jump to large networks.
TRPO provides a monotonic-improvement guarantee under its assumptions, but its second-order machinery is too expensive for large models. PPO keeps the trust-region idea with a cheaper first-order objective.
If you are curious about the gory details (the full FIM memory table, the conjugate gradient trick, why TRPO never scales), see the appendix at the end of the post.
TRPO's exact constrained solution relies on expensive second-order optimization. PPO replaces the hard constraint with a first-order objective.
The standard trick in constrained optimization is the Lagrangian relaxation: instead of enforcing a hard constraint, fold it into the objective as a penalty weighted by a coefficient $\beta$. A constraint violation now costs you objective value rather than being forbidden outright. The optimization problem becomes unconstrained, and you can solve it with plain gradient descent.
TRPO: constrained
$$\max_{\theta}\; \mathbb{E}_{\pi_\text{old}}\!\left[\frac{\pi_\theta(u_t \!\mid\! s_t)}{\pi_{\theta_\text{old}}(u_t \!\mid\! s_t)}\, \hat{A}_t^{\pi_\text{old}}\right]$$ $$\text{subject to}\quad \mathbb{E}_{\pi_\text{old}}\!\left[\mathrm{KL}\!\left(\pi_\theta \;\big\|\; \pi_{\theta_\text{old}}\right)\right] \leq \varepsilon$$ $$\Downarrow \quad \text{Lagrangian relaxation} \quad \times\; \beta$$PPO-penalty: unconstrained
$$\max_{\theta}\; \mathbb{E}_{(s_t,\, u_t) \sim \pi_\text{old}}\!\left[\frac{\pi_\theta(u_t \!\mid\! s_t)}{\pi_{\theta_\text{old}}(u_t \!\mid\! s_t)}\, \hat{A}_t^{\pi_\text{old}}(s_t, u_t)\right] - \beta\left(\mathbb{E}_{(s_t,\, u_t) \sim \pi_\text{old}}\!\left[\mathrm{KL}\!\left(\pi_\theta(u_t \!\mid\! s_t) \;\big\|\; \pi_{\theta_\text{old}}(u_t \!\mid\! s_t)\right)\right]\right)$$This is Proximal Policy Optimization, or PPO, also known as PPO-penalty. The hard constraint is gone. The second-order solver is gone. The conjugate gradient loop is gone. What remains is a single objective you can differentiate and optimize with Adam in one line. The parameter $\beta$ controls how heavily you penalise KL divergence from the old policy: large $\beta$ keeps the update cautious, small $\beta$ lets it move more freely.
This introduces a sensitive hyperparameter. If $\beta$ is too large, the policy barely moves. If it is too small, the policy drifts and importance-sampling ratios become unstable.
One proposed fix is to make $\beta$ adaptive: measure the actual KL divergence after each update, and adjust $\beta$ accordingly.
Adaptive KL Penalty
$$\text{If}\;\; \mathbb{E}_{(s_t,\, u_t) \sim \pi_\text{old}}\!\left[\mathrm{KL}\!\left(\pi_\theta \;\|\; \pi_{\theta_\text{old}}\right)\right] > \mathrm{KL}_\text{max}, \quad \text{then}\;\; \beta \leftarrow \beta \times 2$$ $$\text{If}\;\; \mathbb{E}_{(s_t,\, u_t) \sim \pi_\text{old}}\!\left[\mathrm{KL}\!\left(\pi_\theta \;\|\; \pi_{\theta_\text{old}}\right)\right] < \mathrm{KL}_\text{min}, \quad \text{then}\;\; \beta \leftarrow \beta \,/\, 2$$The adaptive version introduces three settings: $\beta$, $\mathrm{KL}_\text{max}$, and $\mathrm{KL}_\text{min}$. It is also reactive, correcting a KL violation only after an update.
The objective should raise the probability of advantageous actions, lower it for bad ones, and keep the new policy close to the old one. A KL penalty discourages drift but does not impose a hard limit on a single update.
PPO-clip adds that limit directly to the probability ratio $r_t(\theta)$:
The ratio $r_t(\theta) = 1$ when the new policy assigns the same probability as the old one: that is the starting point of every update, marked by the red dot in each plot below. The objective is linear in $r_t(\theta)$. Looking at the two cases:
When the advantage is positive (left), gradient ascent pushes $r_t$ to the right, increasing the probability of taking this action. Keep going, keep gaining reward. There is no ceiling. When the advantage is negative (right), the gradient pushes $r_t$ toward zero. There is no floor. In both cases the gradient is constant: the slope never changes no matter how far $r_t$ has already moved from 1. A single sample with a large advantage can send the policy ratio anywhere, and nothing slows it down.
The wall we want: once $r_t(\theta)$ has moved far enough from 1, the gradient becomes exactly zero. Zero, not "very small". Past that boundary, no sample, however extreme, can push the policy further in that direction.
The clip achieves exactly that. Instead of multiplying directly by $r_t(\theta)$, we cap the ratio to $[1-\varepsilon,\, 1+\varepsilon]$ before multiplying. Past either boundary, the clipped value is constant. A constant times $\hat{A}_t$ is a constant, and a constant has zero gradient. Any sample extreme enough to push the ratio outside the trust region is treated as noise from luck or bad luck, and its gradient is killed entirely.
The same two cases, now with walls on both sides:
Inside the trust region $[1-\varepsilon,\, 1+\varepsilon]$, the gradient behaves exactly as before: the policy learns freely. Outside it, the gradient is zero. This is the wall. No hyperparameter to tune beyond $\varepsilon$, no adaptive controller, no second-order math. Just a hard limit on how far the probability ratio can move in a single update.
But the clip alone has a problem. Consider the left plot with $\hat{A} > 0$: a brilliant action was taken, but the current model is highly pessimistic about it and assigns it a probability ratio of $r_t = 0.5$. That puts us on the far left side of the graph, deep in the flat region: zero gradient. The optimizer sees nothing to learn here. A brilliant move is silently discarded, and the model stays pessimistic. The same issue appears on the right side for $\hat{A} < 0$: if the model was too optimistic about a bad action ($r_t$ very large), it falls off the right edge of the clipped graph and again learns nothing.
We only need the wall on the side that could trash the policy. For $\hat{A} > 0$, the dangerous side is the right (the model over-reinforces a good action, drifting too far). The left side is perfectly safe: learning more from a pessimistic model only corrects the policy in the right direction. For $\hat{A} < 0$, it is the opposite: the dangerous side is the left (collapsing the probability to zero), while the right side safely unlearns an over-represented bad action.
We want the wall only where we need it. Take the minimum of the unclipped surrogate and the clipped surrogate at every point. The $\min$ is pessimistic by construction: it selects whichever objective is lower. On the safe side it picks the unclipped value, preserving the learning signal. On the dangerous side it picks the clipped value, enforcing the wall.
On the left plot, the yellow line rises freely from the origin: a pessimistic model still gets its full learning signal. It only flattens once the ratio exceeds $1+\varepsilon$, blocking over-reinforcement. On the right plot, the yellow line stays flat on the left to prevent probability collapse, then follows the unclipped slope on the right: an over-optimistic model keeps unlearning a bad action past $1+\varepsilon$ without hitting any wall. The $\min$ places the wall exactly where the policy could be trashed, and nowhere else. Awesome, you just invented PPO-clip!
The full pipeline starts with a prompt $s_0$. The old policy $\pi_{\theta_\text{old}}$ generates a trajectory $\tau$, a frozen reward model scores it, and a value network $V_\phi^\pi$ estimates the return from each state. GAE produces per-step advantages, which feed the PPO-clip update:
In PPO, you might have noticed a potential issue. We have created a hard wall, such that a single update cannot completely trash the model. But have you considered what happens after enough training steps? Each update stays within the wall, but those small allowed steps accumulate. The policy ends up drifting far from where it started.
Clipping limits each local update but does not control drift over the full training run. RLHF therefore adds a KL penalty against a fixed reference policy. Here KL is an anchor, not a hard per-step constraint.
Historically, PPO was designed for agents trained to play video games, where we actually wanted the agent to explore freely. Drifting far from the original policy was sometimes the whole point. Language modeling is different. If our model drifts too far from where it started, it can quite literally forget how to speak.
Step 1: Supervised Fine-Tuning (SFT): Train a base language model on high-quality demonstration data to produce a capable assistant baseline πSFT.
Step 2: Reward Model Training: Collect human rankings of different model outputs. Train a separate network rθ(x, y) to score how much a human would prefer response y to prompt x.
Step 3: RL Optimization: Use PPO to optimize πSFT against rθ(x, y), with a KL penalty to limit drift.
The fix is to add a KL term directly to the PPO objective, anchored not to the previous checkpoint (as in PPO-KL) but to the frozen original policy πref, the SFT model from Step 1.
RLHF objective: PPO-clip with a global KL anchor
$$\mathbb{E}_{\pi_\text{old}}\!\left[\min\!\left(r_t(\theta)\,\hat{A}_t,\;\mathrm{clip}\!\left(r_t(\theta),1-\varepsilon,1+\varepsilon\right)\hat{A}_t\right)\right] - \beta\,D_{\text{KL}}\!\left(\pi_\theta \;\big\|\; \pi_{\text{ref}}\right)$$The KL divergence between two distributions over the full vocabulary is defined as:
This requires summing over the entire vocabulary V at every token step. For a 100k-token vocabulary, computing this exactly is expensive.
In practice, we avoid the full sum by computing the KL contribution only for the sampled token at each step. This is the token-level penalty, often called the k1 estimator:
Practical implementation: token-level KL penalty (k1)
$$R_t = r_t - \beta\!\left(\log \pi_\theta(y_t \mid x, y_{\lt t}) - \log \pi_{\text{ref}}(y_t \mid x, y_{\lt t})\right)$$This looks like it only evaluates the sampled token yt. But in expectation under πθ, it recovers the exact KL divergence. It is an unbiased estimator of the true KL gradient.
The two mechanisms have different jobs. Clipping limits a single update, while KL limits cumulative drift from the frozen reference.
If you are still here, you are motivated. It has been a long road, and hopefully each step felt like a natural fix to what came before.
I have bad news. In GRPO, the authors looked at everything we just built and asked a different question: what if we threw a big chunk of it out?
In fairness, the biggest bottleneck in PPO is the Critic. In total, PPO keeps four copies of full model weights in memory at once: the Actor being trained, the Critic providing baselines, the frozen Reward model, and the frozen Reference model for the KL anchor. For large language models, that is not a trivial resource budget.
GRPO takes a different route to reduce PPO's memory cost. Start with the full objective, then examine each term.
where
$$\hat{A}_t^{(i)} = \hat{A}^{(i)} = \frac{r^{(i)} - \mathrm{mean}\!\left(\{r^{(i)}\}\right)}{\mathrm{std}\!\left(\{r^{(i)}\}\right)} \qquad\qquad r_t(\theta) = \frac{\pi_\theta(u_t \!\mid\! s_t)}{\pi_{\theta_\text{old}}(u_t \!\mid\! s_t)}$$$\mathbb{E}_{\tau^{(i)} \sim \pi_\text{old}}$ is the same off-policy setup as before: trajectories are collected from an old snapshot of the policy, and the PPO-clip ratio handles the distribution mismatch. Nothing new there.
For each prompt $s_0$, the policy generates not one but $G$ full trajectories. That is the core of GRPO. Instead of one answer, ask the model the same question $G$ times and collect all the responses.
In PPO, the advantage came from GAE: a weighted mix of n-step returns, requiring a Critic to estimate $V(s_t)$ at every step. GRPO replaces this entirely with a group statistic. The reward of each response is normalized relative to the other $G-1$ responses in the group:
This is a z-score within the group. If your response scored higher than average, you get a positive advantage. If lower, a negative one. The mean of the group is the baseline, taking the place of the Critic's $V_\phi^\pi(s_t)$. The intuition is direct: the Critic was always trying to estimate the expected reward from this state. With $G$ rollouts of the same prompt, the group mean is that estimate, computed by Monte Carlo sampling rather than a learned neural network.
Notice that $\hat{A}_t^{(i)} = \hat{A}^{(i)}$: the same scalar is applied to every token in response $i$, regardless of where in the response that token appears. In PPO with GAE, each timestep $t$ had its own advantage, accounting for how much credit belonged to the actions taken at that specific step. GRPO has no per-step signal. A winning response gets a uniform positive push across all its tokens; a losing response gets a uniform negative push. Credit assignment within a trajectory is left entirely implicit.
There is one more component in the equation: $-\beta\, D_\text{KL}(\pi_\theta \| \pi_\text{ref})$. You already know this term from the RLHF section. The KL rubber band anchors the policy to a frozen reference model $\pi_\text{ref}$, preventing it from drifting away from the language distribution it started with. GRPO carries this anchor forward unchanged.
The wall and the rubber band now work together without conflict. The clip wall disarms any dangerously large gradient before it can overpower the KL penalty: in the flat regions of the clip, the gradient is exactly zero, so it cannot fight anything. The two mechanisms operate at different timescales. The wall handles short-term, per-step absolute safety. The KL handles long-term cumulative drift. The $\beta$ hyperparameter still needs tuning, but its job is now narrower: not preventing policy collapse on a single bad batch (the wall handles that), but ensuring the model does not slowly forget it is a language model.
Computing the true KL divergence requires summing over every token in the vocabulary at every position, which is expensive. In practice, it is approximated per token using an estimator that only requires evaluating the two policies at the token actually generated:
GRPO gives up per-step credit assignment in exchange for eliminating the Critic. For language model tasks where an entire response is evaluated as a unit (did the reasoning lead to the correct answer?), this is often a reasonable trade. The group baseline is noisier than a learned value function, and the uniform per-trajectory advantage is coarser than GAE. But the system is dramatically simpler, faster, and cheaper to run, and in practice it trains reasoning models effectively.
GAE: temporal and local. GAE computes a separate advantage $\hat{A}_t$ for every single token at every step $t$. It knows that token 5 might have been a brilliant word choice while token 90 was a poor one, and it assigns credit accordingly. The Critic's value estimates provide a per-step baseline, and the exponential weighting over n-step returns gives a fine-grained, local signal all the way down the trajectory.
GRPO: global and sparse. GRPO produces a single scalar $\hat{A}^{(i)}$ for the entire trajectory. Every token in the response, from the first word to the last, receives that same number as its advantage. There is no temporal signal, no local credit, no distinction between a token that was pivotal and one that was filler. The only information available is: did this response score above or below the group average?
This is a significant assumption. GRPO implicitly relies on the policy model being capable of distributing a global scalar signal across its attention: reinforcing the specific tokens that actually deserved credit and suppressing those that did not, all without being told which is which. For large, well-pretrained language models this often works in practice. But it is a real loss of precision compared to GAE, and worth keeping in mind clearly before reading what comes next.
What changes. TRPO provides a monotonic-improvement guarantee under its trust-region assumptions, and PPO approximates that constraint. GRPO has no comparable guarantee. Its group mean is a Monte Carlo baseline, and one trajectory-level advantage is applied to every token. Standard-deviation and length normalization are engineering choices rather than consequences of the optimization theory.
GRPO trades formal guarantees and fine-grained credit assignment for a system that scales to large models. Its success depends on the model, reward quality, and training regime; the same trade may fail with small networks or sparse, noisy rewards.
Dr.GRPO targets two terms in this objective: standard-deviation normalization and per-token averaging over $H^{(i)}$. In reasoning-model training, they introduce difficulty and length biases.
GRPO's two normalization terms create length and difficulty biases.
Notice the division by the trajectory length? This means we are actively diluting the reward per token. When the model gets it right (\(\hat{A}>0\)), every additional, non-necessary token increases the denominator and dilutes the reward, so the model learns it is safer to be very direct and short. On the other hand, if the model is wrong (\(\hat{A}<0\)), a long trajectory spreads the negative penalty across so many tokens that the punishment per token becomes microscopically small. The model learns to ramble to avoid concentrated punishment.
We are calculating the z-score over a batch of trajectories for the same prompt. When the question is simple, it is highly likely that over \(G\) trajectories we have very low reward variance. Oppositely, in complex scenarios, the model might not always be correct, so the variance increases. By dividing by the standard deviation over the group, we inadvertently reward simpler questions. Dividing by a tiny variance (easy question) causes the final update scalar to explode, while dividing by a large variance (hard question) shrinks the update.
Looking at the GRPO objective, we can now see exactly where both biases are hiding:
Dr.GRPO removes both normalization terms.
Dr.GRPO
$$\mathbb{E}_{\tau^{(i)} \sim \pi_\text{old}}\!\left[\frac{1}{G}\sum_{i=1}^{G}\sum_{t=0}^{H^{(i)}-1}\min\!\left(\, r_t(\theta)\,\hat{A}^{(i)},\;\; \mathrm{clip}\!\left(r_t(\theta), 1-\varepsilon, 1+\varepsilon\right)\hat{A}^{(i)}\,\right) - \beta\, D_\text{KL}\!\left(\pi_\theta \,\big\|\, \pi_\text{ref}\right)\right]$$ $$\hat{A}^{(i)} = r^{(i)} - \text{mean}\!\left(\{r^{(j)}\}_{j=1}^{G}\right)$$A four-rollout example shows how the terms change each token's gradient signal.
Take a prompt with $G = 4$ rollouts. Rewards are binary: 1 for the right final answer, 0 otherwise.
| Rollout | $H^{(i)}$ | $r^{(i)}$ |
|---|---|---|
| A: “12 × 13 = 156” ✓ short correct | 7 | 1 |
| B: “Let me compute. 12×13 = 12×10 + 12×3 = 120 + 36 = 156” ✓ long correct | 22 | 1 |
| C: “12 × 13 = 144” ✗ short wrong | 7 | 0 |
| D: “Hmm... 12 × 13... 12×10 is 120, plus 12×3 which is... 30? So 150? Actually let me try... 144.” ✗ long wrong | 35 | 0 |
Group statistics: $\mathrm{mean}(r) = 0.5$, $\mathrm{std}(r) = 0.5$. Now compute the per-token signal each token of each rollout actually sees, side by side.
| i | $\hat{A}^{(i)}$ | $\hat{A}^{(i)}/H^{(i)}$ |
|---|---|---|
| A | +1.000 | +0.143 |
| B | +1.000 | +0.045 |
| C | −1.000 | −0.143 |
| D | −1.000 | −0.029 |
| i | $\hat{A}^{(i)}$ | per token |
|---|---|---|
| A | +0.500 | +0.500 |
| B | +0.500 | +0.500 |
| C | −0.500 | −0.500 |
| D | −0.500 | −0.500 |
Under GRPO, rollout A (short and correct) gets $+0.143$ per token, while rollout B (long and correct) gets $+0.045$. The longer correct answer receives a three times weaker signal per token despite the same final reward.
For incorrect answers, rollout C gets $-0.143$ per token and the longer rollout D gets only $-0.029$. Length normalization gives the longer wrong answer a five times weaker penalty per token.
Under Dr.GRPO, every token in every rollout gets the same magnitude of signal regardless of length. Short correct and long correct are reinforced equally per token. Short wrong and long wrong are punished equally per token. The model is judged on whether it got the answer right, not on how many tokens it spent doing so. That is the entire fix.
REINFORCE: all rewards positive, no baseline → noisy, slow convergence
↓ Fix: subtract a baseline
Baseline + Monte Carlo: can't estimate $V(s_t)$ cheaply per state → need a second network
↓ Fix: train a Critic
Actor-Critic: full rollout noisy, pure bootstrapping biased → need a blend
↓ Fix: interpolate with GAE ($\lambda$)
Actor-Critic + GAE: policy can still take a catastrophic step → need a constraint
↓ Fix: hard KL constraint (TRPO)
TRPO: KL constraint requires $O(N^2)$ Fisher Matrix → computationally impossible
↓ Fix: approximate constraint with a clip (PPO)
PPO: clip is a local wall; across thousands of steps the policy drifts, and a language model forgets how to speak
↓ Fix: add a global KL anchor against a frozen reference model (RLHF)
RLHF: requires a Critic, a Reward Model, and a frozen reference all in memory; unstable online RL loop
↓ Fix: replace Critic with group statistics (GRPO)
GRPO: standard deviation normalisation penalises hard prompts, length normalisation penalises long reasoning
↓ Fix: remove both (Dr.GRPO)
Dr.GRPO: removes both normalization terms
Starting from expected reward, each method addresses a specific source of variance, instability, cost, or bias. Dr.GRPO ends this particular chain by removing the two normalization terms from GRPO.
TRPO uses a principled but expensive trust-region update. PPO replaces it with a practical first-order approximation but keeps a critic. GRPO removes the critic for large-model training, accepting coarser credit assignment in return.
Later methods such as GRPO-v2 and DAPO continue this line, while approaches such as DPO avoid the online RL loop. The useful reading strategy is the same: identify the limitation a new objective targets and the trade it makes.
This appendix explains why TRPO's second-order machinery does not scale to large networks.
Recall the setup. TRPO wants to maximize the surrogate objective subject to a KL constraint of size $\varepsilon$. Solving the constrained optimization exactly requires knowing the local curvature of the KL around the current policy. That curvature is the Fisher Information Matrix, defined as the Hessian of the KL divergence evaluated at $\theta_\text{old}$:
For a policy network with $N$ parameters, the FIM is an $N \times N$ matrix.
| Model size | Parameters $N$ | FIM entries ($N^2$) | Memory (float32) |
|---|---|---|---|
| Tiny network | 1 million | 1 trillion | 4 TB |
| Small network | 10 million | 100 trillion | 400 TB |
| Small LLM | 7 billion | 49 quintillion | physically impossible |
At language-model scale, the full matrix cannot be stored, inverted, or computed directly.
The original TRPO paper avoids building the full FIM. TRPO needs to solve the linear system $F^{-1}g$, where $g$ is the surrogate-objective gradient. Conjugate gradient can solve it without materializing $F$, provided it can compute $F \cdot v$ for any vector $v$.
A Fisher-vector product $F \cdot v$ can be computed in $O(N)$ time with two backward passes. Each conjugate-gradient iteration therefore avoids the full matrix but still adds two backpropagations.
Because the price for each policy update is no longer a single gradient step. It is:
This is manageable for small networks. For a multi-billion-parameter language model, the repeated backpropagations, line search, and numerical-stability work make each update too expensive.
TRPO's trust-region framing provides a useful way to control policy updates, but its second-order optimizer does not fit modern language-model training. PPO keeps the same intuition in a cheaper first-order objective.