Policy Gradient Review 2026: From One Answer

FWIW, there are numerous tutorials, videos, and blog posts introducing policy gradient. However, I often find them a bit hard to understand and follow because they become abstract too quickly. As an LLM guy, I want to work through a derivation with the right amount of intuition and help people like me better understand this fundamental idea behind almost all modern post-training.

You do not need much background to follow along. I assume little to no prior knowledge of probability or statistics, and we will unpack each expectation and mathematical step as it appears.

The core policy-gradient move is simple: if an action leads to better future outcomes than expected, make that action more likely next time. The hard part is being precise about what "future outcomes" means. (If you want to know more about RL in general, see my Thinking in RL post.)

And here we are, the policy gradient therom could be seen as chain of equalities:

Each arrow is not a vibe. Each arrow is a legal replacement inside an expectation. We will work through them from the simplest case to the more general one, ending with why we use an advantage instead of a raw reward.

The one-answer case

Start with the smallest language-model RL problem. A prompt $x$ arrives. In this setting, $x$ is often a math problem or a question-answering problem that requires only one complete response to finish the interaction. The model samples that response $y \sim \pi_\theta(\cdot \mid x)$, and a judge gives it one terminal reward $R(x,y)$. If the full response is treated as one action, the objective is:

What does this objective mean? Imagine repeatedly sampling answers from the policy. Some answers appear more often because the model assigns them more probability. $J(\theta \mid x)$ is the average reward across those samples, or equivalently, a weighted average of every possible answer's reward, where the weight is $\pi_\theta(y \mid x)$. We maximize $J$ so the policy places as much probability mass as possible on answers that receive high rewards, such as correct answers to the math or question-answering problem.

To make this weighted average explicit, write the expectation as a sum:

Now differentiate it:

The tiny trick is:

Intuition: we are not differentiating the judge. We are differentiating how likely the model was to produce the sampled answer. High reward says "increase this answer's log-probability." Later, after we subtract a baseline, worse-than-expected reward says "decrease it."

This one-answer case is the whole story when $T=1$. Now we make $T$ bigger.

The trajectory case

A real trajectory has many actions:

For now, use the undiscounted return:

The objective is:

The trajectory probability factors as:

Only the policy terms contain $\theta$. The initial-state distribution $\rho$ and environment transition $P$ do not. So:

Now use the same log-derivative trick as before:

This is the first raw trajectory formula. It says: attach the total return $G_1$ to every action.

From total return to reward-to-go

But should an action receive credit for rewards that happened before the action was chosen?

Before touching the algebra, imagine a game. An agent arrives at a checkpoint with 10 coins already in its pocket. It must now choose left or right:

• Left leads to 2 more coins, for 12 total.
• Right leads to 5 more coins, for 15 total.

The 10 old coins change both totals, but they do not help us decide between left and right. Adding the same 10 to both choices preserves their difference: right is better by 3 either way. The current action cannot travel backward and collect, lose, or deserve credit for those 10 coins.

That is the intuition behind reward-to-go: when training the action chosen at time $t$, keep rewards that can still depend on that action and drop rewards that were already settled.

For a language model trained with only one terminal reward, this particular simplification changes nothing numerically. Before every token, the accumulated reward is zero; the entire reward still lies in the future. Thus $G_t=G_1=R(x,y)$ for every token. We need the general argument because other environments can have intermediate rewards, costs, tool outcomes, or safety penalties.

Now define the reward-to-go from time $t$:

So:

For the gradient term at time $t$:

because the past part has zero expectation when multiplied by the score function.

Expectation is just an average over possible worlds

The expectation symbol can make a simple operation look mysterious. Imagine replaying the random experiment many times. Each replay produces one possible world $w$ and one value $X(w)$. The expectation is the long-run average:

In reinforcement learning, one possible world is indeed one complete trajectory $\tau\sim p_\theta(\tau)$. We never stop sampling complete trajectories.

The key is that one complete trajectory is generated in order. By time $t$, the rollout has already produced a history $H_t$. From that history, its next random draw is

So the action sample is not a new experiment that replaces the trajectory sample. It is one draw inside the trajectory we already sampled. Equivalently:

The picture below shows this without the probability bookkeeping. Start with complete rollouts, pause them at one shared prefix, and inspect which action each rollout takes next.

This is what conditioning does. For a quantity such as $Z_t=\nabla_\theta\log\pi_\theta(a_t\mid s_t)$, which is determined once the history and current action are known, the same trajectory average can be written in two stages:

Read this from right to left: fix the prefix, average over the action that comes next, and then average over all prefixes. That produces the same average as sampling complete trajectories directly. This is the tower property.

In actual Monte Carlo training, we do not need to resample the action separately. One sampled trajectory already gives one sampled $H_t$ and one sampled $a_t$. Across many trajectories, those observed actions estimate the inner action average.

The past reward is not one global constant across trajectories: one rollout may arrive with 2 coins and another with 10. But after fixing one exact $H_t$, its past is fixed. This is why conditioning lets us ask whether that shared past can influence the action sampled next.

Four legal moves inside an expectation

Most expectation manipulations in this post use four rules:

Split sums. Expectation is linear, so $\mathbb{E}[X+Y]=\mathbb{E}[X]+\mathbb{E}[Y]$.

Pull out what is fixed. If $c$ does not vary over the worlds currently being averaged, then $\mathbb{E}[cX]=c\mathbb{E}[X]$. Always ask: fixed with respect to which randomness?

Condition to freeze information. A quantity may vary globally but become fixed inside a conditional expectation. If $Y$ is completely determined by $H$, then $\mathbb{E}[YX\mid H]=Y\mathbb{E}[X\mid H]$.

Average the groups back together. Use the tower property. If the inner average is zero in every group, the outer average is an average of zeros and is also zero.

Freeze the prefix, then average its next actions

Define the history immediately before sampling the current action:

In an autoregressive language model, this is the prompt plus the previously generated tokens. It contains $s_t$, but not $a_t$ or anything that happens after $a_t$.

Also name the two quantities we care about:

$\text{Past}_t$ varies across full trajectories, so we cannot pull it out of the outer expectation. But once we condition on one exact $H_t$, the past is already written and becomes fixed. Only the current action is resampled.

Now follow the picture.

1. Group by history using the tower property:

2. Inside one history bucket, pull out the fixed past:

3. Average the score over the possible current actions:

This is the algebraic version of the last visual panel. The sum over $a$ does not mean that the training procedure stopped sampling trajectories. It lists the possible next-action branches inside one fixed history. A single rollout takes one branch; many rollouts estimate their weighted average.

In words:

Why should this be zero? The action probabilities must always sum to 1. Increasing some probabilities forces others down; the policy cannot create or destroy total probability mass. The score of each individual action is generally not zero, but its probability-weighted average is.

4. Recombine the history buckets:

Finally, use linearity to split total return into past and future:

That is why $G_1$ may be replaced by $G_t$ in the $t$-th gradient term. We did not assume that past reward was globally constant. We grouped trajectories until it became locally constant, proved that its contribution was zero in every group, and then averaged the groups back together.

Applying that equality to every timestep gives:

If using discounting, define:

For the objective $J=\mathbb{E}[G_1]$, the exact derivation gives a factor $\gamma^{t-1}G_t$ in front of the $t$-th score term. Many presentations either include that factor or absorb it into how timesteps are weighted. The important idea is unchanged: past rewards vanish; future rewards stay.

From sampled return to Q

Reward-to-go $G_t$ is still a sampled future. If you take the same state $s_t$ and same action $a_t$, many different futures can happen.

The action-value function averages those futures:

We have finally arrived at the Q-function. It is powerful because it turns all those uncertain futures into one number for each state-action pair. At a fixed state $s$, $Q^\pi(s,a)$ asks: if I take action $a$ now and follow policy $\pi$ afterward, what return should I expect? Once we know these action values, we can compare the available actions and choose the one with the largest expected return:

This is the central action-selection idea behind value-based methods such as Q-learning.Q-learning was introduced in Watkins's 1989 thesis, with its convergence result developed in detail by Watkins and Dayan (1992). The distinction is important: here we use $Q^\pi$, the expected return under the current policy. Classical Q-learning instead learns the optimal action-value function $Q^*$ using a target containing $\max_{a'}Q(s',a')$.

For our policy-gradient derivation, we are not switching algorithms. We are asking one precise question: can the noisy sampled future $G_t$ be replaced by its conditional average $Q^\pi(s_t,a_t)$ without changing the expected gradient?

To answer it, let:

The key is the tower property of expectation:

The second line is legal because $Z_t$ is fully determined once $(s_t,a_t)$ is known. It can be pulled outside the inner expectation.

So:

Intuition: $G_t$ is one roll of the future dice. $Q^\pi(s_t,a_t)$ is the average of many such rolls. They have the same expectation in the gradient, but $Q$ is less noisy if you can estimate it well.

From Q to advantage

The Q-function gives an action an absolute score: how much return do we expect if we take this action? But an absolute score is missing context. The same $Q$-value might be impressive in a difficult state and disappointing in an easy one.

Advantage turns that absolute score into a relative gain:

A positive advantage says that this action gives us more return than we would normally expect at this state. A negative advantage says that choosing it gives up expected return. The magnitude tells us how large that gain or loss is. This is the information we want for a policy update: not merely whether an action is good, but whether it is better than what the policy usually does at this state.

To see why subtracting $V^\pi(s)$ gives us this cleaner signal without changing the expected policy gradient, we need to track two expectation scopes. The notation often hides the switch between them.

We begin at the outer scope:

This expectation still samples complete trajectories. A full trajectory determines which states are visited, which actions are sampled there, and which futures occur.

Temporarily zoom into one fixed state

The value function is the policy's average action-value at state $s$:

This action-only expectation is a conditional slice of the trajectory distribution, not a new training sampler. We pause full trajectories that reached the same state $s$ and ask which action they take next.

There are two nested averages hiding inside $V$:

$Q(s,a)$ first averages the possible futures after a fixed state-action pair. Then $V(s)$ averages those $Q$ values over the actions the policy might choose at the fixed state.

First, watch the cancellation in one state

Suppose a policy reaches one state $s$ and can go left or right. Let

and suppose $Q(s,\text{left})=12$ while $Q(s,\text{right})=8$. The value of the state is the policy-weighted average:

Let $\theta$ be the logit that controls the probability of going left. The two possible score terms point in opposite directions:

Now keep only the baseline part of the expected gradient. It contributes

The baseline says "add the same $9.6$ no matter which action happened." But increasing the left logit moves probability mass from right to left: the left branch pushes up exactly as much as the right branch pushes down after weighting by how often each branch occurs.

Subtracting the baseline therefore changes the numbers attached to individual samples, but not their expected gradient:

The same fact as a general proof

The numbers above are not a coincidence. Because $V^\pi(s)$ is fixed while we average the possible actions at that state, for any fixed $s$:

Writing $Z(s,a)=\nabla_\theta\log\pi_\theta(a\mid s)$, this proves a local statement inside one state bucket:

From one state bucket back to full trajectories

The previous subsection proved a local fact: after we fix one state $s$ and average over the actions available there, subtracting $V^\pi(s)$ contributes zero to the expected gradient. But the policy-gradient objective does not stay inside one fixed state. It averages complete trajectories, and different trajectories visit different states.

This subsection lifts that local result back to the full objective. If every state bucket contributes zero, then weighting those buckets by how often the policy visits them still gives a weighted average of zeros. That is what lets a proof about one fixed state justify the baseline subtraction across the entire trajectory distribution.

The notation can create a false picture here. It may look as if we first sampled a trajectory, threw it away, sampled a state, and then sampled an action. That does not happen.

Fix one timestep $t$. A complete trajectory is a row containing many columns:

Sampling one row automatically samples every column in that row. If we later ignore every column except $(s_t,a_t)$, we have not drawn anything new. We have only projected each full-trajectory sample down to two of its coordinates.

The six-row toy table makes the operation literal. Averaging a function $F(s_t,a_t)$ row by row gives

Putting the same rows into state buckets and then action buckets gives

The numbers did not change. Only the parentheses changed.

The exact probability proof

For one fixed $t$, start from the full-trajectory expectation:

Now partition the set of complete trajectories by the state and action appearing in their $t$-th columns:

The probability in brackets is the total probability mass of all full trajectories whose row has that state-action pair. By the ordinary product rule of probability,

The second factor is $\pi_\theta(a\mid s)$ because the policy is exactly the mechanism that chooses the current action after seeing the current state. For an autoregressive language model, $s_t$ is the entire prompt-plus-token prefix, so it contains everything the next-token policy conditions on.

Substitute that product into the previous sum:

So the right-hand side is not a recipe for collecting different data. It is the same full-trajectory expectation after grouping its probability mass first by state and then by action. This is the tower property of expectation.

Finally choose

and use linearity to sum over timesteps. Inside each fixed-state bucket, $V^\pi(s_t)$ is constant and $\mathbb{E}[Z_t\mid s_t]=0$. Therefore:

Every state bucket contributes zero. The outer trajectory expectation merely weights those buckets by how often full trajectories reach them, so it is a weighted average of zeros.

Therefore, back at the outer scope:

So the policy-gradient theorem becomes:

In actual training there is no separate action-only rollout phase:

1. Sample one complete trajectory $\tau$.
2. For every visited $(s_t,a_t)$, estimate $Q(s_t,a_t)$ and $V(s_t)$.
3. Form $A_t=Q(s_t,a_t)-V(s_t)$ and multiply it by that action's log-probability gradient.

The action expectation appears in the proof and in the definition of $V$. The training sample is still the full trajectory.

Intuition: $Q$ asks, "How good is this action?" Advantage asks, "How good is this action compared with what the policy usually does at this exact state?" That comparison is what gives the clean push-up or push-down signal.

Back to LLMs with terminal reward

For an autoregressive language model, the trajectory is a sequence of token actions:

If there is only one reward at the end, set:

Then, without discounting:

This is the exact reason terminal-reward LLM RL broadcasts one final score across all sampled tokens:

If we train a critic, use rollouts, or use a group baseline, we are trying to replace that raw terminal $R$ with a better estimate of advantage:

The policy-gradient machinery did not change. Only the multiplier attached to each token became more informative:

The whole derivation in one line

The mathematically careful story is:

Each line is the same gradient in expectation:

• $G_1 \to G_t$: remove past rewards; they are action-independent baselines.
• $G_t \to Q^\pi$: average over possible futures using conditional expectation.
• $Q^\pi \to A^\pi$: subtract the state value baseline; it has zero expected score-function term.

That is the math behind policy gradient: do not backprop through the reward. Backprop through the log-probability of the sampled action, weighted by how much better its future was than expected.

@misc{zhou2026policygradient,
  author = {Xuhui Zhou},
  title = {Policy Gradient Review 2026: From One Answer},
  year = {2026},
  howpublished = {\url{https://xuhuiz.com/blog/policy-gradient-from-one-action}},
}