Policy gradient — REINFORCE — Coding notes

WHAT is REINFORCE?

WHY policy gradients at all? In value-based methods (Q-learning) you learn values and act greedily. But for continuous or high-dimensional action spaces (e.g. an aircraft's control-surface deflections), maximizing over actions is painful. Policy methods directly parameterize the action distribution and optimize $\theta$ by gradient ascent — no argmax needed.

HOW: deriving the gradient from scratch

We want $\nabla_\theta J(\theta)$ . The trouble: the expectation is over $\tau$ , whose distribution depends on $\theta$ . We can't just push the gradient inside the expectation naively.

Step 1 — write the expectation as an integral. Let $p_\theta(\tau)$ be the probability of trajectory $\tau$ : $J(\theta)=\int p_\theta(\tau)\,R(\tau)\,d\tau .$ Why this step? An expectation is a probability-weighted integral; making it explicit lets us differentiate.

Step 2 — differentiate. Only $p_\theta$ depends on $\theta$ : $\nabla_\theta J=\int \nabla_\theta p_\theta(\tau)\,R(\tau)\,d\tau .$

Step 3 — the log-derivative trick. We use the identity $\nabla_\theta p_\theta(\tau)=p_\theta(\tau)\,\nabla_\theta \log p_\theta(\tau),$ which is just $\nabla \log f = \frac{\nabla f}{f}$ rearranged. Why this step? It re-inserts $p_\theta(\tau)$ as a weight, turning the integral back into an expectation we can sample: $\nabla_\theta J=\int p_\theta(\tau)\,\nabla_\theta\log p_\theta(\tau)\,R(\tau)\,d\tau=\mathbb{E}_{\tau}\big[\nabla_\theta\log p_\theta(\tau)\,R(\tau)\big].$

Step 4 — expand $\log p_\theta(\tau)$ . The trajectory probability factorizes: $p_\theta(\tau)=\underbrace{p(s_0)}_{\text{env}}\prod_{t=0}^{T}\underbrace{\pi_\theta(a_t\mid s_t)}_{\text{policy}}\;\underbrace{p(s_{t+1}\mid s_t,a_t)}_{\text{env dynamics}}.$ Taking $\log$ turns products into sums; then $\nabla_\theta$ kills every term that doesn't contain $\theta$ (the environment terms $p(s_0)$ and $p(s_{t+1}\mid s_t,a_t)$ ): $\nabla_\theta\log p_\theta(\tau)=\sum_{t=0}^{T}\nabla_\theta\log\pi_\theta(a_t\mid s_t).$ Why this matters: we never need a model of the environment. This is the magic of REINFORCE.

Reward-to-go and baselines (the 20% that gives 80%)

Reward-to-go. An action at time $t$ can't affect rewards that already happened. So replace the full $R(\tau)$ with the return from $t$ onward: $\nabla_\theta J=\mathbb{E}\!\left[\sum_t \nabla_\theta\log\pi_\theta(a_t\mid s_t)\;\underbrace{\sum_{t'\ge t}\gamma^{t'-t}r_{t'}}_{G_t}\right].$ Why? Removing irrelevant future-independent-of-action reward lowers variance without adding bias.

Baseline. Subtract any function $b(s_t)$ that doesn't depend on $a_t$ : $\nabla_\theta J=\mathbb{E}\!\left[\sum_t \nabla_\theta\log\pi_\theta(a_t\mid s_t)\,\big(G_t-b(s_t)\big)\right].$ This is unbiased because $\mathbb{E}_{a}[\nabla_\theta\log\pi_\theta(a\mid s)\,b(s)]=b(s)\,\nabla_\theta\!\underbrace{\sum_a\pi_\theta(a\mid s)}_{=1}=0.$ A good baseline is $b(s)=V(s)$ ; then $G_t-b(s_t)$ is an advantage estimate.

Worked examples

Worked example 1 — Two-action softmax (single step)

Policy $\pi_\theta(a)=\text{softmax}$ over 2 actions with logits $\theta=(\theta_0,\theta_1)$ . Sample action $a=1$ , get reward $R=+2$ .

$\log\pi_\theta(1)=\theta_1-\log(e^{\theta_0}+e^{\theta_1})$ .
$\nabla_{\theta_1}\log\pi_\theta(1)=1-\pi_\theta(1)$ . Why? derivative of log-softmax of the chosen class is $1-p_{\text{chosen}}$ .
$\nabla_{\theta_0}\log\pi_\theta(1)=-\pi_\theta(0)$ .
Gradient $\times R$ : $\theta_1$ pushed up by $2(1-\pi(1))$ , $\theta_0$ pushed down by $2\pi(0)$ → action 1 becomes more likely. ✔ matches intuition.

Worked example 2 — Gaussian policy (continuous control, e.g. thrust)

$\pi_\theta(a\mid s)=\mathcal N(\mu_\theta(s),\sigma^2)$ . Then $\log\pi_\theta(a\mid s)=-\frac{(a-\mu_\theta(s))^2}{2\sigma^2}-\log(\sigma\sqrt{2\pi}).$ $\nabla_{\mu}\log\pi=\dfrac{a-\mu}{\sigma^2}$ . Why this step? only the squared term contains $\mu$ . Update $\propto \dfrac{a-\mu}{\sigma^2}(G_t-b)$ : if a higher-than-mean thrust ( $a>\mu$ ) gave positive advantage, $\mu$ shifts toward that action. Aerospace read: the controller learns to nudge its mean commanded deflection toward whatever worked.

Common mistakes (steel-manned)

Recall Feynman: explain to a 12-year-old

Imagine training a dog. You can't tell it exactly what to do — you just let it try things. When it does a trick and you give a treat, it becomes a bit more likely to do that trick again. REINFORCE is the same: the robot tries random actions, and whichever actions came before a big reward get "voted up" a little. Do this thousands of times and the robot's random guessing slowly turns into skill. The clever math bit ( $\nabla\log\pi$ ) is just how much to turn each knob up so the good action gets more likely.

Active recall

What objective does REINFORCE maximize?

Expected return

J(\theta)=\mathbb{E}_{\tau\sim\pi_\theta}[R(\tau)]

State the REINFORCE gradient.

\nabla_\theta J=\mathbb{E}_\tau[(\sum_t\nabla_\theta\log\pi_\theta(a_t|s_t))\,R(\tau)]

What is the log-derivative (score-function) trick?

\nabla_\theta p_\theta=p_\theta\nabla_\theta\log p_\theta

, which turns a gradient of an integral back into a sampleable expectation.

Why is REINFORCE model-free?

\log p_\theta(\tau)

splits into env terms (no

\theta

) +

\sum_t\log\pi_\theta(a_t|s_t)

;

\nabla_\theta

kills the env terms, so no transition model is needed.

What is "reward-to-go" and why use it?

G_t=\sum_{t'\ge t}\gamma^{t'-t}r_{t'}

; using it instead of full

R(\tau)

removes reward independent of the action, lowering variance without bias.

Why does subtracting a baseline

b(s)

keep the estimate unbiased?

Because

\mathbb{E}_a[\nabla_\theta\log\pi_\theta(a|s)\,b(s)]=b(s)\nabla_\theta\sum_a\pi_\theta=b(s)\nabla_\theta 1=0

Is REINFORCE gradient ascent or descent?

Ascent:

\theta\leftarrow\theta+\alpha\nabla_\theta J

(or minimize

-\log\pi\cdot G

For a Gaussian policy, what is

\nabla_\mu\log\pi_\theta(a|s)

(a-\mu)/\sigma^2

For softmax, what is

\nabla_{\theta_k}\log\pi

of the chosen class

k

1-\pi_\theta(k)

(and

-\pi_\theta(j)

for a non-chosen

j

Main practical weakness of vanilla REINFORCE?

High gradient variance → slow, noisy learning; mitigated by reward-to-go and baselines.

Connections

Value-based methods (Q-learning) — the alternative that learns values, not policies.
Actor-Critic methods — uses a learned $V(s)$ as the baseline/critic.
Advantage function A(s,a) — what $G_t-b(s_t)$ estimates.
Softmax and log-softmax gradients — reused for discrete policies.
Gaussian policies for continuous control — thrust/deflection commands in aerospace.
Monte Carlo estimation and variance reduction — why baselines help.
Gradient ascent / stochastic optimization — the update engine.

Concept Map

Hinglish (regional understanding)

Intuition Hinglish mein samjho

Dekho, REINFORCE ka core idea bilkul dog-training jaisa hai. Agent ke paas ek policy hoti hai $\pi_\theta(a\mid s)$ — yaani state $s$ dekh kar action $a$ ka probability distribution. Hum chahte hain expected reward $J(\theta)$ maximize karna. Problem yeh hai ki reward ka expectation trajectory $\tau$ par hai, aur trajectory ka distribution khud $\theta$ par depend karta hai — toh gradient seedha andar nahi ghusa sakte.

Yahan log-derivative trick kaam aata hai: $\nabla p = p\,\nabla\log p$ . Iski wajah se gradient wapas ek expectation ban jaata hai jise hum samples se estimate kar sakte hain. Aur sabse mast baat — jab hum $\log p_\theta(\tau)$ expand karte hain, toh environment ke transition terms $\theta$ pe depend nahi karte, isliye $\nabla_\theta$ unhe zero kar deta hai. Bache sirf $\sum_t\nabla_\theta\log\pi_\theta(a_t\mid s_t)$ . Matlab koi environment model nahi chahiye — pure model-free. Final formula: gradient $=$ (log-prob ka gradient) $\times$ (return $G_t$ ), aur phir $\theta$ ko us direction mein upar le jao (ascent).

Aerospace context mein socho: ek controller jo thrust ya control-surface deflection command karta hai. Gaussian policy use karo — action ka mean $\mu_\theta(s)$ . Agar mean se zyada thrust dene par acha result mila (positive advantage), toh $\nabla_\mu\log\pi=(a-\mu)/\sigma^2$ mean ko us achhe action ki taraf shift kar deta hai. Dheere-dheere random guessing skill ban jaati hai.

Do practical tips yaad rakho: (1) reward-to-go $G_t$ use karo, full episode reward nahi — variance kam hota hai, bias same. (2) Baseline $b(s)$ subtract karo (jaise $V(s)$ ) — yeh unbiased rehta hai kyunki uska expected contribution zero hai, par variance aur kam. Bas yahi 20% concepts 80% performance dila dete hain.

Test yourself — Machine Learning (Aerospace Applications)

Connections