REINFORCE: Policy Gradient Learning

It is the 19th over of an IPL death-overs chase. The batting side needs 14 off 12 balls, two set batters at the crease, one wicket in hand. The bowling captain has two options left in his mental model: Jasprit Bumrah, who has bowled two overs already and has two left in his quota, or Yuzvendra Chahal, rested since the 14th over. There is no formula that tells the captain, with certainty, which choice wins the match — the outcome depends on the batter's shot selection, the pitch, a dropped catch, a no-ball, a hundred things the captain cannot see into and certainly cannot differentiate. All he gets, twenty minutes later, is one bit of feedback: won or lost.

This is exactly the training signal a reinforcement learning agent gets when it has to learn a policy for choosing among discrete actions in an environment it cannot simulate through symbolically — a recommendation engine choosing which of several response strategies to try, a robot arm choosing among grasp types, or a bowling-change policy choosing among available bowlers. You cannot backpropagate through "we won the match" the way you backpropagate through a cross-entropy loss, because winning is not a differentiable function of the captain's decision — it is the output of an opaque, stochastic simulator (the rest of the match). REINFORCE, introduced by Ronald Williams in 1992, is the algorithm that makes gradient-based learning possible anyway. It is the ancestor of every modern policy-gradient method — from A2C to PPO to the policy networks inside AlphaGo and today's RLHF-tuned language models — and understanding it precisely is the difference between using those tools and merely invoking them.

Why you cannot just backpropagate through the reward

Frame the problem the way a deep learning course would. You have a policy — a function that, given a state s, outputs a probability distribution over actions — parameterized by weights θ. Write it π_θ(a | s). You want to choose θ to maximize the expected total reward earned by following this policy for a full episode (a full match, in the cricket framing). Call an episode's sequence of states and actions a trajectory τ = (s₀, a₀, s₁, a₁, …, s_T), and its total reward R(τ). The objective is:

J(θ) = E_{τ ~ π_θ}[ R(τ) ]

In ordinary supervised deep learning, you compute a loss that is a known, differentiable function of your network's output, and the chain rule carries the gradient straight back through every layer. Here, two things break that plan. First, the action a_t is not the direct, deterministic output of the network — it is a sample drawn from the distribution π_θ(· | s_t), and sampling is not a differentiable operation you can chain through. Second, and more fundamentally, R(τ) is not computed by any function you have access to at all — it comes out of the match simulator, or the real world, which is a black box with respect to θ. You cannot write ∂R/∂θ because you do not have a symbolic expression for R in terms of θ to differentiate.

The instinct many students have at this point is that some clever architecture trick will make the environment differentiable — and that instinct is the misconception worth killing directly, because REINFORCE's entire design is a refusal of that instinct. You do not need the environment to be differentiable, and you never write ∂R/∂θ. What you differentiate is your own policy's probability of having taken the action it took. The reward only ever appears as a scalar multiplier weighting that gradient — never as something you differentiate through. This single reframing is what unlocks the algorithm.

The log-derivative trick

Expand the expectation as a sum (or integral) over all possible trajectories, weighted by how likely the current policy is to produce each one:

J(θ) = Σ_τ P(τ; θ) · R(τ)

Differentiate with respect to θ. Because P(τ; θ) is the only θ-dependent term, and R(τ) is just a fixed number for a given τ, the derivative moves straight inside the sum:

grad_θ J(θ) = Σ_τ grad_θ P(τ; θ) · R(τ)

Now comes the one algebraic move that makes the whole algorithm possible. For any positive function f(θ), the identity grad f = f · grad(log f) holds, because grad(log f) = grad(f) / f by the chain rule of logarithms — just rearrange that. Apply it to P(τ; θ):

grad_θ P(τ; θ) = P(τ; θ) · grad_θ log P(τ; θ)

Substituting back:

grad_θ J(θ) = Σ_τ P(τ; θ) · grad_θ log P(τ; θ) · R(τ)
            = E_{τ ~ π_θ}[ grad_θ log P(τ; θ) · R(τ) ]

This is progress: we have converted a gradient of a sum over an intractably large space of trajectories into an expectation, which we can estimate by sampling — that is, by actually rolling out episodes with the current policy. But we still need to know what log P(τ; θ) is. A trajectory's probability factors into the probability of the starting state, the policy's choice at every step, and the environment's own transition probabilities:

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

Take the log, turning the product into a sum, then differentiate with respect to θ. The starting-state probability and the environment's transition probabilities do not depend on θ at all — the captain's policy does not change how a batter's bat makes contact with the ball — so their gradients vanish entirely, and we are left with only the terms that came from the policy itself:

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

This is the crucial cancellation: the environment dropped out of the gradient completely, without us ever needing to model, simulate, or differentiate it. Substituting back into the expectation gives the policy gradient theorem in its REINFORCE form:

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

In words: roll out an episode under the current policy, sum up the gradient of the log-probability of every action you actually took, and scale that whole sum by however much total reward the episode earned. Do that many times and average, and you have an unbiased estimate of the true gradient of expected reward — computed without ever touching a derivative of the environment.

Turning the formula into an algorithm

For a softmax policy over a discrete action set — exactly the case of the captain choosing among a handful of available bowlers — write the logits as θ_a (one parameter per action, ignoring state-dependence for a moment to keep the algebra visible) so that π_θ(a) = exp(θ_a) / Σ_b exp(θ_b). Differentiating log π_θ(a) with respect to a specific logit θ_i gives a clean, well-known result:

d/dθ_i [ log π_θ(a) ] = 1[i = a] − π_θ(i)

Read this as: the gradient pushes the logit of the action you actually took up by (1 minus its current probability), and pushes every other action's logit down by its own current probability — regardless of whether that other action was ever sampled. That coupling, driven purely by the softmax normalization, is going to matter in the worked example below. The full algorithm, using return-to-go G_t (defined in the next section) instead of the cruder full-episode return, is:

Algorithm: REINFORCE
Initialize policy parameters theta arbitrarily (e.g. zero)
repeat for many episodes:
    roll out one full episode s_0, a_0, r_1, s_1, a_1, r_2, ..., s_T
      by sampling a_t from pi_theta(. | s_t) at every step
    for t = 0 to T-1:
        G_t = sum of r_{t+1} through r_T   (return-to-go from step t)
    theta = theta + alpha * sum_t [ grad_theta log pi_theta(a_t | s_t) * G_t ]

The mechanism, end to end

Worked example: one episode, one update

Take the softmax policy from the diagram, reduced to its simplest possible form so every number can be checked by hand: two actions (Bumrah = action 0, Chahal = action 1), two logits θ = (θ_B, θ_C), initialized to (0, 0), giving π(B) = π(C) = 0.5 before any learning happens. To keep the arithmetic self-contained, treat the state as unchanging across the two overs bowled (a stationary policy reused at every step, which is exactly what a weight-shared network does across timesteps).

Episode: over 19 — sampled action is Bumrah. Over 20 — sampled action is Bumrah again. The team wins: R(τ) = +1 (a single terminal reward, so G₀ = G₁ = R(τ) = 1 here — the return-to-go and full-episode return coincide because nothing rewarding happens until the very end).

Using the softmax gradient identity derived above, grad log π(a) at each timestep is the vector (1[a = B] − π(B), 1[a = C] − π(C)). Since π(B) = π(C) = 0.5 throughout (θ has not changed yet), and the action taken at both steps is Bumrah:

step t=0 (over 19, action=Bumrah):  grad log pi = (1 - 0.5, 0 - 0.5) = (0.5, -0.5)
step t=1 (over 20, action=Bumrah):  grad log pi = (1 - 0.5, 0 - 0.5) = (0.5, -0.5)

sum over both steps: (0.5+0.5, -0.5-0.5) = (1.0, -1.0)

The REINFORCE gradient estimate for this single episode is this sum scaled by the return: (1.0, −1.0) × 1 = (1.0, −1.0). With learning rate α = 0.1:

theta_B = 0 + 0.1 * 1.0  =  0.1
theta_C = 0 + 0.1 * (-1.0) = -0.1

Recomputing the softmax: exp(0.1) ≈ 1.10517, exp(−0.1) ≈ 0.90484, sum ≈ 2.01001, so π(B) ≈ 1.10517 / 2.01001 ≈ 0.5498 and π(C) ≈ 0.4502. A single winning episode with Bumrah bowling both overs nudges the policy from a 50/50 coin flip to roughly 55/45 in Bumrah's favour — exactly the direction intuition demands, produced entirely mechanically by the log-derivative trick, with no hand-coded "reward Bumrah" rule anywhere in the algorithm.

Here is that exact update as runnable code, with the sampling step replaced by the fixed actions from the worked episode so the numbers can be checked line by line:

import numpy as np

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

theta = np.array([0.0, 0.0])     # logits: [Bumrah, Chahal]
alpha = 0.1
actions_taken = [0, 0]           # over 19 -> Bumrah(0), over 20 -> Bumrah(0)
R = 1.0                          # team won

grad = np.zeros(2)
for a in actions_taken:
    probs = softmax(theta)       # still [0.5, 0.5]; theta not updated mid-episode
    onehot = np.zeros(2)
    onehot[a] = 1.0
    grad += (onehot - probs)     # d log pi(a) / d theta

theta = theta + alpha * R * grad
print(theta)          # approximately [0.1, -0.1]
print(softmax(theta)) # approximately [0.5498, 0.4502]

Tracing it: the loop runs twice, and because θ is only updated after the loop (matching the algorithm — you finish the whole episode before touching the parameters), probs is [0.5, 0.5] on both iterations, so grad accumulates to exactly [1.0, −1.0], matching the by-hand derivation above.

Why an unlucky bowler's teammate benefits too

Run the same episode but with the opposite outcome to see a subtlety that trips students up. Suppose instead over 19 and over 20 are both bowled by Chahal, and the team loses: R(τ) = −1. Then grad log π(a=Chahal) at each step is (0 − 0.5, 1 − 0.5) = (−0.5, 0.5), summing to (−1.0, 1.0), and the update is α · R · grad = 0.1 × (−1) × (−1.0, 1.0) = (0.1, −0.1). θ_B rises to 0.1 and θ_C falls to −0.1 — Bumrah's selection probability goes up to roughly 55%, even though Bumrah did not bowl a single ball in this losing match. This is not a bug. Softmax is a competitive, normalized distribution: pushing one logit down is mathematically identical to pushing every other logit up in relative terms, because the probabilities must still sum to one. REINFORCE never reasons about Bumrah individually in this episode; it only ever computes ∂ log π(a_t)/∂θ for the action actually sampled, and the softmax's own structure does the rest.

Reducing variance: reward-to-go and baselines

The worked example used a terminal-only reward, so return-to-go and full-episode return were identical. In general they are not, and the difference matters. If rewards arrive throughout the episode — a small penalty for every run conceded, say, on top of the terminal win/loss bonus — then weighting the over-19 log-probability term by the reward earned in over 14, which happened before over 19 was even chosen, adds pure noise: that earlier reward could not possibly have been caused by a decision made later. The causality (or reward-to-go) trick fixes this by weighting grad log π(a_t | s_t) only by rewards from step t onward, G_t = r_t+1 + r_t+2 + … + r_T, exactly as written in the algorithm box above. Formally, the terms involving rewards realized strictly before t have zero expectation once multiplied by grad log π(a_t | s_t), since those rewards are independent of a decision made later — dropping them removes variance from the gradient estimate without introducing any bias.

A second, complementary fix subtracts a baseline b(s_t) — often an estimate of the average return achievable from that state — from G_t before multiplying by the log-probability gradient. This matters in practice because a return that is always positive (say, every episode nets some reward between 0 and 1) makes the raw REINFORCE update always push every sampled action's probability up, never down, which slows learning to a crawl; subtracting a baseline recenters the signal so genuinely below-average actions get pushed down. The proof that this does not bias the gradient is short and worth doing once: for any state s, E_a~π[grad log π(a|s) · b(s)] = b(s) · Σ_a π(a) · grad log π(a) = b(s) · Σ_a grad π(a) = b(s) · grad(Σ_a π(a)) = b(s) · grad(1) = 0. Since probabilities always sum to 1, that sum's gradient is exactly zero regardless of what b(s) is, so subtracting any state-dependent baseline changes only the variance of the estimate, never its expected value.

Active recall

Attempt every question before reading its answer.

1. Why can gradient descent not be applied directly to J(θ) the way it is applied to a supervised loss?
2. Starting from grad P(τ;θ) = P(τ;θ) · grad log P(τ;θ), show the two algebraic steps that turn grad_θ J(θ) into E_τ[grad_θ log P(τ;θ) · R(τ)].
3. With θ = (0, 0), the captain plays Chahal for both overs 19 and 20, and the team loses (R = −1). Compute the updated θ and the new π(Bumrah) at α = 0.1.
4. What is the reward-to-go trick, and why does dropping the rewards earned before step t not bias the gradient estimate?
5. Prove that subtracting a baseline b(s) from G_t leaves the REINFORCE gradient estimate unbiased.
6. If Bumrah keeps being associated with wins and Chahal with losses over many episodes, what happens to the policy as training continues, and what practical problem does that create?

Worked answers

1. Because two links in the chain rule are missing: the sampled action a_t is drawn from a distribution rather than computed as a deterministic differentiable function of θ, and the reward R(τ) is produced by the environment — a black-box simulator or the real world — for which no symbolic expression in terms of θ exists. There is nothing to differentiate through.
2. Step one: divide and multiply by P(τ;θ) inside the sum, i.e. grad_θJ(θ) = Σ_τ grad_θP(τ;θ)·R(τ) = Σ_τ P(τ;θ)·[grad_θP(τ;θ)/P(τ;θ)]·R(τ) = Σ_τ P(τ;θ)·grad_θlog P(τ;θ)·R(τ). Step two: recognize this sum, weighted by P(τ;θ), as an expectation over trajectories sampled from the current policy: E_{τ~π_θ}[grad_θlog P(τ;θ)·R(τ)].
3. grad log π(Chahal) at each step is (1[Chahal=B]−0.5, 1[Chahal=C]−0.5) = (0−0.5, 1−0.5) = (−0.5, 0.5); summed over both steps, (−1.0, 1.0). The update is α·R·grad = 0.1×(−1)×(−1.0,1.0) = (0.1, −0.1), so θ_B=0.1, θ_C=−0.1. Softmax gives π(Bumrah) = exp(0.1)/(exp(0.1)+exp(−0.1)) ≈ 1.10517/2.01001 ≈ 0.5498 — about 55%, up from 50%, even though Bumrah never bowled that match.
4. Reward-to-go replaces the full-episode return R(τ) with G_t = r_t+1+…+r_T, the sum of rewards from step t onward only. Rewards earned before step t are statistically independent of the action taken at t (a decision cannot affect the past), so their contribution to E[grad log π(a_t|s_t)·r_earlier] is exactly zero — including them adds variance to the estimate without shifting its expected value, so dropping them is a pure variance reduction, not an approximation.
5. E_a~π[grad log π(a|s)·b(s)] = b(s)·Σ_a π(a)·grad log π(a) = b(s)·Σ_a grad π(a) (using grad π = π·grad log π) = b(s)·grad(Σ_a π(a)) = b(s)·grad(1) = 0, because probabilities sum to 1 for every value of θ, so the gradient of that sum is identically zero. Since the extra term contributes zero in expectation, subtracting b(s) leaves the gradient estimate unbiased while typically reducing its variance.
6. θ_B keeps growing relative to θ_C with every winning episode that used Bumrah, so π(Bumrah) drifts toward 1 and π(Chahal) toward 0 — the policy becomes increasingly deterministic. The practical problem is exploration collapse: once π(Chahal) is near zero, the policy almost never samples Chahal again, so it can never gather the evidence needed to discover that Chahal might in fact be the better choice in some other match state. Production policy-gradient methods counter this with an entropy bonus added to the objective, or a slowly decaying exploration rate, to keep probability mass from collapsing prematurely.

Think About It

Think about this: How would you explain reinforce: policy gradient learning 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.

Key Takeaways — Summary and Recap

Let us recap what we covered: the core ideas behind reinforce: policy gradient learning, 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.