Visual walkthrough — Policy gradient methods

This is the visual companion to Policy gradient methods. Read that first for the why; read this for the how it actually falls out.

---

Step 0 — The words and symbols we must agree on first

Before any math: a few plain-word ideas, each pinned to a picture.

• A policy is a dial-controlled coin. Given a situation (a state $s$), it flips a weighted coin to choose an action $a$. The weights on that coin are set by a bunch of numbers we call parameters $\theta$ (think: the settings of the dials on a machine). We write ${\pi }_{\theta }(a\mid s)$ = "the chance the machine picks action $a$ when it sees state $s$, given the dial settings $\theta$."
• A game runs in discrete time steps $t=0,1,2,\dots$. At step $t$ the machine sees state $s_t$, picks action $a_t$, and the environment hands back a reward at that step, written $r_t$ (a time-indexed treat — one number per step). This subscript $t$ is important: $r_3$ is the treat collected at step 3 specifically.
• The game ends after a final step $T$, called the horizon (the episode length — how many steps one game lasts). So a game touches steps $t=0$ up to $t=T$.
• The very first state $s_0$ has to come from somewhere — the environment drops us into a starting situation. The rule for that drop is a fixed probability list $\rho (s_0)$, the initial-state distribution ("how likely each possible start is"). It is set by the environment, not by our dials $\theta$ — we don't get to choose where the game begins.
• A trajectory $\tau$ is one whole game played start to finish: $\tau =(s_0,a_0,r_0,s_1,a_1,r_1,\dots ,s_T,a_T,r_T)$. It is exactly the list of states, actions and step-rewards from $t=0$ to $t=T$.
• The total treat of that one game is the sum of its step-rewards: $R(\tau )={\sum }_{t=0}^T{r}_t$. So $R(\tau )$ and the individual $r_t$ are the same treats, just packaged whole vs. step-by-step.

Keep that last sentence in mind — it is the whole puzzle.

---

Step 1 — WHAT are we even trying to maximize?

WHAT. We want dial settings $\theta$ that make high-treat games happen often. So we score $\theta$ by the average treat over all the games those dials tend to produce: $J(\theta )=\underset{\text{average over games the dials produce}}{\underbrace{{\mathbb{E}}_{\tau \sim {\pi }_{\theta }}}}\left[\underset{\text{treat of one game}}{\underbrace{R(\tau )}}\right],R(\tau )={\sum }_{t=0}^T{r}_t.$

WHY the average (expectation) symbol ${\mathbb{E}}_{\tau \sim {\pi }_{\theta }}$? Using the definition just above: the coin is random — the same dials give different games each time. We don't care about one lucky game; we care about the typical treat. ${\mathbb{E}}_{\tau \sim {\pi }_{\theta }}[\cdot ]$ reads "the long-run mean of the thing in brackets, when games $\tau$ are drawn using policy ${\pi }_{\theta }$." We carry that subscript $\tau \sim {\pi }_{\theta }$ on every expectation on this page — it names which random games we average over, and it is never anything else.

PICTURE. Below: each dial setting $\theta$ produces a whole cloud of possible games (dots), some tall-treat, some short-treat. $J(\theta )$ is the height of that cloud's centre of mass. Turning the dials slides and reshapes the cloud.

Figure s01 — Alt text: a scatter of magenta dots (games $\tau$) drawn from dial settings $\theta$; a dashed violet line marks the cloud's mean treat $J(\theta )$; an orange arrow shows that turning the dials slides the whole cloud upward.

Our job: push the cloud's centre of mass upward by turning dials. That means we need the direction to turn them — a gradient.

---

Step 2 — WHY a gradient, and where's the trouble?

WHAT. To improve $J$ we ask: "if I nudge each dial a tiny bit, does the average treat go up or down, and how fast?" That question is the gradient ${\nabla }_{\theta }J$ — an arrow in dial-space pointing uphill. We then step that way: $\theta \leftarrow \theta +\alpha {\nabla }_{\theta }J(\theta ),$ where $\alpha$ (the learning rate) is just the size of the step. The ${\nabla }_{\theta }$ symbol ("nabla") means "collect the slope of $J$ with respect to every dial into one arrow." We go up (ascent) because $J$ is a treat we want more of.

WHY it's hard. Recall from the boxed definition that ${\mathbb{E}}_{\tau \sim {\pi }_{\theta }}[R(\tau )]$ is by definition the weighted sum ${\sum }_{\tau }{P}_{\theta }(\tau )R(\tau )$, where $P_{\theta }(\tau )$ = "chance the dials produce game $\tau$." Trajectories are actually discrete lists (Step 0), so this is a genuine sum over all possible games: $J(\theta )={\sum }_{\tau }{P}_{\theta }(\tau )R(\tau ).$ When actions or states can vary continuously the sum becomes an integral, which we write $\int P_{\theta }(\tau )R(\tau )d\tau$; here "$\int \cdots d\tau$" means exactly the same thing as "${\sum }_{\tau }$" — add over all possible games — and "$d\tau$" is the continuous analogue of "one term per game." Everything below reads identically whether you keep the ${\sum }_{\tau }$ or the $\int d\tau$; we use $\int$ only so the formulas match the parent note. Either way the snag is the same: $J(\theta )=\int \underset{\text{how likely game }\tau }{\underbrace{P_{\theta }(\tau )}}\underset{\text{its treat}}{\underbrace{R(\tau )}}d\tau .$ $\theta$ sits inside $P_{\theta }$ — how likely a game is — but not inside $R$, the treat. The dials never touch the reward; they only re-weight which games show up.

PICTURE. Two bars: $P_{\theta }$ (blue, moves when we turn dials) and $R$ (orange, frozen — the environment owns it). We must differentiate through the movable one only.

Figure s02 — Alt text: two bars side by side; a magenta bar labelled $P_{\theta }(\tau )$ "moves with dials," an orange bar labelled $R(\tau )$ carries a lock icon "frozen by environment, no $\theta$."

---

Step 3 — Differentiate, but only the movable part

WHAT. Take the gradient of the sum. Sum-of-things → sum-of-gradients, and since $R(\tau )$ is a frozen constant (no $\theta$ in it) it just rides along: ${\nabla }_{\theta }J=\int \underset{\text{only this moves}}{\underbrace{{\nabla }_{\theta }{P}_{\theta }(\tau )}}R(\tau )d\tau .$

WHY. This is the pay-off of Step 2's picture: the orange bar was frozen, so ${\nabla }_{\theta }$ passes straight over $R(\tau )$ and lands entirely on $P_{\theta }(\tau )$. This kills the tempting mistake "$R$ has no $\theta$ so the whole gradient is zero" — the movement lives in how likely each game is, not in the treats.

PICTURE. Same two bars; a magenta arrow of change sits on the $P_{\theta }$ bar only, the $R$ bar shows a lock icon.

Figure s03 — Alt text: the same two bars; a violet upward arrow labelled ${\nabla }_{\theta }{P}_{\theta }$ rises from the magenta probability bar, while the orange reward bar stays locked and untouched.

But ${\nabla }_{\theta }{P}_{\theta }(\tau )$ is not an average of anything — we can't estimate it by playing games. We need to turn it back into an average. Enter the trick.

---

Step 4 — The log-derivative trick (the heart)

WHAT. Use one identity from calculus: ${\nabla }_{\theta }\log f=\frac{{\nabla }_{\theta }f}{f}⟹{\nabla }_{\theta }f=f{\nabla }_{\theta }\log f.$ Apply it with $f=P_{\theta }(\tau )$: ${\nabla }_{\theta }{P}_{\theta }(\tau )=\underset{\text{a probability again!}}{\underbrace{P_{\theta }(\tau )}}\underset{\text{a score direction}}{\underbrace{{\nabla }_{\theta }\log P_{\theta }(\tau )}}.$

WHY this exact tool and not another? We needed a $P_{\theta }(\tau )$ to reappear out front, because $\int P_{\theta }(\tau )(\text{stuff})d\tau$ is by definition an average over games — the one thing we can estimate by simply playing. The log-derivative identity is the unique move that pulls a $P_{\theta }$ back out front while leaving behind something differentiable ($\log$). This is the Log-derivative trick / Score function estimator.

Substitute: ${\nabla }_{\theta }J=\int P_{\theta }(\tau ){\nabla }_{\theta }\log P_{\theta }(\tau )R(\tau )d\tau ={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}[{\nabla }_{\theta }\log P_{\theta }(\tau )R(\tau )].$ Read the boxed idea: gradient of a probability → average of (score × reward), which we can sample.

PICTURE. A flow: the "un-samplable" $\nabla P$ object passes through the log-trick gate and comes out as "$P\times \nabla \log P$" — the left half is a probability (a coin we can flip = sample), the right half a slope we can compute.

Figure s04 — Alt text: a left box "${\nabla }_{\theta }{P}_{\theta }$ (can't sample)" flows through an orange "log-derivative gate" into a right box "$P_{\theta }{\nabla }_{\theta }\log P_{\theta }$ (sample × slope)"; the identity ${\nabla }_{\theta }f=f{\nabla }_{\theta }\log f$ is written below.

---

Step 5 — Crack open $\log P_{\theta }(\tau )$: the dynamics vanish

WHAT. A whole game's probability is a chain of independent choices multiplied together: $P_{\theta }(\tau )=\underset{\text{where we start}}{\underbrace{\rho (s_0)}}{\prod }_{t=0}^{T-1}\underset{\text{world’s physics}}{\underbrace{P(s_{t+1}\mid s_t,a_t)}}\underset{\text{our coin}}{\underbrace{{\pi }_{\theta }(a_t\mid s_t)}}.$ The product runs $t=0,1,\dots ,T-1$ (each step feeds the next state $s_{t+1}$, so it stops one short of the horizon $T$). Take $\log$ (which turns a product into a sum, because $\log (xy)=\log x+\log y$): $\log P_{\theta }(\tau )=\log \rho (s_0)+{\sum }_{t=0}^{T-1}\log P(s_{t+1}\mid s_t,a_t)+{\sum }_{t=0}^{T-1}\log {\pi }_{\theta }(a_t\mid s_t).$ Now ${\nabla }_{\theta }$. The starting term $\log \rho (s_0)$ (environment's, defined in Step 0) and the physics term $\log P(s_{t+1}\mid s_t,a_t)$ contain no dial $\theta$, so their slope is zero. Only our coin survives: $\boxed{{\nabla }_{\theta }\log P_{\theta }(\tau )=\sum _{t=0}^{T-1}{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)}$

WHY it matters enormously. The environment's physics $P(s_{t+1}\mid s_t,a_t)$ — which we usually don't know — evaporated, and so did the start distribution $\rho (s_0)$. We never need a model of the world. This is what "model-free" means, and it's why policy gradients work in games where the rules are a black box.

PICTURE. The trajectory drawn as alternating boxes: state, action, state, action. Colour the action-coins magenta ("$\theta$ lives here, keep") and the physics-arrows grey ("no $\theta$, slope 0, delete"). The kept magenta pieces sum into the score.

Figure s05 — Alt text: a trajectory as alternating state boxes and magenta action coins; grey "physics" arrows between states are marked $\nabla =0$ and deleted, while the magenta ${\nabla }_{\theta }\log {\pi }_{\theta }$ coins are kept and summed.

---

Step 6 — Assemble the Policy Gradient Theorem

WHAT. Drop the boxed sum from Step 5 into the average from Step 4: ${\nabla }_{\theta }J(\theta )={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}\left[({\sum }_{t=0}^{T-1}{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t))R(\tau )\right].$ Term by term, right where each lives:

• ${\mathbb{E}}_{\tau \sim {\pi }_{\theta }}[\cdot ]$ — "play many games drawn from policy ${\pi }_{\theta }$ and average"; we estimate it with $N$ real rollouts.
• ${\sum }_{t=0}^{T-1}$ — sum over every action-step of the game, from start $t=0$ to the last action before the horizon $T$.
• ${\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)$ — the score at step $t$: the nudge that makes the action we actually took more probable.
• $R(\tau )={\sum }_{t=0}^T{r}_t$ — the treat weight: how hard to push. Big treat → push all this game's actions up hard; negative treat → push them down.

WHY this is beautiful. It says: reinforce every choice in a game in proportion to how good that game turned out. Exactly the dog-and-treats intuition, now exact.

PICTURE. Two sampled games side by side: a high-treat game (magenta, thick "make more likely" arrows on each action) and a low-treat game (violet, thin arrows). The theorem = average of these pushes.

Figure s06 — Alt text: two sampled games; a high-treat game with thick magenta upward arrows on each action, and a low-treat game with thin violet arrows; the theorem formula is written between them.

---

Step 7 — Edge case A: causality trims the reward (and where the discount enters)

WHAT. In Step 6 every action is weighted by the whole game's treat $R(\tau )={\sum }_{t=0}^T{r}_t$, including treats collected before that action happened. But an action at time $t$ cannot change the past. So the honest weight for the step-$t$ score is only the reward-to-go — the step-rewards $r_{t^{\prime }}$ from $t$ onward: ${\nabla }_{\theta }J={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}\left[{\sum }_{t=0}^{T-1}{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)G_t\right],G_t={\sum }_{t^{\prime }=t}^T{\gamma }^{t^{\prime }-t}{r}_{t^{\prime }}.$

A note on the discount $\gamma$ (why the objective quietly changed). The parent note's practical objective is the discounted return $R(\tau )={\sum }_t{\gamma }^t{r}_t$ with a discount factor $\gamma \in [0,1]$ that shrinks far-future treats (so "a treat now beats a treat later"). Setting $\gamma =1$ recovers the plain undiscounted total ${\sum }_t{r}_t$ used in Steps 0–6 — that was just the special case $\gamma =1$, chosen to keep the derivation clean. Everything from Step 3 to Step 6 goes through identically for any fixed $\gamma$, because $\gamma$ lives inside the (frozen, $\theta$-free) rewards, not inside $P_{\theta }$. Here in the reward-to-go, ${\gamma }^{t^{\prime }-t}$ discounts each future step counting from now ($t^{\prime }-t$ = steps ahead). So the objective is unchanged in form; we simply expose the $\gamma$ that was always allowed to sit inside the treats.

WHY dropping the past is still correct. Past step-rewards $r_{t^{\prime }}$ (with $t^{\prime }<t$) are constants with respect to the action at $t$; in expectation they multiply the score to give zero (the same zero-mean argument we prove in Step 8). Dropping them doesn't bias the answer — it just removes pure noise. This is variance reduction.

PICTURE. A timeline of step-rewards $r_{t^{\prime }}$; for the action at $t$, grey out the rewards to the left ("can't affect these") and highlight the rewards to the right ("$G_t$, the only ones this action earns").

Figure s07 — Alt text: a timeline of reward bars; those to the left of a magenta dashed line at the chosen action ($t=3$) are greyed "can't affect," those to the right are orange, labelled reward-to-go $G_t$.

---

Step 8 — Edge case B: baseline subtraction (with the full zero-mean proof)

WHAT. Subtract any function $b(s_t)$ that depends on the state only, not the action: ${\nabla }_{\theta }J={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}\left[{\sum }_{t=0}^{T-1}{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)(G_t-b(s_t))\right].$

WHY it changes nothing (the full proof). We must show the baseline's own contribution averages to exactly zero, for each state $s$. Work with a single term, averaging over the action $a$ drawn from ${\pi }_{\theta }(\cdot \mid s)$:

= \sum_{a}\pi_\theta(a\mid s)\,\nabla_\theta\log\pi_\theta(a\mid s)\,b(s).$$ Now use the log-derivative identity *backwards* — $\pi_\theta\,\nabla_\theta\log\pi_\theta = \nabla_\theta\pi_\theta$ — inside the sum: $$= b(s)\sum_{a}\nabla_\theta\pi_\theta(a\mid s) = b(s)\,\nabla_\theta\!\underbrace{\sum_{a}\pi_\theta(a\mid s)}_{=\,1\ \text{(probabilities sum to one)}} = b(s)\,\nabla_\theta 1 = b(s)\cdot 0 = 0.$$ Reading it: $b(s)$ pulls out front (it doesn't depend on $a$); the sum of a probability distribution over all its actions is always $1$; and the slope of the constant $1$ is $0$. So subtracting $b(s)$ removes **exactly zero** in expectation — the estimate stays **unbiased**, we only reduce variance. **The degenerate case this fixes.** Suppose three actions give returns $G\in\{100,102,101\}$ — *all huge and nearly equal*. Every score gets multiplied by ~$100$, so the update screams "make everything more likely!" and the tiny meaningful differences drown in noise. Subtract the mean $b=101$: weights become $\{-1,+1,0\}$ — now we clearly push the winner up and the loser down. Best baseline: $b(s)=V(s)$, giving the [[Advantage function|advantage]] $A(s,a)=G_t-V(s)$, the doorway to [[Actor-Critic methods]] and [[Generalized Advantage Estimation]]. **PICTURE.** Left panel: three tall near-equal bars all pushed up (noisy). Right panel: same bars after subtracting the mean line — one up, one down, one flat (clean signal). ![[deepdives/dd-ai-ml-5.2.05-d2-s08.png]] *Figure s08 — Alt text: left, three tall near-equal return bars all pushed upward (noisy); right, after subtracting the mean baseline, one bar up (orange), one down (violet), one flat (grey) — a clean learning signal.* --- ## Step 9 — Edge case C: infinite-horizon / continuing tasks **WHAT.** Steps 0–8 assumed every game *ends* at a finite horizon $T$ (an **episodic** task — chess, a maze, one Atari life). But some tasks *never stop*: a server balancing load, a robot walking forever. These are **continuing** (infinite-horizon) tasks, where $T\to\infty$. **WHY the plain total return breaks, and the fix.** With $T\to\infty$ the undiscounted total $\sum_{t=0}^{\infty} r_t$ can blow up to infinity — you cannot maximize "infinity." Two standard repairs, each keeping the *exact same* gradient machinery: - **Discounting saves us.** With $\gamma<1$ the geometric weights $\gamma^t$ shrink fast enough that $\sum_{t=0}^{\infty}\gamma^t r_t$ stays finite (bounded rewards give a finite sum, since $\sum_t \gamma^t = \tfrac{1}{1-\gamma}$). The reward-to-go $G_t=\sum_{t'\ge t}\gamma^{t'-t}r_{t'}$ from Step 7 is then well-defined even with no last step, so the theorem carries over verbatim — just replace the finite $\sum_{t=0}^{T-1}$ by $\sum_{t=0}^{\infty}$. - **Average-reward objective.** Alternatively, maximize the long-run reward *per step*, $J(\theta)=\lim_{T\to\infty}\tfrac1T\mathbb{E}[\sum_{t=0}^{T-1} r_t]$ (finite because it's an average, not a total). The policy-gradient theorem has an analogous form here too, with $G_t-b$ replaced by a *differential* return; the derivation shape — score × weight, averaged — is identical. **WHY nothing structural changes.** Every step from 3 to 8 only ever used two facts: (i) $\theta$ lives inside $\pi_\theta$ and nowhere in $\rho,P,r$; (ii) probabilities sum to $1$. Neither fact cares whether the game is length $T$ or endless. So the boxed theorem is the *same picture* — we only reinterpret the sum's upper limit and how we keep the return finite. **PICTURE.** Two timelines: an **episodic** one that stops at a red "END" flag at $T$, and a **continuing** one running off the edge with fading (discounted) bars showing $\gamma^t$ shrinking future treats. ![[deepdives/dd-ai-ml-5.2.05-d2-s09.png]] *Figure s09 — Alt text: top, an episodic timeline of reward bars ending at a red END flag at horizon $T$; bottom, a continuing timeline running off the right edge with reward bars fading in height as $\gamma^t$ discounts the far future.* --- ## The one-picture summary Every step, compressed into one flow: start → score → drop dynamics → weight by (reward-to-go − baseline) → average → step uphill. ![[deepdives/dd-ai-ml-5.2.05-d2-s10.png]] *Figure s10 — Alt text: a flow chart of the whole derivation — "$J=\mathbb{E}[R]$" → "$\nabla$ hits $P_\theta$" → "log-trick" → "dynamics die" → "$\sum_t\nabla\log\pi$" → "$\times(G_t-b)$" → "step uphill" — ending in the ascent update $\theta\leftarrow\theta+\alpha\,\widehat{\nabla_\theta J}$.* > [!recall]- Feynman: the whole walkthrough in plain words > You've got a machine full of dials that plays a game by flipping weighted coins. You want it to win more treats. The problem: turning the dials changes *which games happen*, but not the treats themselves — so at first it looks like the dials don't affect the score at all. The fix is a bookkeeping trick (log-derivative): "how much more likely would this exact game get if I nudged the dials?" You compute that for every choice you actually made, then multiply by how many treats that game earned. Play lots of games, average these nudges, and turn the dials that way. Two clean-ups: (1) only count treats that came *after* each choice — a choice can't change the past; (2) subtract the "typical" treat so you push good moves up and bad moves down instead of shouting "yes!" at everything. And notice: the world's physics never appeared in the final formula — that's why this works without knowing the rules. If the game never ends, a discount $\gamma$ makes the treat-total finite and the very same recipe still runs. > [!mnemonic] The nine steps as a chant > **Score → Log-trick → Dynamics-die → Sum-the-coins → Weight-by-treat → Trim-the-past → Subtract-the-average → (works even) Forever.** (Steps 4, 5, 6, 7, 8 are the load-bearing ones.) --- ## Connections - [[Policy gradient methods]] — the parent; this page derives its boxed theorem. - [[REINFORCE algorithm]] — the Monte-Carlo estimator we built here. - [[Log-derivative trick]] · [[Score function estimator]] — the engine of Step 4. - [[Advantage function]] · [[Generalized Advantage Estimation]] · [[Actor-Critic methods]] — where Step 8 leads. - [[Trust Region and PPO]] — controlling the step size $\alpha$ from Step 2. - [[Variance reduction in Monte Carlo]] — the theme of Steps 7–8. - [[Value-based methods (Q-learning)]] — the alternative we bypassed. - 🇮🇳 [[5.2.05 Policy gradient methods (Hinglish)]]