Policy Gradient Estimation

Chapter 1: Finite Difference

The simplest idea: compute the gradient numerically by perturbing one parameter at a time and observing how U changes. For the i-th component of θ:

where e_i is the i-th standard basis vector (all zeros except a 1 in position i) and δ is a small step size. To estimate the full gradient for a policy with n parameters, you need n+1 evaluations of U: one at θ and one perturbed evaluation per parameter.

The trouble is that U(θ) itself is an expectation — you can't evaluate it exactly. You estimate it with m rollouts. So the total rollout budget is m(n+1) just to get one gradient estimate. For a neural network with 10,000 parameters that's 10,001 evaluation packs.

The centered (two-sided) version is more accurate: perturb in both directions and take the symmetric difference ( U(θ + δe_i) − U(θ − δe_i) ) / 2δ. This cancels the leading-order error, so the bias drops from O(δ) to O(δ²). The cost doubles (2n evaluations instead of n+1), but the gain is worth it in smooth problems.

The Taylor expansion makes the bias explicit. Forward difference: U(θ + δe_i) = U(θ) + δ ∂U/∂θ_i + (δ²/2) ∂²U/∂θ_i² + O(δ³). Dividing by δ and rearranging: (U(θ + δe_i) − U(θ))/δ = ∂U/∂θ_i + (δ/2) ∂²U/∂θ_i² + O(δ²). The error is δ/2 times the second derivative — proportional to δ and to the curvature. The centered version: (U(θ+δe_i) − U(θ−δe_i)) / 2δ = ∂U/∂θ_i + O(δ²), because the second-derivative terms cancel. The curvature bias vanishes to leading order.

python
import numpy as np

def finite_diff_gradient(policy, theta, delta, n_rollouts):
    """Estimate gradient of U(theta) via finite differences."""
    n = len(theta)
    grad = np.zeros(n)
    U0 = estimate_utility(policy, theta, n_rollouts)  # baseline

    for i in range(n):
        theta_plus = theta.copy()
        theta_plus[i] += delta
        U_plus = estimate_utility(policy, theta_plus, n_rollouts)
        grad[i] = (U_plus - U0) / delta  # forward difference

    return grad  # cost: (n+1) * n_rollouts rollouts total

def estimate_utility(policy, theta, n_rollouts):
    """Average return over m rollouts."""
    returns = [run_episode(policy, theta) for _ in range(n_rollouts)]
    return np.mean(returns)

For neural network policies (π_θ(a|s) implemented as a network with PyTorch/JAX), the score function ∇_θ log π(a|s) is computed by backpropagating through the network from the log-probability output. You don't compute it analytically — autograd does it. The workflow is: (1) run the forward pass to get action logits, (2) compute log π(a|s) for the taken action, (3) multiply by the rtg weight, (4) call .backward(). PyTorch accumulates gradients. After processing all steps in the batch, call optimizer.step(). One subtlety: you want gradient ascent, not descent. Either negate the loss before calling backward(), or use -logπ × rtg as the loss and standard gradient descent. Both are equivalent. Most RL libraries define the REINFORCE loss as -E[logπ × rtg] for this reason.

A useful sanity check when implementing any gradient estimator: run your code on a simple convex function where you know the exact gradient (e.g., U(θ) = −θ², gradient = −2θ). Compare your estimator's output to the exact gradient. If they agree within statistical noise, your implementation is probably correct. This gradient checking pattern is standard practice in RL codebases before deploying on real problems.

Kochenderfer also presents a smarter variant: regression gradient. Instead of probing one parameter at a time, draw m random perturbations Δθ⁽¹⁾, ..., Δθ^(m) (each perturbing all parameters simultaneously), collect the utility change ΔU⁽ⁱ⁾ for each, and fit a linear model:

where ΔΘ is the m×n matrix of perturbations and ΔΘ⁺ is its Moore-Penrose pseudoinverse. Each random perturbation probes all parameter directions at once, so the gradient estimate becomes efficient: roughly 2n perturbations suffice for an n-parameter policy, compared to n+1 for finite differences. The recommended rule of thumb from the book: use 2n perturbations with a perturbation magnitude of δ ≈ 1% of the parameter norm.

python
import numpy as np

def regression_gradient(policy, theta, delta, m_perturb, n_rollouts):
    """
    Regression gradient: random perturbations + least-squares fit.
    theta:     current parameters (n,)
    delta:     perturbation magnitude (radius of hypersphere)
    m_perturb: number of random perturbations (use 2*n rule of thumb)
    """
    n = len(theta)
    U0 = estimate_utility(policy, theta, n_rollouts)

    # Sample m random perturbations from unit sphere, scale by delta
    Delta_Theta = []
    Delta_U = []
    for _ in range(m_perturb):
        direction = np.random.randn(n)
        direction /= np.linalg.norm(direction)   # unit sphere
        d_theta = delta * direction
        d_U = estimate_utility(policy, theta + d_theta, n_rollouts) - U0
        Delta_Theta.append(d_theta)
        Delta_U.append(d_U)

    Delta_Theta = np.array(Delta_Theta)          # shape: (m, n)
    Delta_U = np.array(Delta_U)                  # shape: (m,)

    # Least squares: find grad that best explains Delta_U = Delta_Theta @ grad
    grad, _, _, _ = np.linalg.lstsq(Delta_Theta, Delta_U, rcond=None)
    return grad   # cost: m * n_rollouts rollouts total (vs (n+1)*n_rollouts)

Numerical example. Consider a 3-parameter policy (θ = [0.5, −0.2, 1.1]) in a simple environment. We want ∂U/∂θ₂ using finite differences with δ = 0.01 and m = 50 rollouts each:

Step 1: Estimate U(θ) ≈ 4.23 (mean of 50 rollouts).

Step 2: Perturb θ₂: θ + 0.01 e₂ = [0.5, −0.19, 1.1]. Estimate U(perturbed) ≈ 4.41 (50 rollouts).

Step 3: ∂U/∂θ₂ ≈ (4.41 − 4.23) / 0.01 = 18.0.

But wait — with only 50 rollouts, each estimate has standard error σ/√50 ≈ 0.14 (assuming reward standard deviation ~1). The difference (4.41 − 4.23) = 0.18 might be 100% noise. The finite difference estimate's standard error is 0.14√2 / 0.01 ≈ 20 — larger than the signal. This is exactly the regime where likelihood-ratio methods shine: they compute ∇logπ analytically, avoiding this amplification of rollout noise.

Chapter 2: The Log-Derivative Trick

Finite differences treats U as a black box. But U has structure we can exploit: it's an integral involving the policy π_θ, which we designed. We know its derivative analytically. The log-derivative trick is how we harvest that knowledge.

Start from the definition. A trajectory τ is a sequence of states and actions: (s₁, a₁, s₂, a₂, ..., s_d, a_d). The expected utility is:

where p_θ(τ) is the probability of trajectory τ under policy θ, and R(τ) is its total discounted reward. Moving the gradient inside the integral:

We can't sample from ∇_θp because it's not a probability distribution (it can be negative!). Here's the trick: multiply and divide by p_θ(τ):

So the final gradient estimator is:

Now what is ∇_θ log p_θ(τ)? Expand the trajectory probability:

Take the log: it turns the product into a sum. Then take the gradient w.r.t. θ. The initial state distribution p(s₁) doesn't depend on θ. The transitions T(s'|s,a) don't depend on θ (we're model-free). They vanish. Only the policy terms survive:

There's a deeper reason the transition terms cancel that's worth dwelling on. The Markov property tells us that state s_t+1 is determined entirely by (s_t, a_t) and the environment's randomness — not by θ. So when we take ∇_θ, the transition T(s_t+1 | s_t, a_t) contributes zero. Only π_θ(a_t | s_t) carries θ dependence. If the environment were not Markov (say, a complicated simulator with internal state), the transitions might depend on θ indirectly (through action history affecting hidden state), and this clean cancellation would break down. The Markov assumption is load-bearing for the likelihood ratio trick.

Formally: log p_θ(τ) = log p(s₁) + ∑_t [log T(s_t+1|s_t,a_t) + log π_θ(a_t|s_t)]. Taking ∇_θ: the log p(s₁) term vanishes (no θ). Each log T term vanishes (no θ). What remains is exactly ∑_t ∇_θ log π_θ(a_t|s_t). The derivation is three lines; the implication is enormous — the entire field of model-free RL depends on this cancellation holding.

Each term ∇_θ log π_θ(a_t | s_t) is called the score function at time t. It points in the direction in parameter space that most increases the log-probability of taking action a_t in state s_t. When multiplied by R(τ), high-reward trajectories push θ to make those actions more likely, and low-reward trajectories push in the opposite direction.

Concrete score functions for two common policy families:

Both score functions have a deep property: E_a∼π[∇ log π(a|s)] = 0. The expected score under the policy itself is always zero. This is why baselines (Chapter 5) work: you can subtract any state-dependent quantity from R and the expected gradient change is zero.

What happens at convergence? When the policy has converged to the optimal π*, the gradient should be zero (we're at a maximum). For the softmax: at the optimum π(best_action|s) → 1. The score for the best action is 1 − 1 = 0. All other actions have probability 0, so their contribution is also 0 × (−1) = 0. The entire gradient is zero. Gradient ascent correctly stops improving. For Gaussian policies, convergence means the mean μ has moved to the optimal action and the score (a − μ)/σ² is centered around zero, so its expectation is also zero. The gradient signal dies out as the policy peaks.

This is not a bug — it's the correct behavior. A peaked policy that's fully converged needs no further updates. The challenge is premature convergence: the policy might peak on a locally optimal but globally suboptimal action. Entropy regularization (adding H(π) to the objective) counteracts this by penalizing overconfident policies, keeping σ from shrinking to zero or the softmax from becoming too sharp too soon.

Chapter 3: REINFORCE — Live Demo

Williams (1992) combined the likelihood ratio gradient into a clean algorithm: REINFORCE. Collect trajectories from the current policy, estimate the gradient from each, take a gradient ascent step, repeat. The full gradient estimator is:

where m trajectories τ⁽¹⁾, ..., τ^(m) are sampled from the current policy, R(τ⁽ⁱ⁾) is the total return of trajectory i, and the inner sum is the summed score over all timesteps. After computing the gradient estimate, we take a gradient ascent step: θ ← θ + α ∇_θ U(θ).

Worked example from scratch. Suppose you have a 1-step MDP: one state, two actions (Left = reward −1, Right = reward +5). Your policy is a softmax parameterized by a single scalar θ: π(Right) = e^θ/(e^θ + 1) = σ(θ). Currently θ = 0, so π(Right) = 0.5.

You run m = 3 episodes. Suppose the outcomes are: episode 1 goes Right (reward +5), episode 2 goes Left (reward −1), episode 3 goes Right (reward +5). First, compute each score:

Now sum R × score over the three episodes:

With learning rate α = 0.1: θ ← 0 + 0.1 × 1.83 = 0.183. Now π(Right) = σ(0.183) ≈ 0.546. The policy has shifted slightly toward Right. After many such updates, θ grows large and π(Right) → 1.0.

Below is a live simulation of exactly this MDP. The agent starts at θ = 0 and you can watch REINFORCE drive it toward the optimal policy.

The policy is a softmax over a single parameter: π(Right | S) = σ(θ) = e^θ / (1 + e^θ). The gradient is:

Run episodes, watch the gradient accumulate, and see θ climb toward the optimal policy. Click "Auto-run" to let REINFORCE converge on its own.

An aside on deterministic policies. The likelihood ratio trick requires ∇_θ log π_θ(a|s) to exist. For a deterministic policy, π(a|s) is either 1 (for the greedy action) or 0 (for others). The log of 0 is undefined, and the gradient doesn't exist in the usual sense. This is why stochastic policies are essential for REINFORCE — without them, the log-derivative trick breaks. Deterministic policy gradient (DPG) methods (Ch 17) handle this via a different trick: they differentiate the deterministic action output directly through the Q-function, bypassing the log-probability entirely. But that requires a differentiable environment model or a differentiable Q approximation.

The practical takeaway: always use a stochastic policy for REINFORCE. For discrete action spaces, a softmax with a temperature parameter τ controls how peaked the distribution is: π(a|s) ∝ exp(θ_sa/τ). High temperature (large τ) = uniform/exploratory; low temperature = nearly deterministic. You can anneal τ from high to low over training to transition from exploration to exploitation. For continuous action spaces, a Gaussian policy π(a|s) = N(μ_θ(s), σ_θ²(s)) maintains stochasticity throughout training, with the learned σ naturally decreasing as the policy converges — though you may need a minimum σ floor to prevent premature convergence to a degenerate distribution.

Notice that at θ = 0, the policy is 50/50. As REINFORCE accumulates gradient signal from high-reward Right trajectories, θ grows. The policy becomes increasingly biased toward Right. Eventually π(Right) → ~1.0 and the updates plateau — the policy has converged.

Note that the gradient estimate from REINFORCE, even when computed correctly, is extremely high-variance in the beginning of training when the policy is nearly uniform. Every episode has nearly equal probability, so the gradient weights are nearly equal for all trajectories, making the net gradient close to zero. As the policy concentrates (high θ magnitude), the score function becomes larger for the preferred actions and the gradient signal clarifies. This is why REINFORCE often shows very slow initial progress followed by a "breakthrough" when the policy suddenly tips past 50% toward the optimal action — a phase transition in the gradient landscape.

The convergence rate depends on learning rate and batch size. A large learning rate (α = 0.5) converges faster but can overshoot — the policy oscillates between "mostly Right" and "all Right" before stabilizing. A small learning rate (α = 0.01) is stable but slow. Batch size matters too: with 1 episode per update, the gradient estimate is noisy and training is erratic; with 20 episodes, the estimate is smoother and the θ trajectory is much more monotone. This is the same bias-variance trade-off that appears everywhere in statistics, now applied to the gradient estimator.

One subtlety: in this 1-step MDP, the Left action always gives −1. If early rollouts happen to produce many Right episodes by chance (giving θ a head start), convergence is fast. If they produce many Left episodes (unlucky start), θ dips negative before climbing. The initial policy's randomness is the primary driver of early training variance — which is exactly why the baseline (subtracting the mean return ≈ 2.0 from the batch) stabilizes the gradient signal significantly.

Chapter 4: Reward-to-Go

REINFORCE as stated multiplies every step's score by the entire trajectory return R(τ). But can action a₂ (taken at time 2) affect reward r₁ (collected at time 1)? No. Causality forbids it. Past rewards are already fixed before action a₂ is chosen.

This means we're polluting the gradient signal for each step t with irrelevant noise from rewards before t. The reward-to-go (also called the "causality trick") substitutes the total return with only the future rewards, from step t onward:

The full policy gradient becomes:

The benefit compounds across long trajectories. For a trajectory of length d, step t uses d−t+1 future reward terms instead of d total terms. Early steps see the largest reduction. Step d uses only r_d; step 1 uses all d rewards. The average number of terms per step drops from d to (d+1)/2.

In practice, for long-horizon problems like Atari games (episode length ~1000 steps) or robotic locomotion (500+ steps), reward-to-go can reduce variance by a factor of 5-10x compared to total return. This translates directly to needing 5-10x fewer samples to achieve the same gradient signal quality — a massive practical advantage that costs nothing in terms of bias or additional computation. Every policy gradient implementation that doesn't use reward-to-go is leaving free variance reduction on the table.

Another subtlety: the discount γ in the reward-to-go has a dual role. Mathematically, it down-weights future rewards to prioritize immediate returns, reflecting time preference. Practically, it also prevents the reward-to-go from blowing up in continuing (non-episodic) tasks where the sum could otherwise be infinite. For γ < 1, the geometric series ∑_t'=t^∞ γ^t'−t r_t' converges as long as rewards are bounded. The choice γ = 0.99 (vs 0.95 or 1.0) can significantly affect learned behavior: a lower γ makes the agent myopic and greedy; γ = 1.0 requires finite-horizon tasks to guarantee convergence of the sum.

One implementation subtlety: when computing reward-to-go across a batch of trajectories with different lengths, be careful to truncate correctly. Each trajectory has its own rtg array; they should not be mixed or padded incorrectly. A common bug is computing the running sum across trajectory boundaries, which mixes future rewards from one trajectory into past steps of another. Always maintain a separate rtg list per trajectory and concatenate after.

With a discount factor γ < 1, the reward-to-go also has a natural "effective horizon" of 1/(1−γ) steps — the geometric series converges to a finite value. For γ = 0.99, the effective horizon is ~100 steps: rewards more than 100 steps in the future are down-weighted by more than e⁻¹. This means in very long episodes (1000+ steps), the policy is effectively optimized for the near-term horizon, not the true long-run return. Choosing γ is thus a design decision about how far ahead the agent should plan, not just a mathematical convenience.

python
def compute_rewards_to_go(rewards, gamma=0.99):
    """
    Input:  rewards = [r1, r2, ..., rd] from a single trajectory
    Output: rtg[t] = gamma^0 * r_t + gamma^1 * r_{t+1} + ... + gamma^{d-t} * r_d
    """
    d = len(rewards)
    rtg = [0.0] * d
    running = 0.0
    for t in range(d - 1, -1, -1):  # iterate backward
        running = rewards[t] + gamma * running
        rtg[t] = running
    return rtg

# REINFORCE with reward-to-go:
for trajectory in batch:
    rtg = compute_rewards_to_go(trajectory.rewards)
    for t, (s, a, r) in enumerate(trajectory.steps):
        score = grad_log_pi(theta, a, s)  # nabla_theta log pi
        grad += score * rtg[t]             # weight by future rewards only

theta += alpha * grad / len(batch)

Worked example: reward-to-go vs. total return. Suppose a trajectory has 3 steps with rewards [1, 3, 2] and γ = 1.0 (no discounting for simplicity). Total return: R(τ) = 6.

With total return, every step gets multiplied by 6. Step 1's score gets multiplied by 6 — even though 5 of those 6 reward units came from steps 2 and 3, which had nothing to do with action a₁.

With reward-to-go: step 1 gets 1 + 3 + 2 = 6, step 2 gets 3 + 2 = 5, step 3 gets only 2. Now the weight for step 3 is tight (only the immediate reward), and the weight for step 1 is the most informative possible. Steps at the end of the trajectory contribute more signal per rollout, because their weights are smaller and thus less contaminated by the environment's randomness across future steps.

Chapter 5: Baseline Subtraction

Even with reward-to-go, the gradient can still be noisy. Here's a concrete failure mode: suppose all trajectories have positive return — the environment is generous. Then every action gets pushed to be more likely, just by different amounts. The gradient signal is drowned by the absolute scale of returns, not the relative quality of actions.

The fix: subtract a baseline b(s_t) from the reward-to-go at each step. The gradient becomes: