Policy Gradient Methods: Teaching Agents to Win

When there is nothing to take an argmax over

ISRO's attitude-control problem for a geostationary communication satellite looks deceptively like the grid-world and Atari problems you have already met: a state (current orientation, angular velocity, fuel remaining), an action (fire a thruster), a reward (staying pointed at the correct patch of Earth while burning as little propellant as possible). But the action here is a continuous thrust fraction on each of several axes — not "move up" or "move down," but a real number like 0.0347 Newtons on the yaw thruster. A value-based controller, the kind you have seen elsewhere in this curriculum, works by learning Q(s,a) for every action and then acting greedily: pick argmax_a Q(s,a). That recipe requires enumerating actions. With a continuous thrust axis there are infinitely many candidate values of a, and "loop over every action and compare Q-values" is not a procedure — it is a contradiction.

Policy gradient methods sidestep the enumeration problem entirely by never computing Q(s,a) for every action in the first place. Instead they parameterize the policy itself, π_θ(a|s), as a function that outputs a probability distribution over actions directly — a Gaussian with a learned mean and variance for continuous control, or a softmax over logits for discrete choices — and then adjust θ by gradient ascent on expected return. There is no argmax step anywhere in the algorithm. This chapter derives that gradient from first principles, turns it into the REINFORCE algorithm, diagnoses why the naive version trains badly, and fixes it with actor-critic methods. Assume the MDP formalism, reward functions, and value functions V^π and Q^π are already familiar territory.

The policy gradient theorem

Let τ = (s0, a0, r0, s1, a1, r1, …) denote a trajectory, R(τ) = Σ_t γ^t r_t its discounted return, and J(θ) = E_{τ~π_θ}[R(τ)] the quantity we want to maximize. The obstacle to computing ∇θJ(θ) directly is that the expectation is taken over a distribution, P(τ;θ), that itself depends on θ — you cannot just push the gradient inside the sum without accounting for that.

The standard trick, sometimes called the score-function estimator or the log-derivative trick, rewrites the gradient of an expectation as an expectation of a gradient:

∇θ J(θ) = ∇θ Σ_τ P(τ;θ) R(τ)
        = Σ_τ P(τ;θ) · [∇θ P(τ;θ) / P(τ;θ)] · R(τ)
        = Σ_τ P(τ;θ) · ∇θ log P(τ;θ) · R(τ)
        = E_τ [ ∇θ log P(τ;θ) · R(τ) ]

The move from line one to line two just multiplies and divides by P(τ;θ); the identity ∇θ log P = ∇θP / P then collapses it into a log. The payoff is in what P(τ;θ) expands to. The probability of a specific trajectory factors as the initial-state probability, times the product of transition probabilities, times the product of action probabilities under the policy:

P(τ;θ) = P(s0) · Π_t P(s_{t+1}|s_t,a_t) · Π_t π_θ(a_t|s_t)

log P(τ;θ) = log P(s0) + Σ_t log P(s_{t+1}|s_t,a_t) + Σ_t log π_θ(a_t|s_t)

Take ∇θ of that sum. The initial-state term and every transition-probability term have no θ in them — the environment's physics do not care how the policy is parameterized — so their gradients are exactly zero. Everything the environment controls vanishes, and what survives is only the policy's own contribution:

∇θ log P(τ;θ) = Σ_t ∇θ log π_θ(a_t|s_t)

∇θ J(θ) = E_τ [ ( Σ_t ∇θ log π_θ(a_t|s_t) ) · R(τ) ]

Applying a standard causality argument — reward earned before time t cannot depend on the action chosen at time t, so those terms drop out of the sum when you distribute R(τ) across the per-timestep gradients — collapses this into the per-timestep form most implementations use, with G_t as the return from time t onward:

∇θ J(θ) = E [ Σ_t ∇θ log π_θ(a_t|s_t) · G_t ]

This is the policy gradient theorem. Replacing G_t by its expectation, Q^π(s_t,a_t), gives the textbook statement ∇θJ(θ) = E_π[∇θ log π_θ(a|s) · Q^π(s,a)], but for a Monte Carlo algorithm the sampled return G_t is what you actually have on hand after playing out an episode, and it is an unbiased estimate of Q^π(s_t,a_t) by definition.

REINFORCE and a fully worked update

REINFORCE (Williams, 1992) is exactly this equation turned into a training loop: play a full episode under the current policy, compute each timestep's return G_t, and take a gradient-ascent step on Σ_t log π_θ(a_t|s_t) · G_t. This is precisely the algorithm that underlies the self-play refinement stage of AlphaGo (Silver et al., "Mastering the game of Go with deep neural networks and tree search," Nature, 2016): after supervised pretraining on human games, AlphaGo's policy network was improved by playing full games against earlier versions of itself and applying a REINFORCE-style update, with the game outcome z ∈ {−1, +1} standing in for the return. A win nudges every move played that game to become more likely; a loss nudges every move to become less likely. The toy example below is that exact mechanism at miniature scale.

Take a one-decision episode with two actions and a softmax policy, logits θ = [θ1, θ2]. For a softmax, the score function has a clean closed form worth memorizing: ∇θ_i log π(a_k) = 𝟙[i=k] − π(a_i) — the gradient with respect to the taken action's own logit is "1 minus its probability," and with respect to every other logit it is "minus that action's probability." Start with θ = [0, 0], so π = [0.5, 0.5]. The agent samples a1 and receives reward G = +1 (a win). Then:

import numpy as np

def softmax(logits):
    exp = np.exp(logits - np.max(logits))
    return exp / exp.sum()

theta = np.array([0.0, 0.0])
alpha = 0.1
action_taken = 0        # a1
reward = 1.0             # G, the observed return

pi = softmax(theta)                      # [0.5, 0.5]
grad_log_pi = -pi.copy()                 # [-0.5, -0.5]
grad_log_pi[action_taken] += 1.0         # [0.5, -0.5]

theta = theta + alpha * reward * grad_log_pi
pi_new = softmax(theta)

print(pi, grad_log_pi, theta, pi_new)

Tracing it by hand: π = [0.5, 0.5]; grad_log_pi starts as −π = [−0.5, −0.5], then the taken action's entry gets +1, giving [0.5, −0.5], matching the closed form above. The update is θ ← [0,0] + 0.1 · 1 · [0.5, −0.5] = [0.05, −0.05]. Recomputing the softmax: e^0.05 ≈ 1.05127, e^−0.05 ≈ 0.95123, and their ratio gives π_new ≈ [0.525, 0.475]. The rewarded action's probability rose from 0.500 to 0.525 in one step — the entire mechanism of REINFORCE compressed into two numbers. Scale that up to thousands of episodes and hundreds of moves per game, and it is the same arithmetic that trained AlphaGo's policy network.

Why the naive version trains so badly

G_t in the update above is a full-episode Monte Carlo return: the sum of every future reward the agent happened to collect after that action, under whatever stochastic choices the policy and environment made for the rest of the episode. Two agents can take the identical action in the identical state and receive wildly different G_t simply because of what happened next — one opponent blunders, one does not; one dice roll favors you, one does not. That means the quantity multiplying ∇θ log π(a|s) is a noisy, high-variance random variable, and the noise grows with the length of the episode, since G_t accumulates every downstream reward and every downstream source of randomness. A gradient estimate built from a single noisy sample per state-action pair swings the parameters around erratically; averaging over many episodes eventually converges, but "eventually" can mean orders of magnitude more environment interaction than a lower-variance estimator would need.

There is a second, more subtle failure mode. If every reward in the environment happens to be positive (a common setup: +1 for progress, 0 otherwise, never negative), then G_t ≥ 0 always, and the REINFORCE update ∇θ log π(a|s) · G_t always pushes the probability of whatever action was actually sampled upward — never downward — regardless of whether that action was any good relative to the alternatives. The algorithm still works on average, because better actions get pushed up by more than worse ones, but the raw magnitude of every update is inflated by the same additive reward-scale constant, which is exactly the kind of thing that should not affect where the optimum sits.

Baseline subtraction and the advantage function

Both problems have the same fix: subtract a baseline b(s) from the return before using it as the multiplier, replacing G_t with G_t − b(s_t). The best choice of b(s) is the state-value function V^π(s), because then the multiplier becomes G_t − V(s_t), a sample of the advantage function A(s,a) = Q(s,a) − V(s): how much better than average this specific action was in this specific state. Actions that merely matched the state's baseline expectation get zero update instead of always being pushed up.

Common misconception: does subtracting a baseline change what the algorithm optimizes?

A student's first reaction is often that subtracting anything from the reward signal must bias the result toward whatever b(s) happens to be — surely you are changing the quantity being maximized. This is false, and the proof is short enough to carry around: for any function b(s) that does not depend on the action a,

E_{a~π}[ ∇θ log π(a|s) · b(s) ] = b(s) · Σ_a π(a|s) ∇θ log π(a|s)
                                = b(s) · Σ_a ∇θ π(a|s)          [since π∇logπ = ∇π]
                                = b(s) · ∇θ Σ_a π(a|s)
                                = b(s) · ∇θ (1)
                                = 0

The extra term contributes exactly zero to the expected gradient no matter what b(s) is, so it can only change variance, never the direction the parameters are being pushed in expectation. The requirement that b cannot depend on a is what makes the sum-to-one identity apply; that is also why V(s) is a legal baseline (it depends only on state) while, say, Q(s,a) itself would not be.

Actor-critic: making the baseline learn online

REINFORCE with a baseline still needs a full episode to compute G_t before any update can happen. Actor-critic methods replace the Monte Carlo return with a bootstrapped one-step estimate, the temporal-difference error, which is available after a single transition:

δ_t = r_t + γ V_w(s_{t+1}) − V_w(s_t)

δ_t is a one-sample, one-step estimate of the advantage A(s_t,a_t): it compares "what actually happened plus what we now expect from here" against "what we expected before we acted." Two functions are trained simultaneously — the actor π_θ, updated by α·δ_t·∇θ log π_θ(a_t|s_t), and the critic V_w, updated to reduce δ_t² by ordinary gradient descent (equivalently, nudged toward the TD target r_t + γV_w(s_{t+1})). Because δ_t only depends on one transition, updates happen every step, not once per episode, and the bootstrapped critic estimate has far lower variance than a full-episode return — at the cost of bias, since V_w is itself only an approximation that is wrong early in training.

Trace one full episode of two transitions through both updates. States s0 → s1 → terminal T, reward r0 = 0 on the first transition and r1 = 10 on the second, γ = 0.9. The critic currently estimates V(s0) = 2, V(s1) = 5, V(T) = 0 (terminal states are always valued at zero). Each state has its own softmax actor over two actions, both starting at logits [0, 0]. Learning rates: α_actor = α_critic = 0.1.

Transition 1 — at s0 the actor samples a1, environment returns r0 = 0 and next state s1.

δ0 = r0 + γV(s1) − V(s0) = 0 + 0.9(5) − 2 = 2.5

critic:  V(s0) ← 2 + 0.1(2.5) = 2.25

actor at s0 (a1 taken, index 0): ∇logπ = [1,0] − [0.5,0.5] = [0.5, −0.5]
θ_s0 ← [0,0] + 0.1(2.5)[0.5,−0.5] = [0.125, −0.125]

Transition 2 — at s1 the actor samples a2, environment returns r1 = 10 and terminates.

δ1 = r1 + γV(T) − V(s1) = 10 + 0.9(0) − 5 = 5

critic:  V(s1) ← 5 + 0.1(5) = 5.5

actor at s1 (a2 taken, index 1): ∇logπ = [0,1] − [0.5,0.5] = [−0.5, 0.5]
θ_s1 ← [0,0] + 0.1(5)[−0.5,0.5] = [−0.25, 0.25]

Both TD errors came out positive, meaning both transitions turned out better than the critic expected, so both actor updates increase the probability of the action actually taken (a1 at s0, a2 at s1) — the same directional logic as REINFORCE, but computed from a single transition each, immediately, without waiting for the episode to end. Note also that the two states' logits are entirely independent parameter vectors here (a tabular/per-state policy for tractability); in a real deep actor-critic network both states would share parameters through a neural network, and updates at s1 would also nudge the shared weights' predictions for s0.

Constraining how far one gradient step is allowed to move

Vanilla actor-critic and REINFORCE both take a fixed step of size α in parameter space, but a fixed step in θ-space can correspond to an enormous, destabilizing step in policy-distribution space if the softmax happens to be near-saturated — the same α that changes π by a fraction of a percent near [0.5, 0.5] can flip a near-deterministic policy's action entirely. Trust Region Policy Optimization (Schulman, Levine, Moritz, Jordan & Abbeel, ICML 2015) constrains each update to stay within a bounded KL-divergence "trust region" of the old policy, solving a constrained optimization problem at every step. Proximal Policy Optimization (Schulman, Wolski, Dhariwal, Radford & Klimov, arXiv:1707.06347, 2017) achieves a similar effect far more cheaply by clipping the probability ratio π_new(a|s)/π_old(a|s) in the surrogate objective, so a single gradient step cannot be credited for pushing an already-large probability change even further. Both are, mechanically, the same policy gradient theorem derived above with an added guardrail on step size — not a different algorithm family.

Actor-critic update loop

Active recall

Attempt every question before reading its answer.

1. Why does the policy gradient theorem let us compute ∇θJ(θ) without ever knowing ∇θ P(s'|s,a), the gradient of the environment's own transition dynamics?
2. A softmax actor has logits θ = [1, 0, −1] over three actions. Compute the policy π and the gradient ∇θ log π(a2) for the middle action (index 1, logit 0).
3. Using the gradient from question 2, suppose the observed return for taking a2 was G = −3 and α = 0.05. Compute the updated θ, and explain in one sentence why all three logits move, not just the one for a2.
4. Prove that subtracting a state-dependent baseline b(s) from the return in a REINFORCE update leaves the expected gradient unchanged.
5. In the actor-critic worked example, suppose α_critic had been 0.5 instead of 0.1, but only for the V(s0) update. (a) What is the new value of V(s0)? (b) Does this change δ0, the actor update at s0, δ1, or the actor update at s1?
6. A classmate says: "Actor-critic uses a TD error instead of a full Monte Carlo return, so it is strictly better than REINFORCE." What is wrong with this claim?

Answers

1. Expanding log P(τ;θ) into its factors — the initial-state probability, the product of transition probabilities, and the product of action probabilities — shows that only the action-probability terms contain θ. The transition-probability terms P(s'|s,a) belong to the environment and do not depend on the policy's parameters, so their gradient with respect to θ is exactly zero and they drop out of the sum entirely. What is left is Σ_t ∇θ log π_θ(a_t|s_t), which requires only the policy's own gradient, computable by ordinary backpropagation through the policy network — never the environment's dynamics.

2. e^1 ≈ 2.71828, e^0 = 1, e^−1 ≈ 0.36788, sum ≈ 4.08616. π ≈ [0.6652, 0.2447, 0.0900]. ∇θ log π(a2) = onehot(index 1) − π = [−0.6652, 0.7553, −0.0900].

3. Δθ = α·G·∇logπ = 0.05·(−3)·[−0.6652, 0.7553, −0.0900] = [0.0998, −0.1133, 0.0135]. θ_new ≈ [1.0998, −0.1133, −0.9865]. All three move because the softmax normalization couples every logit: raising or lowering one action's probability mechanically redistributes probability mass across the other two, so the score function ∇θ log π has a nonzero entry for every action, not just the one sampled — a negative return on a2 pushes θ2 down (making a2 less likely) while pushing θ1 and θ3 up (making the other two more likely), even though neither of them was chosen this step.

4. E_{a~π}[∇θ log π(a|s)·b(s)] = b(s)·Σ_a π(a|s)∇θ log π(a|s) = b(s)·Σ_a ∇θ π(a|s) = b(s)·∇θ Σ_a π(a|s) = b(s)·∇θ(1) = 0, since Σ_a π(a|s) = 1 for every θ. The baseline term's expected contribution to the gradient is zero regardless of what b(s) is, as long as it does not depend on the action a.

5. (a) V(s0) ← 2 + 0.5(2.5) = 3.25. (b) No to both. δ0 = r0 + γV(s1) − V(s0) is computed from the value estimates as they stood before this update (V(s0)=2, V(s1)=5), and the actor's update at s0 already consumed that same δ0 = 2.5, so raising α_critic changes only the stored V(s0) that future episodes will see — it does not retroactively change this step's TD error or this step's actor gradient. δ1 and the s1 actor update are untouched for an even more direct reason: their formula involves V(s1) and V(T), and neither one references V(s0) at all, so a change confined to how V(s0) is updated cannot propagate into a computation that never reads V(s0).

6. The comparison is a bias-variance tradeoff, not a strict improvement. REINFORCE's Monte Carlo return G_t is unbiased — its expectation is exactly Q^π(s,a) — but high variance. The TD error δ_t is lower variance because it only depends on one transition, but it is biased whenever the critic V_w is inaccurate, which is guaranteed early in training and whenever function approximation cannot represent V^π exactly. A poorly initialized or slow-learning critic can therefore feed the actor a systematically wrong signal that REINFORCE, despite its noise, would never produce in expectation.

Think About It

Think about this: How would you explain policy gradient methods: teaching agents to win to a friend who has never seen a computer? What real-world analogy would you use? Imagine you had to build a system using these concepts — what would be your first step? Try this: before moving on, write down three things you learned and one question you still have.

Practice Exercises

Now it is time to practice! Complete these challenges to solidify your understanding:

• Exercise 1: Write a short program that demonstrates the core concept from this chapter. Test it with at least 3 different inputs.
• Exercise 2: Find a real-world example where policy gradient methods: teaching agents to win is used in an Indian company (like TCS, Infosys, Flipkart, or ISRO). Write a paragraph explaining the connection.
• Exercise 3: Create a mind-map connecting policy gradient methods: teaching agents to win to at least 3 other topics you have studied.

Key Takeaways — Summary and Recap

Let us recap what we covered: the core ideas behind policy gradient methods: teaching agents to win, how they connect to real-world applications, and why they matter for your journey in computer science. Remember these key points as you move forward. For competitive exam preparation (CBSE, JEE, BITSAT), focus on understanding the WHY behind each concept, not just the WHAT.