Until now the agent has been scoring moves and picking the highest. Policy gradients flip the script — the agent directly learns a strategy and adjusts it based on outcomes. Robots learn continuous motion this way. ChatGPT gets fine-tuned this way. The log-probability trick at the heart of REINFORCE is the same gradient that trains modern RLHF.
Learning Objectives
After this lesson, you will be able to:
Understand why policy gradient methods are needed: value-based methods cannot handle continuous or stochastic action spaces, but learning the policy directly can
Follow the REINFORCE algorithm and the log-probability gradient trick: ∇J(θ) = E[∇ log π(a|s) · G_t], which converts the unmathematical "do more of what works" intuition into a differentiable gradient
Understand why raw REINFORCE has high variance (single-episode Monte Carlo estimates) and how baselines and advantage functions reduce it without introducing bias
See how actor-critic combines the best of value-based and policy-gradient methods: the critic provides low-variance advantage estimates, the actor learns the policy directly
Before we flip the script: recall from the Q-learning lesson, write down (a) the Bellman update for Q(s,a) in one line, (b) what 'argmax over actions' is doing geometrically, and (c) one specific kind of action space where this argmax falls apart. Once you've committed an answer, the motivation for policy gradients below will land much harder.
Write your answer in your own words — don't look back at the lesson. This is the most effective way to remember what you just learned.
Everything we have seen so far -- Q-learning, DQN -- learns a value function (Q or V) and derives a policy from it (pick the action with the highest value). This is the value-based approach.
Policy gradient methods flip this around: they learn the policy directly. Instead of asking "how valuable is each action?", they ask "what is the probability of each action, and how should I adjust those probabilities to get more reward?"
Why learn a policy directly?
Continuous actions: DQN outputs Q-values for discrete actions. But what if the action is a continuous steering angle, a robot joint torque, or a continuous token probability? You cannot take a max over infinite actions. Policy gradients output continuous action distributions naturally.
Stochastic policies: Sometimes the optimal policy is stochastic (mixed strategies in games like rock-paper-scissors). Value-based methods produce deterministic policies (always pick the max). Policy gradients naturally represent probability distributions over actions.
Simpler representation: In some problems, the optimal policy is simpler than the optimal value function. A policy might just be "go toward the goal," while the value function encodes complex contours of expected return across the state space.
Better convergence properties: Policy gradients update the policy smoothly (small changes to parameters = small changes to action probabilities). Value-based methods can have discontinuous policy changes when Q-value rankings shift.
A policy network takes a state as input and outputs a probability distribution over actions:
πθ(a∣s)=P(At=a∣St=s;θ)
The goal is to maximize the expected return:
J(θ)=Eτ∼πθ[t=0∑TγtRt+1]
The policy gradient theorem tells us how to compute the gradient of J with respect to theta:
∇θJ(θ)=Eτ∼πθ[t=0∑T∇θlogπθ(At∣St)⋅Gt]
Try it! Imagine you played 3 episodes of a game. Episode 1: you went left and scored +10. Episode 2: you went right and scored -5. Episode 3: you went left and scored +8. What should you adjust? Increase the probability of going left (it led to positive outcomes) and decrease the probability of going right (negative outcome). That is exactly what the equation above does mathematically.
This is the most important equation in policy gradient methods. Let us unpack it:
log pi_theta(A_t|S_t): The log-probability of the action actually taken. Its gradient tells us how to nudge the parameters to make this action more likely.
G_t: The return (total discounted reward) from time step t onward. This tells us whether the action was good or bad.
The product: If the return was positive (good outcome), increase the probability of the action. If negative, decrease it. The magnitude of the return scales the update.
What Do You Think?
In REINFORCE, if an action leads to a return of 0, what happens to that action's probability?
If the return is exactly 0, the gradient contribution is zero -- the action probability does not change. In practice, returns are rarely exactly zero, and using a baseline (subtracting the average return) means a return of 0 would actually decrease the probability if the average return is positive.
The equation above is the punchline. Below is the derivation that earns it. We work in a discounted infinite-horizon Markov decision process M = (S, A, P, R, gamma) with state space S, action space A, transition kernel P(s' | s, a), reward function R(s, a), and discount factor gamma in (0, 1). The agent acts according to a parameterized stochastic policy pi_theta(a | s). A trajectory tau = (s_0, a_0, r_0, s_1, a_1, r_1, ...) is sampled by drawing s_0 from the initial distribution rho(s_0), then iterating a_t ~ pi_theta(. | s_t) and s_ ~ P(. | s_t, a_t). The performance objective is the expected discounted return:
J(θ)=Eτ∼πθ[t=0∑∞γtrt]=∫R(τ)pθ(τ)dτ
Step 1 — the log-derivative trick (score function). For any density p(x; theta) that is differentiable in theta and any bounded function f, observe that nabla_theta p = p * nabla_theta log p (just divide and multiply by p). Therefore:
The quantity nabla_theta log p_theta(X) is called the score, and it has zero mean under p_theta (the expectation of the gradient of the log-density is always zero, because the density integrates to one). This identity is older than RL: it powers maximum likelihood estimation, variational inference, and the reparameterization-free path in modern generative modeling.
Step 2 — apply the trick to trajectories. The trajectory density factorizes as
Taking the log turns the product into a sum, and the dynamics terms rho and P do not depend on theta — so their gradient is zero. We get:
∇θlogpθ(τ)=t=0∑∞∇θlogπθ(at∣st)
Combining steps 1 and 2 with f(tau) = R(tau) = sum_t gamma^t r_t yields the canonical policy gradient:
∇θJ(θ)=Eτ∼πθ[R(τ)⋅t=0∑∞∇θlogπθ(at∣st)]
Step 3 — REINFORCE with return-to-go (causality). A key observation: future rewards cannot depend on past actions, in expectation. Formally, for k < t,
E[r_k * nabla_theta log pi_theta(a_t | s_t)] = 0,
because conditional on (s_t, history), the score has zero mean (it is a score) and r_k is measurable with respect to the past. So the cross-terms between past rewards and present scores vanish. We may therefore replace the full trajectory return with the return-to-go from time t, G_t = sum_ gamma^ r_k, getting the REINFORCE estimator of Williams (1992):
∇θJ(θ)=Eτ∼πθ[t=0∑∞γt⋅Gt⋅∇θlogπθ(at∣st)]
Step 4 — baseline invariance. Subtract any function b(s_t) that depends only on the state (not the action) from G_t. The resulting estimator is still unbiased:
The variance-minimizing choice is b*(s_t) = E[G_t * ||nabla log pi||^2 | s_t] / E[||nabla log pi||^2 | s_t] — the score-squared-weighted conditional return. In practice nobody computes that; the canonical choice is the state value V^pi(s_t), which is near-optimal, easy to learn with a second network, and produces the advantage A^pi(s_t, a_t) = Q^pi(s_t, a_t) - V^pi(s_t).
What Do You Think?
Suppose we subtract a baseline b(s_t, a_t) that depends on BOTH state AND action. Is the resulting policy gradient still unbiased?
Step 5 — Generalized Advantage Estimation (Schulman et al., 2016). Even with V^pi as the baseline, the Monte-Carlo G_t is still high-variance. The other extreme — bootstrapping after one step with the TD residual
delta_t = r_t + gamma * V(s_) - V(s_t)
— gives a low-variance but biased advantage estimate (the bias comes from V being an imperfect approximator). GAE interpolates between these. Define:
The two limits make the trade-off explicit. When lambda = 0, A_hat_t = delta_t — the TD(0) advantage. When lambda = 1, the sum telescopes: A_hat_t = sum_ gamma^l * (r_ + gamma V(s_) - V(s_)), and the interior V terms cancel pairwise, leaving G_t - V(s_t) — the Monte-Carlo advantage. Lambda dials the bias-variance trade-off: small lambda trusts the critic (low variance, biased), large lambda trusts the actual rewards (high variance, unbiased). Schulman observed empirically that lambda ≈ 0.95 hits the sweet spot for most continuous-control tasks, and that hyperparameter has stuck.
The chain — score-function trick, dynamics cancellation, return-to-go via causality, baseline invariance, GAE — is the full derivation. Williams (1992) gave the original score-function argument for REINFORCE. Sutton, McAllester, Singh, and Mansour (2000) gave the formal Policy Gradient Theorem in the stationary-distribution form (Sutton and Barto, 2nd ed., Chapter 13). Schulman, Moritz, Levine, Jordan, and Abbeel (2016) introduced GAE as the practical advantage estimator that is now the default everywhere from PPO to RLHF.
Initialize the policy network with random weights. The network takes a state as input and outputs a probability distribution over actions. For discrete actions, this is a softmax layer. For continuous actions, it outputs the mean and standard deviation of a Gaussian.
At initialization, the policy is essentially random -- it assigns roughly equal probability to all actions regardless of state. The agent has no idea what to do.
The agent runs a complete episode using the current policy. At each time step, it samples an action from the policy distribution, executes it, and records the state, action, and reward. The episode produces a full trajectory: (s0, a0, r1, s1, a1, r2, ..., sT).
REINFORCE requires complete episodes -- it cannot update mid-episode. This is a Monte Carlo method: it waits for the final outcome before learning.
Working backward from the end of the episode, compute the discounted return at each time step: G_t = r_ + gamma * r_ + gamma^2 * r_ + ... The return at each step tells us the total future reward from that point onward.
Steps that led to high cumulative reward will have high G_t. Steps that led to poor outcomes will have low G_t. This is the "scorecard" for every action in the episode.
For each time step, compute the gradient of log pi(a_t | s_t) -- the log-probability of the action actually taken. Multiply it by the return G_t. This gives the policy gradient: actions that led to high returns get their probabilities increased, and actions that led to low returns get decreased.
The log-probability trick converts the problem into a form amenable to standard backpropagation. The magnitude of G_t scales how strongly each action is reinforced or punished.
Sum up the per-step gradients across the entire episode and perform a gradient ascent step. The policy parameters theta shift to make high-return actions more probable and low-return actions less probable.
This is gradient ascent, not descent -- we are maximizing expected reward. In practice, the loss is negated so standard optimizers (which minimize) can be used: loss = -sum(log_prob * G_t).
Discard the episode data and run a new episode with the updated policy. The new policy is slightly better -- it tends to take actions that worked well in previous episodes and avoids actions that led to poor outcomes.
Over hundreds or thousands of episodes, the policy converges toward optimal behavior. Each episode provides a noisy but unbiased estimate of the gradient. Averaging over many episodes smooths out the noise.
Loading visualization...
Try it: Watch the Policy ShiftInteractive
See how a policy gradient agent adjusts its action distribution over time. The policy starts uniform (equal probability for all directions), and as the agent collects trajectories, it gradually shifts probability mass toward actions that lead to higher returns. Watch the probability arrows grow and shrink in real time.
REINFORCE works in theory but has a critical practical issue: high variance. The return G_t is a noisy, Monte Carlo estimate. It depends on the entire future trajectory, which includes randomness from the policy, the environment, and the episode length. This noise makes the gradient estimates wildly variable, causing slow and unstable learning.
The key insight for reducing variance: subtract a baseline b(s) from the return. This does not change the expected gradient (it remains unbiased) but can dramatically reduce variance:
∇θJ(θ)=E[t=0∑T∇θlogπθ(At∣St)⋅(Gt−b(St))]
The best baseline is the value function V(s) -- the expected return from state s. This gives us the advantage:
Try it: See how subtracting a baseline changes the gradient signalInteractive
Actor-critic methods combine policy gradients with value function learning. Two networks work together:
Actor: The policy network -- chooses actions
Critic: The value network -- evaluates how good the current state is
The critic provides the baseline for the actor, reducing variance. The actor provides the behavior for the critic to evaluate.
Actor update: θ←θ+αθ∇θlogπθ(At∣St)⋅A^tCritic update: w←w−αw∇w(Gt−Vw(St))2where A^t=Rt+1+γVw(St+1)−Vw(St)
The advantage estimate A-hat_t = R + gamma * V(s') - V(s) is called the TD error. It uses one step of actual reward and then bootstraps off the value estimate, balancing bias and variance:
REINFORCE (Monte Carlo): Low bias, high variance -- waits for the full episode
TD Actor-Critic: Some bias (from the value estimate), much lower variance -- updates after each step
n-step Actor-Critic: Adjustable tradeoff -- uses n real steps before bootstrapping
Tests · Verify that after training, P(action 1) > 0.8. Implement the baseline version and verify it converges faster (higher P(action 1) after 500 episodes).
Policy gradients learn the policy directly. Instead of estimating values and deriving actions, policy gradient methods directly optimize the probability of taking good actions, enabling handling of continuous and stochastic action spaces
REINFORCE uses the log-probability trick. By multiplying the log-probability of actions by the return, high-return trajectories increase the probability of their actions while low-return ones decrease it
High variance is the main challenge. Raw policy gradients are extremely noisy because a single trajectory is a poor estimate of the expected return; baselines and advantage functions reduce this variance
Actor-critic combines value and policy learning. The actor learns the policy while the critic learns a value function; the advantage (reward minus baseline) provides low-variance gradient estimates while remaining unbiased
Interactive Lab
Toggle between a value-based agent (Q-table) and a policy-based agent (softmax over states) on the same task — see why policy gradients can express continuous and stochastic strategies that argmax-over-Q cannot.
What is the key advantage of policy gradient methods over value-based methods like DQN?
Policy gradients give us the ability to learn policies directly, but REINFORCE is too high-variance for complex problems and basic actor-critic can take unstably large policy steps. Next up: Proximal Policy Optimization (PPO) -- the algorithm that tames policy gradients and powers everything from robotics to ChatGPT.