Policy Gradient

Core Question

Reinforcement learning asks an agent to improve its behavior from experience, without a supervisor that names the correct action at every state. Policy gradient answers this by treating the policy itself as the object being optimized, rather than treating value estimation as a separate problem to solve first and behavior as something extracted from it afterward. The question motivating everything below is:

How can an agent update a policy

$$\pi_\theta(a \mid s)$$

so that actions leading to higher future return become more likely?

This sounds simple, but it hides two real difficulties. First, the policy assigns probability to actions, not directly to the returns those actions eventually produce, so there is no ready-made loss to differentiate. Second, the connection between a small change in $\theta$ and the resulting change in expected return runs through an entire trajectory of future states, actions, and rewards, not through a single number computed in one step. The rest of this note is essentially an answer to how that connection can be made precise, and differentiable.

DefinitionPolicy, Trajectory, and Return

The policy $\pi_\theta(a \mid s)$ is the object being optimized: a parameterized conditional distribution over actions given a state, controlled entirely by $\theta$.

$$\pi_\theta(a \mid s)$$

Rolling this policy out in the environment for one episode produces a trajectory: the full sequence of states and actions visited from the start of the episode to its end.

$$\tau = (s_1, a_1, s_2, a_2, \ldots, s_T, a_T, s_{T+1})$$

Each trajectory carries a trajectory return, the discounted sum of rewards collected along it.

$$R(\tau)$$

Because the trajectory itself is random (it depends on which actions the policy happened to sample, and on how the environment happened to respond), the quantity actually worth optimizing is not the return of any single rollout, but the expected trajectory return, averaged over every way a rollout under $\pi_\theta$ could unfold.

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

Relationship to Value-Based Methods

Before working through the derivation, it helps to place policy gradient against the family of methods it departs from. Value-based methods (dynamic programming, Monte Carlo control, temporal-difference learning) all work by first estimating how good a state or a state-action pair is, and only afterward deriving behavior from those estimates. Formally, they learn values such as:

$$Q^\pi(s,a)$$

and extract a policy from them after the fact, typically by acting greedily with respect to the learned value.

Policy-gradient methods skip that intermediate step entirely. Instead of first learning “how good is this action” and then converting that judgment into a decision rule, they parameterize the decision rule itself and adjust its parameters directly:

$$\pi_\theta(a \mid s)$$

The key shift is from “estimate which action is valuable” to “increase or decrease the probability of sampled actions based on return or advantage.” This is not merely a stylistic preference between two equally good options. Section 1 below works through the specific practical problems that value-based extraction runs into, and why parameterizing the policy directly sidesteps them.

Derivation

Here $J(\theta)$ is the expected return of the policy $\pi_\theta$, and $\nabla_\theta J(\theta_t)$ is the gradient of the expected return with respect to the parameters $\theta$ at time step $t$. The goal of gradient ascent is to maximize the expected return by iteratively updating the parameters in the direction of the positive gradient.

1. Objective, Notation, and Intuition

Motivation

In a deterministic policy, the policy maps a state directly to an action:

$$\mu_\theta : \mathcal{S} \to \mathcal{A}.$$

However, policy-gradient methods usually start from a stochastic policy:

$$\pi_\theta(a \mid s),$$

which maps a state to a probability distribution over actions. The sampled action is written as:

$$a_t \sim \pi_\theta(\cdot \mid s_t),$$

where the dot means that the action variable is left open: given state $s_t$, $\pi_\theta(\cdot \mid s_t)$ is the full probability distribution over possible actions.

In practice, $\pi_\theta$ may be implemented by a neural network, a linear function approximator, or another parameterized controller. The important point is that $\theta$ controls the behavior of the policy, even if the policy is not a simple table.

This distinction matters because policy gradient does not update a value table first. It directly updates the parameters of the policy that assign probabilities to sampled actions.

Earlier dynamic programming, Monte Carlo, and temporal-difference methods usually learn value estimates such as:

$$V^\pi(s) \quad \text{or} \quad Q^\pi(s,a).$$

In tabular methods, these values are stored explicitly for states or state-action pairs. For example, a tabular Q-learning style update has the form:

$$Q(s_t,a_t) \leftarrow Q(s_t,a_t) + \alpha \left[ r_{t+1} + \gamma \max_{a'} Q(s_{t+1},a') - Q(s_t,a_t) \right].$$

This family of methods is conceptually important, but it has practical limitations. A table-based value function is only natural when the state and action spaces are small enough to enumerate. Once states become high-dimensional observations, images, videos, or continuous feature vectors, maintaining a separate value for every possible $(s,a)$ pair becomes infeasible.

Even when $Q(s,a)$ can be approximated by a function, the policy is still usually derived indirectly, for example by choosing:

$$a^*(s)=\arg\max_a Q(s,a).$$

This indirect extraction can also be awkward in continuous action spaces, where computing $\arg\max_a Q(s,a)$ may require solving a difficult optimization problem at every decision step. It is also less natural when the desired behavior is a stochastic policy rather than a deterministic greedy action.

Policy gradient addresses this by parameterizing the policy directly and optimizing its parameters:

$$\theta_{t+1} = \theta_t + \alpha \nabla_\theta J(\theta_t).$$

This is gradient ascent on the expected return objective $J(\theta)$. It is analogous in form to neural-network gradient descent,

$$\theta_{t+1} = \theta_t - \alpha \nabla_\theta L(\theta_t),$$

but the sign is different because supervised learning usually minimizes a loss $L(\theta)$, while policy gradient maximizes expected return $J(\theta)$.

Before Diving In

As its name suggests, policy gradient is an on-policy method for finding an optimal policy $\pi_\theta$, where $\theta$ is the parameter vector of the policy. The goal is still an optimization problem, but we do not directly differentiate a supervised loss against a ground-truth label. Instead, we maximize expected return under trajectories sampled from the current policy.

DefinitionMarkov Decision Process

Consider a complete system for MDP defined by tuple

$$(S, A, P, R, \gamma)$$

where:

• $S$ is the state space;
• $A$ is the action space;
• $P$ is the transition probability;
• $R$ is the reward function;
• $\gamma$ is the discount factor.

The one-step reward function $R$ can be written as:

$$R(s, a) = \mathbb{E}[R_{t+1} \mid s_t = s, a_t = a]$$

When the action is sampled by the policy,

$$a_t \sim \pi_\theta(\cdot\mid s_t),$$

the observed reward becomes random through the sampled action and the next environment transition. The reward function itself is not directly parameterized by $\theta$; $\theta$ changes the distribution over actions and therefore the distribution over trajectories.

The key point is that $a_t$ is a random variable, and the one-step reward function is not usually differentiable with respect to $\theta$ directly.

DefinitionReturn and Trajectory Return

The return $G_t$ is defined as the total discounted reward from time step $t$ to the end of the episode. One common convention is:

$$G_t = \sum_{k=t}^{T-1} \gamma^{k-t} R_{k+1} = R_{t+1} + \gamma R_{t+2} + \gamma^2 R_{t+3} + \ldots + \gamma^{T-t-1} R_T$$

Recall the trajectory $\tau = (s_1, a_1, s_2, a_2, \ldots, s_T, a_T, s_{T+1})$. The final state is included here because the transition term $p(s_{t+1}\mid s_t,a_t)$ appears in the trajectory probability; some texts omit it from the shorthand notation.

This allows us to define the return of a trajectory as:

$$R(\tau) = \sum_{t=1}^{T} \gamma^{t-1} R_t = G_0$$

RemarkIndexing Conventions

This equality uses the convention that the first sampled reward is $R_1$ and the full-trajectory return is $G_0$. If a text starts the trajectory at $(s_1,a_1)$ and names the first-step return $G_1$, the same quantity may be written as $R(\tau)=G_1$.

Different textbooks use slightly different indexing, but they refer to the same discounted sum after reindexing. In this note, action $a_t$ is taken at state $s_t$, the next reward is $R_{t+1}$, and the reward-to-go from decision time $t$ is:

$$G_t = \sum_{k=t}^{T-1} \gamma^{k-t} R_{k+1}.$$

Under this convention, the full-episode trajectory return can be written as:

$$R(\tau) = \sum_{t=1}^{T} \gamma^{t-1} R_t = G_0.$$

Other common conventions include:

• starting from $t=0$ and writing $R(\tau)=\sum_{t=0}^{T-1}\gamma^t R_{t+1}$;
• writing reward-to-go as $G_t=\sum_{k=0}^{T-t-1}\gamma^k R_{t+k+1}$;
• using lowercase rewards, such as $r_t$ or $r_{t+1}$, instead of $R_t$;
• omitting the final state $s_{T+1}$ from the shorthand trajectory notation.

The formulas are equivalent only after the indices are shifted consistently. Mixing conventions inside one derivation is the main thing to avoid, including across notes: the interaction sequence sketched in The Reinforcement Learning Problem used the $0$-indexed convention above ($s_0,a_0,\dots,s_T,r_T$) for informal illustration, while this note fixes the $1$-indexed convention throughout for the actual derivation.

To sum up, intuitively (using $\to$ to denote causality):

$$\theta \rightarrow \pi_\theta(s_t) \rightarrow a_t \rightarrow s_{t+1}, R_{t+1} \rightarrow \tau \rightarrow R(\tau)\ \text{or}\ G_t$$

2. Trajectory Probability

A trajectory is not fixed in advance. It is sampled by rolling out the current policy $\pi_\theta$ in the environment, so we write:

$$\tau \sim p_\theta$$

where $p_\theta$ is the trajectory distribution induced by the policy and the environment dynamics.

DefinitionTrajectory Distribution

Here, $p_\theta$ denotes the whole trajectory distribution, while $p_\theta(\tau)$ denotes the probability mass or density assigned to one particular trajectory. The trajectory probability factorizes as:

$$\begin{aligned} p_\theta(\tau) &= p(s_1) \pi_\theta(a_1 \mid s_1) p(s_2 \mid s_1,a_1) \pi_\theta(a_2 \mid s_2) p(s_3 \mid s_2,a_2) \cdots \\ &= p(s_1) \prod_{t=1}^{T} \pi_\theta(a_t \mid s_t) p(s_{t+1}\mid s_t,a_t). \end{aligned}$$

Here, $p(s_1)$ and $p(s_{t+1}\mid s_t,a_t)$ come from the environment. The policy terms $\pi_\theta(a_t\mid s_t)$ are the parts controlled by $\theta$.

This is important because it shows that the expected return $J(\theta)$ is a function of $\theta$ through the trajectory distribution, not through the reward function directly. The expected return objective is:

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

RemarkDefinition of Expectation

This is only the definition of expectation. If $\tau$ is a discrete random variable with probability mass function $p_\theta(\tau)$, then for any function $f$:

$$\mathbb{E}_{\tau \sim p_\theta}[f(\tau)] = \sum_\tau p_\theta(\tau) f(\tau).$$

Taking $f(\tau)=R(\tau)$ gives the expected trajectory return above. If the trajectory space is continuous, the summation is replaced by an integral, but the expectation notation means the same thing.

The key point is that $\tau$ is random, and $R(\tau)$ is therefore a random return. However, for a fixed parameter vector $\theta$, $J(\theta)$ is not a random variable; it is a deterministic scalar objective: the average return of the policy $\pi_\theta$.

From the optimization point of view, policy gradient is trying to solve:

$$\theta^* \in \arg\max_\theta J(\theta).$$

In practice, we do not know how to enumerate all possible trajectories, and the environment dynamics are usually not differentiable through the agent’s parameters. The trick is to estimate the gradient of $J(\theta)$ from sampled trajectories, then apply stochastic gradient ascent:

$$\theta_{k+1} = \theta_k + \alpha \widehat{\nabla_\theta J(\theta_k)}.$$

AssumptionRegularity Conditions

The derivation below assumes that $R(\tau)$ is not directly parameterized by $\theta$, the environment dynamics $p(s_{t+1}\mid s_t,a_t)$ do not depend on $\theta$, and the trajectory distribution is differentiable with respect to $\theta$ on the sampled support. It also assumes we can exchange gradient and summation or expectation:

$$\nabla_\theta \sum_\tau p_\theta(\tau)R(\tau) = \sum_\tau R(\tau) \nabla_\theta p_\theta(\tau).$$

For the log-derivative trick, we also need $p_\theta(\tau)>0$ on trajectories where $\log p_\theta(\tau)$ is evaluated.

3. Log-Derivative Trick

We start from the objective in summation form:

$$J(\theta) = \sum_\tau p_\theta(\tau)R(\tau).$$

Taking the gradient with respect to $\theta$:

$$\nabla_\theta J(\theta) = \nabla_\theta \sum_\tau p_\theta(\tau)R(\tau).$$

Since the reward function is treated as independent of $\theta$, the gradient acts on the trajectory probability:

$$\nabla_\theta J(\theta) = \sum_\tau R(\tau) \nabla_\theta p_\theta(\tau).$$

At this point the expression is not convenient, because it asks us to differentiate the probability of a complete trajectory. The log-derivative trick rewrites this term using a positive differentiable function $f_\theta(x)$.

LemmaLog-Derivative (Score-Function) Identity

For a positive, $\theta$-differentiable function $f_\theta(x)$,

$$\nabla_\theta f_\theta(x) = f_\theta(x) \nabla_\theta \log f_\theta(x).$$

Read the left-hand side carefully. The outcome $x$ is fixed; only the parameter $\theta$ is moving. So $\nabla_\theta f_\theta(x)$ means:

if we slightly change $\theta$, how does the probability or density assigned to this fixed outcome $x$ change?

It is not a derivative with respect to $x$. In policy gradient, $x$ will later become a sampled trajectory $\tau$, and we ask how the current policy parameters change the probability assigned to that trajectory.

ProofProof of the Log-Derivative Identity

For a scalar parameter $\theta$, the chain rule gives:

$$\frac{\partial}{\partial \theta} \log f_\theta(x) = \frac{1}{f_\theta(x)} \frac{\partial}{\partial \theta} f_\theta(x).$$

For a vector parameter $\theta=(\theta_1,\ldots,\theta_d)$, the same statement holds coordinate by coordinate:

$$\frac{\partial}{\partial \theta_j} \log f_\theta(x) = \frac{1}{f_\theta(x)} \frac{\partial}{\partial \theta_j} f_\theta(x), \qquad j=1,\ldots,d.$$

Stacking these partial derivatives back into a gradient vector:

$$\nabla_\theta \log f_\theta(x) = \frac{1}{f_\theta(x)} \nabla_\theta f_\theta(x).$$

Multiplying both sides by $f_\theta(x)$ gives:

$$\nabla_\theta f_\theta(x) = f_\theta(x) \nabla_\theta \log f_\theta(x).$$

The condition $f_\theta(x)>0$ is only there because $\log f_\theta(x)$ must be defined. In probability settings, this means we apply the identity on the support where the sampled outcome has nonzero probability or density.

CorollaryLog-Derivative Trick for Trajectory Probabilities

Applying the lemma with $f_\theta(x)=p_\theta(\tau)$:

$$\nabla_\theta p_\theta(\tau) = p_\theta(\tau) \nabla_\theta \log p_\theta(\tau).$$

Here $\tau$ is treated as a fixed sampled trajectory inside the derivative. We are not differentiating the realized states and actions themselves. We are differentiating how much probability the current policy assigns to seeing that complete trajectory.

RemarkTerminology and the Diffusion-Model Connection

This identity is commonly called the log-derivative trick. In stochastic-gradient literature, the resulting estimator is also called the score-function estimator, likelihood-ratio estimator, or REINFORCE estimator. It is related in spirit to DDPM and score-based diffusion derivations because both work with gradients of log probabilities. The object is different, though: here the score is a parameter score, $\nabla_\theta \log p_\theta(\tau)$, used to estimate $\nabla_\theta J(\theta)$ without differentiating through the environment. In DDPM-style derivations, the score usually means a data-space score such as $\nabla_x \log p_t(x)$, and training is usually presented through denoising score matching plus a reparameterized noising process.

TheoremPolicy-Gradient Identity (Score-Function Form)

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

Proof

$$\begin{aligned} \nabla_\theta J(\theta) &= \sum_\tau R(\tau) p_\theta(\tau) \nabla_\theta \log p_\theta(\tau) \\ &= \sum_\tau p_\theta(\tau) R(\tau) \nabla_\theta \log p_\theta(\tau) \\ &= \mathbb{E}_{\tau\sim p_\theta} \left[ R(\tau) \nabla_\theta \log p_\theta(\tau) \right]. \end{aligned}$$

The first line substitutes the corollary above into $\nabla_\theta J(\theta)=\sum_\tau R(\tau)\nabla_\theta p_\theta(\tau)$; the last line is the definition of expectation.

This is the central policy-gradient identity. It turns an optimization problem over expected return into an expectation of a score function, $\nabla_\theta \log p_\theta(\tau)$. In optimization language, this gives us a stochastic gradient estimator: sample trajectories from the current policy, compute the return, compute the score of the sampled trajectory, and move $\theta$ in the estimated ascent direction.

4. Removing Environment Terms

Here, “expand” does not mean Taylor expansion. It means substituting the factorized trajectory probability from the previous section and then using log-product rules.

PropositionScore of the Trajectory Probability Reduces to Policy Log-Probabilities

$$\nabla_\theta \log p_\theta(\tau) = \sum_{t=1}^{T} \nabla_\theta \log \pi_\theta(a_t\mid s_t).$$

Proof

Recall the trajectory probability:

$$p_\theta(\tau) = p(s_1) \prod_{t=1}^{T} \pi_\theta(a_t \mid s_t) p(s_{t+1}\mid s_t,a_t).$$

Taking $\log$ on both sides:

$$\begin{aligned} \log p_\theta(\tau) &= \log \left[ p(s_1) \prod_{t=1}^{T} \pi_\theta(a_t \mid s_t) p(s_{t+1}\mid s_t,a_t) \right] \\ &= \log p(s_1) + \sum_{t=1}^{T} \log \left[ \pi_\theta(a_t \mid s_t) p(s_{t+1}\mid s_t,a_t) \right] \\ &= \log p(s_1) + \sum_{t=1}^{T} \log \pi_\theta(a_t\mid s_t) + \sum_{t=1}^{T} \log p(s_{t+1}\mid s_t,a_t). \end{aligned}$$

The only algebra used here is:

$$\log(xy)=\log x+\log y, \qquad \log\prod_{t=1}^{T}x_t = \sum_{t=1}^{T}\log x_t.$$

Taking the gradient with respect to $\theta$:

$$\begin{aligned} \nabla_\theta \log p_\theta(\tau) &= \nabla_\theta \log p(s_1) + \sum_{t=1}^{T} \nabla_\theta \log \pi_\theta(a_t\mid s_t) \\ &\quad+ \sum_{t=1}^{T} \nabla_\theta \log p(s_{t+1}\mid s_t,a_t). \end{aligned}$$

The environment terms do not depend on $\theta$, so their gradients are zero:

$$\nabla_\theta \log p(s_1)=0, \qquad \nabla_\theta \log p(s_{t+1}\mid s_t,a_t)=0.$$

Thus:

$$\nabla_\theta \log p_\theta(\tau) = \sum_{t=1}^{T} \nabla_\theta \log \pi_\theta(a_t\mid s_t).$$

CorollaryPractical Policy-Gradient Formula

Substituting the proposition above into the policy-gradient identity gives:

$$\nabla_\theta J(\theta) = \mathbb{E}_{\tau\sim p_\theta} \left[ R(\tau) \sum_{t=1}^{T} \nabla_\theta \log \pi_\theta(a_t\mid s_t) \right].$$

This is why policy gradient is practical: we do not need to differentiate through the environment transition model. We only need the log probability that the current policy assigned to the actions it actually sampled.

5. Sampled Gradient Estimate

The expectation above is still not computable exactly in most environments.

AlgorithmMonte Carlo Policy-Gradient Estimate

Sample $N$ trajectories by rolling out the current policy,

$$\tau^{(i)}\sim p_\theta, \qquad i=1,\ldots,N,$$

and approximate the expectation by a sample average:

$$\widehat{\nabla_\theta J(\theta)} = \frac{1}{N} \sum_{i=1}^{N} R(\tau^{(i)}) \sum_{t=1}^{T} \nabla_\theta \log \pi_\theta(a_t^{(i)}\mid s_t^{(i)}).$$

Then perform gradient ascent:

$$\theta_{k+1} = \theta_k + \alpha \widehat{\nabla_\theta J(\theta_k)}.$$

This is the basic REINFORCE idea: actions that appear in high-return trajectories are reinforced more strongly. Through the policy parameterization and normalization over actions, increasing the probability of some sampled actions also changes the probability mass assigned to other actions.

So the mathematical nature is the same as many stochastic optimization algorithms:

• define an objective $J(\theta)$;
• derive a gradient identity;
• replace the true expectation by samples;
• update parameters using stochastic gradient ascent.

6. After the Core Derivation: Credit Assignment

The derivation above gives the cleanest form:

$$R(\tau) \sum_{t=1}^{T} \nabla_\theta \log \pi_\theta(a_t\mid s_t).$$

But this assigns the same whole-trajectory return $R(\tau)$ to every action in the episode. That is mathematically valid, but it is noisy. If an action happens early in the episode, using rewards that occurred before the action does not make causal sense for credit assignment.

A common refinement is to replace the full trajectory return with reward-to-go. Using the convention fixed above, this means using the future return after the sampled action, $G_t=\sum_{k=t}^{T-1}\gamma^{k-t}R_{k+1}$. Unlike the derivation so far, this refinement is a genuine claim: that trading $R(\tau)$ for $G_t$ leaves the estimator’s expectation essentially unchanged, and it is worth actually proving rather than just asserting.

LemmaThe Score Function Has Zero Expectation

For any state $s$,

$$\mathbb{E}_{a\sim\pi_\theta(\cdot\mid s)}\big[\nabla_\theta\log\pi_\theta(a\mid s)\big] = 0.$$

Proof

$$\mathbb{E}_{a\sim\pi_\theta(\cdot\mid s)}\big[\nabla_\theta\log\pi_\theta(a\mid s)\big] = \sum_a \pi_\theta(a\mid s)\,\nabla_\theta\log\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,$$

using the log-derivative trick from Section 3 for the second equality, and the fact that $\pi_\theta(\cdot\mid s)$ sums to $1$ for every $\theta$ in the last.

PropositionCausality: Only Future Reward Depends on the Present Action

$$\mathbb{E}_{\tau\sim p_\theta}\left[R(\tau)\sum_{t=1}^T\nabla_\theta\log\pi_\theta(a_t\mid s_t)\right] = \mathbb{E}_{\tau\sim p_\theta}\left[\sum_{t=1}^T \gamma^t\,G_t\,\nabla_\theta\log\pi_\theta(a_t\mid s_t)\right]$$

Proof

Expand $R(\tau) = \sum_{t'=1}^T \gamma^{t'-1}R_{t'}$ and swap the order of summation:

$$\mathbb{E}_{\tau\sim p_\theta}\left[R(\tau)\sum_{t=1}^T\nabla_\theta\log\pi_\theta(a_t\mid s_t)\right] = \sum_{t=1}^T\sum_{t'=1}^T\gamma^{t'-1}\,\mathbb{E}_{\tau\sim p_\theta}\big[R_{t'}\nabla_\theta\log\pi_\theta(a_t\mid s_t)\big].$$

Fix $t$ and split the inner sum at $t'=t$. For $t' \le t$, the reward $R_{t'}$ is determined by $(s_1,a_1,\dots,s_{t'-1},a_{t'-1})$, states and actions sampled strictly before $a_t$, so conditioning on the trajectory prefix up to $s_t$ leaves $R_{t'}$ fixed while $a_t\sim\pi_\theta(\cdot\mid s_t)$ is still random, and the lemma above gives

$$\mathbb{E}_{\tau\sim p_\theta}\big[R_{t'}\nabla_\theta\log\pi_\theta(a_t\mid s_t)\big] = 0 \qquad (t'\le t).$$

Only the terms with $t' \ge t+1$ survive:

$$\sum_{t=1}^T\sum_{t'=t+1}^T\gamma^{t'-1}\,\mathbb{E}_{\tau\sim p_\theta}\big[R_{t'}\nabla_\theta\log\pi_\theta(a_t\mid s_t)\big] = \sum_{t=1}^T\gamma^t\,\mathbb{E}_{\tau\sim p_\theta}\left[\Big(\sum_{t'=t+1}^T\gamma^{t'-1-t}R_{t'}\Big)\nabla_\theta\log\pi_\theta(a_t\mid s_t)\right],$$

factoring $\gamma^t$ out of the surviving terms. The remaining inner sum is exactly $G_t = \sum_{t'=t+1}^T\gamma^{t'-1-t}R_{t'}$ by definition, which gives the claimed identity.

RemarkThe Conventional Reward-to-Go Formula Drops a Factor of $\gamma^t$

The proposition just proved shows that the exactly unbiased reward-to-go estimator carries an extra $\gamma^t$ weight on the term for time step $t$. The formula below, used throughout the rest of this note (including the REINFORCE box and every downstream advantage-based estimator), instead uses $G_t$ without the $\gamma^t$ factor. This is not an oversight; it is the essentially universal convention in the policy-gradient literature, including Sutton and Barto’s own presentation and OpenAI Spinning Up, both cited in this note’s references. Dropping $\gamma^t$ makes the estimator the exact gradient of a closely related, undiscounted-per-step objective rather than of $J(\theta)$ itself, a small, well-documented bias traded for lower variance and an update rule that weights early and late time steps equally. Since this is a standard, deliberate simplification rather than an error, the rest of this note follows the convention.

PropositionReward-to-Go Policy-Gradient Estimator (Conventional Form)

$$\widehat{\nabla_\theta J(\theta)} = \frac{1}{N} \sum_{i=1}^{N} \sum_{t=1}^{T} G_t^{(i)} \nabla_\theta \log \pi_\theta(a_t^{(i)}\mid s_t^{(i)}).$$

This does not change the basic policy-gradient idea. It just gives each sampled action a more relevant measure of future return.

7. Baseline and Advantage as Variance Reduction

Baseline and advantage are useful, but they are not the core log-trick derivation. They answer a numerical optimization question: how can we reduce the variance of the sampled gradient estimator without changing its expected direction?

DefinitionBaseline

A baseline subtracts a state-dependent reference value:

$$G_t \quad\rightarrow\quad G_t-b(s_t).$$

The common choice is $b(s_t)=V^\pi(s_t)$.

PropositionBaseline Invariance

If the baseline $b(s_t)$ does not depend on the sampled action $a_t$, it does not bias the policy-gradient direction in expectation.

Proof

Subtracting a baseline changes the reward-to-go estimator’s expectation by exactly

$$\mathbb{E}_{\tau\sim p_\theta}\left[\sum_{t=1}^T b(s_t)\,\nabla_\theta\log\pi_\theta(a_t\mid s_t)\right] = \sum_{t=1}^T \mathbb{E}\Big[b(s_t)\,\underbrace{\mathbb{E}_{a_t\sim\pi_\theta(\cdot\mid s_t)}\big[\nabla_\theta\log\pi_\theta(a_t\mid s_t)\big]}_{=0 \text{ by the lemma above}}\Big] = 0,$$

using that $b(s_t)$ depends only on $s_t$ and can therefore be pulled outside the inner expectation over $a_t\sim\pi_\theta(\cdot\mid s_t)$, since it does not depend on $a_t$. The baseline term contributes exactly zero to $\mathbb{E}[\widehat{\nabla_\theta J(\theta)}]$, for any choice of $b$ satisfying this condition.

DefinitionAdvantage Function

$$A^\pi(s_t,a_t) = Q^\pi(s_t,a_t) - V^\pi(s_t).$$

DefinitionSample Advantage Estimate

Writing $\widehat A_t^{(i)} := G_t^{(i)} - b(s_t^{(i)})$ for the baselined reward-to-go, and choosing $b(s_t)=V^\pi(s_t)$ as above, $\widehat A_t^{(i)}$ is a Monte Carlo estimate of $A^\pi(s_t^{(i)},a_t^{(i)})$: the sampled reward-to-go $G_t^{(i)}$ is an unbiased (if high-variance) estimate of $Q^\pi(s_t^{(i)},a_t^{(i)})$, since $Q^\pi(s,a)=\mathbb{E}_\pi[G_t\mid S_t=s,A_t=a]$, so $\widehat A_t^{(i)}$ estimates $Q^\pi(s_t^{(i)},a_t^{(i)}) - V^\pi(s_t^{(i)}) = A^\pi(s_t^{(i)},a_t^{(i)})$.

CorollaryAdvantage-Based Policy-Gradient Estimator

$$\widehat{\nabla_\theta J(\theta)} \approx \frac{1}{N} \sum_{i=1}^{N} \sum_{t=1}^{T} \widehat A_t^{(i)} \nabla_\theta \log \pi_\theta(a_t^{(i)}\mid s_t^{(i)}).$$

This is the bridge to actor-critic and PPO. For this note, the main point is only that advantage replaces raw return with a more stable signal for stochastic gradient ascent.

REINFORCE

REINFORCE is the Monte Carlo policy-gradient algorithm corresponding to the derivation above. It samples complete episodes under the current policy and uses their returns to estimate the ascent direction for $J(\theta)$.

AlgorithmREINFORCE

$$\nabla_\theta J(\theta) \approx \frac{1}{N} \sum_{i=1}^{N} \sum_{t=1}^{T} G_t^{(i)} \nabla_\theta \log \pi_\theta(a_t^{(i)} \mid s_t^{(i)}).$$

In words: sample trajectories from the current policy, estimate which sampled actions led to good future return, then increase the log probability of those actions by gradient ascent.

PPO Bridge

PPO keeps the policy-gradient idea but constrains how far the new policy can move from the old policy.

DefinitionProbability Ratio

$$r_t(\theta) = \frac{\pi_\theta(a_t \mid s_t)} {\pi_{\theta_{\text{old}}}(a_t \mid s_t)}$$

Suggested Order

Do not study PPO deeply until the policy-gradient derivation is explainable.

Further Reading

This note assumes familiarity with the policy notation built up in Policies as Probability Distributions, and with the value functions and Bellman equations of Return, Value, and the Bellman Equation. The Bellman equation in particular is the useful background for the value-based methods this note contrasts itself against, and $V^\pi$, $Q^\pi$, and $A^\pi$ are exactly the objects the baseline and advantage sections above put to work. The estimator derived here is on-policy, which means each batch of trajectories funds exactly one update before going stale. Proximal Policy Optimization picks that problem up directly: it buys sample reuse with importance sampling, and spends the rest of its length paying for it.

References

• Datawhale Easy RL, Chapter 4
• OpenAI Spinning Up: Policy Optimization Introduction
• Mohamed et al., Monte Carlo Gradient Estimation in Machine Learning
• Ho et al., Denoising Diffusion Probabilistic Models