REINFORCE algorithm: Visual walkthrough — AI-ML deep dive

Prerequisite trails, if a symbol here feels unfamiliar: Policy Gradient Methods, Log-Derivative Trick, Monte Carlo Methods, Value Function V(s), Advantage Function.

The characters (define every symbol before we use it)

Before Step 1, meet the cast. Nothing below is used until it appears in a picture.

Definition Words first, symbols second

A state $s$ — the situation the agent is in (a picture of the world right now).
An action $a$ — a move the agent can make.
A reward $r$ — a number the environment hands back after a move: bigger = better.
A policy $π_{θ} (a ∣ s)$ — a dial-controlled recipe that reads " $a$ 's probability, given we are in state $s$ ". The little $θ$ (theta) is the bundle of tunable dials (e.g. the weights of a neural net). Turn the dials → the recipe's probabilities change.
A trajectory $τ$ (tau) — one full played-out episode: $s_{0}, a_{0}, r_{0}, s_{1}, a_{1}, r_{1}, \dots$ . Think "one game, start to finish."
The return $R (τ) = \sum_{t} γ^{t} r_{t}$ — the total (discounted) reward of that game. $γ$ (gamma), between 0 and 1, means "reward now is worth a bit more than reward later."
$E_{τ \sim π_{θ}} [\cdot]$ — the expectation (the average) of a quantity over all the games our current policy tends to produce.
$J (θ) = E_{τ} [R (τ)]$ — the objective: the average return we get if we keep playing with policy $π_{θ}$ . This is the single number we want to push up.
$\nabla_{θ}$ (nabla-theta) — the gradient: the arrow in dial-space pointing in the direction that increases whatever follows it, fastest. We will draw this literally as an arrow.

Two more characters belong to the environment (the world), not to us — they carry no dials $θ$ and we will meet them again in Step 5:

Our whole quest: find $\nabla_{θ} J$ — the arrow telling us which way to turn the dials to earn more reward.

Step 1 — See the landscape we are climbing

WHAT. Picture the dials $θ$ as a flat floor. Above each spot on the floor, the height is $J (θ)$ — how much reward that dial-setting earns on average. This makes a hill.

WHY. "Maximize $J$ " is now just "walk uphill." The gradient $\nabla_{θ} J$ is the compass that points straight up the slope from wherever we stand. Everything after this step is a fight to compute that compass arrow — because the hill's height involves an average over random games, we cannot read the slope directly.

PICTURE. The green surface is $J$ . The black dot is our current $θ$ . The orange arrow is $\nabla_{θ} J$ , the uphill direction we are hunting for.

Step 2 — Why we can't just differentiate the average

WHAT. Write the average out honestly. An average over games is a probability-weighted sum over all possible games: $J (θ) = \int how likely game τ is p_{θ} (τ) its total reward R (τ) d τ$

WHY. $\int \dots d τ$ is just "add up over every possible game, weighting each by how probable it is." Here is the catch that motivates the whole trick: the weights $p_{θ} (τ)$ themselves depend on our dials $θ$ . Turning a dial doesn't only change the reward of a fixed game — it changes which games happen. So the slope has two intertwined effects, and a naive derivative gets tangled.

PICTURE. Bars = possible games. Bar height = reward $R (τ)$ . Bar width = probability $p_{θ} (τ)$ . $J$ is the shaded area. When we turn a dial, the widths shift — that's the hard part.

Step 3 — Move the gradient inside the sum

WHAT. The gradient of a sum is the sum of the gradients (differentiation is linear — it distributes over addition). And $R (τ)$ has no $θ$ in it (rewards are decided by the environment, not our dials). So only the widths get differentiated: $\nabla_{θ} J = \int \nabla_{θ} p_{θ} (τ) R (τ) d τ$

WHY. This isolates the entire difficulty into one object: $\nabla_{θ} p_{θ} (τ)$ , "how fast does game $τ$ 's probability grow as we turn the dials." If we can rewrite that as something we can sample, we win.

PICTURE. Same bars as Step 2, now with small arrows on each bar showing which way its width is changing as $θ$ nudges. Heights (rewards) stay frozen.

Step 4 — The log-derivative trick: turn "gradient of probability" into "probability × gradient of log"

WHAT. Use the identity (proven in Log-Derivative Trick): $\nabla_{θ} p_{θ} (τ) = p_{θ} (τ) \nabla_{θ} lo g p_{θ} (τ)$

WHY this exact tool, not another? We want an expectation — a shape $\int p_{θ} (τ) [something] d τ$ — because expectations can be estimated by sampling (play games, average). Right now we have $\int (\nabla p) R d τ$ , which is not in that shape: there is no clean $p_{θ}$ multiplying a bracket. The identity manufactures that missing $p_{θ}$ out front. Where does the identity come from? Calculus: $\frac{d}{d θ} lo g p = \frac{1}{p} \frac{d p}{d θ}$ , so $\nabla lo g p = \frac{\nabla p}{p}$ ; multiply both sides by $p$ . That's it.

Substitute and the shape appears: $\nabla_{θ} J = \int p_{θ} (τ) \nabla_{θ} lo g p_{θ} (τ) R (τ) d τ = E_{τ} [\nabla_{θ} lo g p_{θ} (τ) R (τ)]$

PICTURE. Left box: the tangled form $\nabla p \cdot R$ . A labelled arrow "×1 in disguise: multiply and divide by $p$ ". Right box: the clean $p \cdot (\nabla lo g p) \cdot R$ , tagged "now it's an average — samplable!"

Step 5 — Crack open $lo g p_{θ} (τ)$ : the environment vanishes

WHAT. A game's probability is a chain of two kinds of factors — our choices $π_{θ}$ and the world's own pieces (the initial-state distribution $ρ (s_{0})$ and the transition kernel $P (s_{t + 1} ∣ s_{t}, a_{t})$ , both defined in "The characters" above): $p_{θ} (τ) = ρ (s_{0}) \prod_{t} we choose π_{θ} (a_{t} ∣ s_{t}) world responds P (s_{t + 1} ∣ s_{t}, a_{t})$ Take $lo g$ (a product becomes a sum, because $lo g (x y) = lo g x + lo g y$ ): $lo g p_{θ} (τ) = lo g ρ (s_{0}) + \sum_{t} lo g π_{θ} (a_{t} ∣ s_{t}) + \sum_{t} lo g P (s_{t + 1} ∣ s_{t}, a_{t})$

WHY. Now differentiate w.r.t. $θ$ . The start-state term $ρ$ and every world-response $P$ contain no $θ$ — the environment doesn't care about our dials. Their gradients are exactly zero. Only our own choices survive: $\nabla_{θ} lo g p_{θ} (τ) = \sum_{t} \nabla_{θ} lo g π_{θ} (a_{t} ∣ s_{t})$ This is the punchline that makes REINFORCE model-free: we never need to know the world's physics $P$ .

PICTURE. The chain of factors laid in a row. Blue tiles = $π_{θ}$ (kept). Gray tiles = $ρ, P$ (crossed out, "no $θ$ inside → gradient 0").

Plugging back gives the raw REINFORCE gradient: $\nabla_{θ} J = E_{τ} [(\sum_{t} \nabla_{θ} lo g π_{θ} (a_{t} ∣ s_{t})) R (τ)]$

Step 6 — Causality: replace $R (τ)$ with reward-to-go $G_{t}$

WHAT. Inside the sum, each term $\nabla_{θ} lo g π_{θ} (a_{t} ∣ s_{t})$ is currently multiplied by the whole-game reward $R (τ)$ . Replace it with only the reward that comes at or after time $t$ : $G_{t} = \sum_{k = t}^{T} γ^{k - t} r_{k} (“reward-to-go”)$

WHY. An action $a_{t}$ cannot reach back in time and change rewards $r_{0}, \dots, r_{t - 1}$ that already happened. Those past rewards contribute zero on average to $a_{t}$ 's gradient term — they are pure noise. Dropping them keeps the gradient's direction identical in expectation while shrinking its random jitter (its variance). This is the causality argument; more on why lowering variance matters in Variance Reduction in RL.

PICTURE. A timeline. For action $a_{t}$ (blue), rewards before it are grayed out ("can't affect these"), rewards from $t$ onward are orange and summed into $G_{t}$ .

Step 7 — From a true average to a sample estimate (Monte Carlo)

WHAT. We can't compute the exact average $E_{τ} [\cdot]$ (infinitely many possible games). So we play $N$ games and average — the Monte Carlo idea: $\nabla_{θ} J \approx \frac{1}{N} \sum_{i = 1}^{N} \sum_{t = 0}^{T_{i}} \nabla_{θ} lo g π_{θ} (a_{t}^{i} ∣ s_{t}^{i}) G_{t}^{i}$

WHY. Step 4 was worth all its effort precisely so we could reach this line: an average over sampled games is directly computable. More games $N$ → the noisy estimate settles onto the true gradient.

PICTURE. Three sampled trajectories, each producing its own $\sum_{t} \nabla lo g π G_{t}$ ; the three arrows average into one (bold) estimate of $\nabla_{θ} J$ .

Step 8 — Degenerate & edge cases (the ones people trip on)

WHAT & WHY & PICTURE, four situations the formula must survive:

All returns large & positive (e.g. every $G_{t} \approx 500$ ). Then every $\nabla lo g π$ is pushed up — good and bad actions alike. Learning crawls. The fix is a baseline $b (s)$ : use $G_{t} - b (s_{t})$ so the signal becomes "better or worse than average." Because $E_{a} [\nabla_{θ} lo g π_{θ} (a ∣ s) b (s)] = b (s) \nabla_{θ} \sum_{a} π_{θ} (a ∣ s) = b (s) \cdot \nabla_{θ} 1 = 0$ , the baseline adds zero bias. Choosing $b (s) = \hat{V} (s)$ (see Value Function V(s)) turns $G_{t} - b (s_{t})$ into the advantage $A_{t}$ , the bridge to Actor-Critic Methods.
Zero return, $G_{t} = 0$ . That action's term contributes nothing — neither pushed up nor down. Correct: no evidence, no update.
Deterministic-ish policy, $π_{θ} (a ∣ s) \to 1$ . Then $lo g π \to 0$ and $\nabla lo g π \to 0$ : the gradient dries up. This is why we keep policies stochastic during training (e.g. entropy bonuses) — a fully confident policy stops learning.
Sign of the update. Optimizers minimize, but we want to maximize $J$ . So we code the loss as $L = - \sum_{t} lo g π_{θ} (a_{t} ∣ s_{t}) A_{t}$ ; minimizing $L$ ascends $J$ . Forget the minus and you climb downhill.

The one-picture summary

Everything above, compressed: from the hill $J (θ)$ (Step 1), through the log-trick reshaping (Step 4), the environment cancelling (Step 5), causality trimming to $G_{t}$ (Step 6), and Monte-Carlo sampling (Step 7) — landing on the boxed gradient and the ascent update.

Recall Feynman: retell the whole walkthrough in plain words

We want to turn some dials so a game pays out more on average. The trouble: turning a dial changes which games happen, so we can't just read the slope. The log-derivative trick rewrites "how the game-probabilities shift" as "the probability of the game times how the log-probability shifts" — and that magic reshaping makes the whole thing an average, which we can estimate just by playing games and averaging. When we open up a game's probability, it's our choices multiplied by the world's responses; the world's part has none of our dials in it, so it disappears when we take the gradient — that's why we never need to know how the world works. Then causality says: judge each move only by what happened afterward ( $G_{t}$ ), never before. Finally, play $N$ games, average the "nudge each action by how much reward followed it," and step the dials that way. Add a baseline (the average day) so we push good moves up and bad moves down instead of pushing everything up. Mind the minus sign, because the optimizer walks downhill and we want up.

Active Recall

Recall Why is the log-derivative trick used at all?

To turn $\int (\nabla p) R d τ$ — which is not an average — into $E [\nabla lo g p \cdot R]$ , which is an average and can be estimated by sampling games. ::: It manufactures the missing $p_{θ}$ factor so the integral becomes an expectation.

Recall Which terms of

$lo g p_{θ} (τ)$ survive the gradient, and why? Only $\sum_{t} \nabla_{θ} lo g π_{θ} (a_{t} ∣ s_{t})$ . ::: The start-state $ρ$ and transition $P$ contain no $θ$ , so their gradients are zero → REINFORCE is model-free.

Recall Why replace

$R (τ)$ with $G_{t}$ ? Because an action can't change past rewards; those contribute zero in expectation. ::: Dropping them keeps the direction unbiased and lowers variance (causality).

Recall Does subtracting a baseline

$b (s)$ change the gradient's expected value? No. ::: $E_{a} [\nabla lo g π \cdot b (s)] = b (s) \nabla \sum_{a} π = b (s) \cdot 0 = 0$ — zero bias, only variance shrinks.