The Policy-Gradient Theorem

In policy-based reinforcement learning, we choose a parameterized policy π(a|s,θ), typically a neural network, and define the performance objective ``` J(θ) = E_{τ ~ π_θ} [ R(τ) ] ``` where τ = (s₀,a₀,s₁,a₁,…) is a trajectory sampled by running policy π_θ in the environment, and R(τ) is the total (possibly discounted) return along that trajectory. We want ∇_θ J(θ) so we can perform gradient ascent: θ ← θ + α∇_θJ(θ). The difficulty is structural. J(θ) is an expectation over a distribution p_θ(τ) that itself depends on θ, because the policy determines which trajectories are likely. Naively writing ``` ∇_θ J(θ) = ∇_θ ∫ p_θ(τ) R(τ) dτ ``` requires differentiating p_θ(τ), which factors as a product of the (unknown, non-differentiable) environment transition probabilities and the (known, differentiable) policy probabilities. We cannot compute ∇_θ p(s'|s,a) because we don't have a model of the world — and even if we did, it would generally not be differentiable in any useful closed form. The Policy-Gradient Theorem states that despite this, the gradient can be rewritten in a form that involves only ∇_θ log π(a|s,θ), the score function of the policy, and never requires differentiating the transition dynamics. Formally, for the episodic or average-reward setting, using the discounted state distribution d^π(s) induced by π, ``` ∇_θ J(θ) = E_{s~d^π, a~π_θ}[ ∇_θ log π(a|s,θ) · Q^π(s,a) ] ``` This is a landmark identity: it converts a derivative of an expectation over an intractable, θ-dependent measure into an expectation (over that same measure) of a tractable integrand. That means we can approximate it with Monte Carlo sampling — run the policy, observe returns, and average — which is exactly what REINFORCE: Monte Carlo Policy Gradient does.

Step 1 — Set up the objective. Consider the episodic case with trajectories of length T (finite or infinite-horizon with discount γ). Define ``` J(θ) = E_{τ~π_θ}[ R(τ) ], R(τ) = Σ_{t=0}^{T} γ^t r_t ``` The trajectory distribution factors as ``` p_θ(τ) = μ(s₀) · Π_{t=0}^{T-1} π_θ(a_t|s_t) · p(s_{t+1}|s_t,a_t) ``` where μ is the initial state distribution and p is the environment's transition kernel — both independent of θ.

Step 2 — Apply the log-derivative trick. For any differentiable, positive density p_θ(x), ``` ∇_θ p_θ(x) = p_θ(x) · ∇_θ log p_θ(x) ``` This is just the chain rule: ∇_θ log p_θ(x) = ∇_θ p_θ(x) / p_θ(x). It converts a gradient of a density into the density itself times a gradient of a log — turning an unknown normalizing structure into something we can sample from.

Step 3 — Differentiate the objective. ``` ∇_θ J(θ) = ∇_θ ∫ p_θ(τ) R(τ) dτ = ∫ ∇_θ p_θ(τ) R(τ) dτ = ∫ p_θ(τ) ∇_θ log p_θ(τ) R(τ) dτ = E_{τ~π_θ}[ ∇_θ log p_θ(τ) · R(τ) ] ``` (Swapping ∇ and ∫ is justified under mild regularity — dominated convergence — standard in this setting.)

Step 4 — Expand log p_θ(τ) and cancel dynamics. ``` log p_θ(τ) = log μ(s₀) + Σ_t log π_θ(a_t|s_t) + Σ_t log p(s_{t+1}|s_t,a_t) ``` Taking ∇_θ, only the middle sum survives, since μ and p do not depend on θ: ``` ∇_θ log p_θ(τ) = Σ_{t=0}^{T-1} ∇_θ log π_θ(a_t|s_t) ``` This is the crucial cancellation: every unknown environment term vanishes under the gradient. We are left with ``` ∇_θ J(θ) = E_τ [ ( Σ_t ∇_θ log π_θ(a_t|s_t) ) · R(τ) ] ```

Step 5 — Causality: replace total return with future return. The action a_t cannot causally affect rewards r_0,...,r_{t-1} that occurred before it. A more refined argument (using the tower property of expectation, sometimes called the "future reward" or "reward-to-go" argument) shows that the expectation of ∇log π(a_t|s_t) against past rewards is zero, so they may be dropped without changing the expectation: ``` ∇_θ J(θ) = E_τ [ Σ_t ∇_θ log π_θ(a_t|s_t) · G_t ], G_t = Σ_{k=t}^{T} γ^{k-t} r_k ``` This reduces variance relative to Step 4's form while remaining exactly unbiased.

Step 6 — Recognize the action-value function. Since E[G_t | s_t, a_t] = Q^π(s_t,a_t) by definition of the action-value function, we can replace G_t with its conditional expectation without changing the overall expectation (by the tower property again): ``` ∇_θ J(θ) = E_{s~d^π,a~π_θ}[ ∇_θ log π_θ(a|s) · Q^π(s,a) ] ``` This is the Policy-Gradient Theorem in its canonical form, where d^π(s) is the (discounted) stationary state-visitation distribution induced by following π from the start state. ```pseudocode

# Estimating the policy gradient (single-episode Monte Carlo form)

# Input: policy network pi_theta, environment env, discount gamma

# Output: gradient estimate g approximating grad_theta J(theta) function ESTIMATE_POLICY_GRADIENT(pi_theta, env, gamma): trajectory = [] s <- env.reset() done <- False while not done: a <- sample from pi_theta(. | s) # forward pass, sample action s_next, r, done <- env.step(a) trajectory.append((s, a, r)) s <- s_next T <- length(trajectory) G <- 0 g <- zero_vector(size = dim(theta))

# walk backward to compute reward-to-go G_t efficiently: O(T) for t from T-1 down to 0: (s_t, a_t, r_t) <- trajectory[t] G <- r_t + gamma G # G is now G_t score <- grad_theta( log pi_theta(a_t | s_t) ) # backprop through policy net g <- g + (gamma^t) score * G return g

# Usage in training loop: for episode in 1..N: g <- ESTIMATE_POLICY_GRADIENT(pi_theta, env, gamma) theta <- theta + alpha * g # gradient ASCENT on J(theta) ``` The pseudocode implements Step 5's causal form directly: it never requires a model of `env`, only the ability to sample from it and to differentiate log π_θ through automatic differentiation (a single backward pass per time step, exactly as in supervised learning).

Computational complexity. Per episode of length T, computing the Monte Carlo gradient estimate requires:

• one forward pass of the policy network per time step to sample an action: O(T · C_fwd), where C_fwd is the cost of one forward pass (network-size dependent, not algorithm-dependent);
• one backward pass per time step to obtain ∇_θ log π(a_t|s_t): O(T · C_bwd), with C_bwd typically comparable to C_fwd in modern autodiff frameworks;
• an O(T) backward sweep to compute all reward-to-go values G_t via the recurrence G_t = r_t + γG_{t+1}, avoiding the naive O(T²) resummation.

Total time per episode is O(T · (C_fwd + C_bwd)), i.e., linear in trajectory length and in network cost — no worse than ordinary supervised training over T labeled examples. Space is O(T · dim(θ)) if gradients are accumulated per step before summing, though in practice frameworks accumulate a single scalar loss and backpropagate once, using O(dim(θ)) space plus O(T) for stored log-probabilities and rewards.

Correctness — why the estimator is unbiased. The entire derivation is a chain of exact algebraic identities, not approximations: 1. The swap of gradient and integral (∇_θ ∫ = ∫ ∇_θ) is valid whenever p_θ(τ) is differentiable in θ and the integrand is dominated by an integrable function independent of θ in a neighborhood of θ — a standard regularity condition satisfied by essentially all smooth policy parameterizations (softmax, Gaussian policies) with bounded rewards. 2. The log-derivative identity ∇_θ p_θ = p_θ ∇_θ log p_θ is pure calculus (chain rule), valid wherever p_θ(τ) > 0. 3. The cancellation of environment terms is exact: log p(s_{t+1}|s_t,a_t) genuinely contains no θ, so ∇_θ of that term is identically zero — not approximately small. 4. The causality reduction (replacing total return R(τ) with reward-to-go G_t) uses the fact that for k < t, E[∇_θ log π_θ(a_t|s_t) · r_k] = 0. This follows because, conditioned on the history up to time t (which includes s_t but is generated before a_t is sampled), E_{a_t~π_θ}[∇_θ log π_θ(a_t|s_t)] = 0 for any fixed s_t — a general fact about score functions: ``` E_{x~p_θ}[∇_θ log p_θ(x)] = ∫ p_θ(x) ∇_θ log p_θ(x) dx = ∫ ∇_θ p_θ(x) dx = ∇_θ ∫ p_θ(x) dx = ∇_θ 1 = 0 ``` Since r_k for kidentity rather than an approximation, the raw Monte Carlo estimator built directly from sampled trajectories — using the observed return G_t in place of Q^π(s_t,a_t) — is an unbiased estimator of ∇_θJ(θ). This is the theoretical guarantee that justifies calling it a "gradient" method at all: stochastic gradient ascent using this estimator converges (under standard step-size conditions, Robbins–Monro) to a stationary point of J(θ).

Honest tradeoff — variance. Unbiasedness is not the whole story. The estimator's variance can be very high, because G_t is a noisy sum of many random future rewards, and because ∇_θ log π_θ(a_t|s_t) can have large magnitude in regions where the policy is close to deterministic. High variance means many episodes may be needed before the sample average concentrates near the true gradient direction — this is precisely the motivation for subtracting a baseline b(s_t) from Q^π (a technique developed in Baselines and Variance Reduction), which provably does not change the expectation (again by E[∇log π]=0) but can sharply reduce variance.

We trace the theorem on a tiny two-state, two-action MDP to make every quantity concrete.

Setup. States {A, B}. From A, actions {left, right}: `left` deterministically transitions to a terminal state with reward +1; `right` deterministically transitions to terminal with reward 0. Episodes are one step long (T=1), γ=1. Policy is a one-parameter softmax over the two actions at state A: ``` π_θ(left | A) = σ(θ) = 1 / (1 + e^{-θ}) π_θ(right| A) = 1 - σ(θ) ``` We never actually reach B in this trivial example (it exists only to emphasize that d^π assigns nonzero mass to A as the start state). We want ∇_θ J(θ) at θ = 0.

Exact objective. With one-step episodes, ``` J(θ) = π_θ(left|A)·1 + π_θ(right|A)·0 = σ(θ) ∇_θ J(θ) = σ(θ)(1-σ(θ)) ``` At θ=0: σ(0)=0.5, so the true gradient is 0.5 × 0.5 = 0.25.

Now apply the theorem's formula. Q^π(A, left) = 1, Q^π(A, right) = 0 (the one-step returns). The score functions of a softmax over two actions are: ``` ∇_θ log π_θ(left|A) = ∇_θ log σ(θ) = 1 - σ(θ) ∇_θ log π_θ(right|A) = ∇_θ log(1-σ(θ)) = -σ(θ) ``` The theorem says ``` ∇_θJ(θ) = Σ_a π_θ(a|A) · ∇_θ log π_θ(a|A) · Q^π(A,a) = σ(θ)(1-σ(θ))·1 + (1-σ(θ))(-σ(θ))·0 = σ(θ)(1-σ(θ)) ``` At θ=0 this again gives 0.25, matching the direct computation exactly — confirming the theorem's identity on a case simple enough to check by hand.

Monte Carlo estimate from sampled episodes. Suppose we roll out 4 episodes at θ=0 and observe:

Average contribution = (0.5 + 0.0 + 0.5 + 0.0)/4 = 0.25 — matching the true gradient in this small sample purely because the sample happens to mirror the population 50/50 split; with only 4 samples this is somewhat lucky, illustrating both that the estimator is unbiased in expectation and that finite-sample variance is real (a run of 4 unlucky samples, e.g. three "right"s and one "left", would give a noticeably different estimate).

Policy-gradient methods built on this theorem power continuous-control robotics, where action spaces are continuous and value-based methods struggle to take an argmax over infinitely many actions. They underlie game-playing agents such as those trained with actor-critic architectures for real-time strategy and Atari-style environments. In modern large language model alignment (RLHF), the policy-gradient theorem is the mathematical backbone of the PPO update used to fine-tune language models against a learned reward model — the "policy" is the language model's next-token distribution, and the score function ∇log π(token|context) is computed via ordinary backpropagation through the transformer. Autonomous-vehicle simulators, dexterous manipulation research, and financial trading agents that must act in continuous or combinatorially large action spaces all rely, directly or through a descendant algorithm, on this single identity. Any time you see an "actor" network updated by "log-probability times advantage," you are watching the Policy-Gradient Theorem in action.

:::quiz[Why is E_{a~π_θ}[∇_θ log π_θ(a|s)] equal to zero for any fixed state s?]

• Because π_θ is always a softmax function
• Because the gradient of a probability distribution's total integral (which is always 1) is zero, and this equals E[∇log π_θ] by the log-derivative identity ✓
• Because rewards are bounded
• Because it only holds at convergence, θ = θ*

Ask the tutor: Can you walk me through why replacing G_t with Q^π(s_t,a_t) in the gradient formula doesn't change the expectation?