Policy Gradients: Learning the Policy Directly, from a Bandit to Actor-Critic

The throughline: The value of where I am is the reward I just got, plus a discounted value of where I’ll land next. The SARSA, Q-learning & DQN post used that sentence to learn $Q(s,a)$ and read a policy off it by argmax. This post throws out the middleman: learn the policy $\pi(a \mid s;\theta)$ directly and climb expected reward by gradient ascent. The throughline becomes the critic’s job, while the actor learns what to do.

1. The intuition

For two lectures we estimated a value and let it choose actions for us: learn $Q(s,a)$, then act by $\arg\max_a Q(s,a)$. That pipeline has a hidden assumption: the action set is small and discrete, so the argmax is just “compare a few numbers and pick the biggest.”

What if we skip the value and optimize the policy itself?

There are two routes to a policy:

• Value-based (DQN, last lecture). Learn $Q(s,a)$, then read the policy off it: $\pi(s) = \arg\max_a Q(s,a)$. The policy is implicit, a by-product of the values.
• Policy-based (this post). Parameterize the policy itself, $\pi_\theta(a \mid s)$ (a network that outputs action probabilities), and adjust $\theta$ to make good actions more likely. The policy is the thing you learn.

Three concrete situations break the argmax:

1. Continuous actions. A robot arm’s torque is a real number. You cannot take an argmax over infinitely many actions. A policy sidesteps this by outputting the mean and spread of a Gaussian, $a \sim \mathcal{N}(\mu_\theta(s), \sigma_\theta(s))$, and sampling. One forward pass gives you the move.

2. Stochastic optimal policies. In rock-paper-scissors, the only unbeatable strategy is uniform $\frac{1}{3}, \frac{1}{3}, \frac{1}{3}$. A deterministic argmax can never output “play each move a third of the time.” A stochastic policy can.

3. Smooth, direct optimization. A tiny change in $Q$ can flip the argmax from one action to another, a discontinuous jump. Policy gradients move smoothly: a small step in $\theta$ nudges $\pi$, it never snaps it. And they climb expected return $J(\theta)$ directly, not a Bellman-error surrogate. (Recall the deadly triad from the SARSA, Q-learning & DQN post: function approximation + bootstrapping + off-policy data can diverge. Policy gradients avoid off-policy data entirely.)

The honest cost. Policy gradients are on-policy (each batch of experience is used once, then thrown away, so no replay buffer) and high-variance (the gradient is estimated from noisy sampled returns). The rest of this post earns the method and then pays that cost down.

The Archer: the simplest possible RL problem

For most of this post the world has exactly one state. The Archer stands at a fixed spot and picks one of 9 discrete release angles $a_1 \dots a_9$. After shooting, the episode is over: reward = how close the arrow lands to the target (+1 dead centre, falling off with distance).

A one-state, one-shot problem is called a bandit. Your action changes the reward, but never the next state, because there is no next state. That single property strips away returns, discounting, and credit assignment. What is left is the pure policy-gradient mechanism. We add states back in the later sections, once the mechanism is airtight.

flowchart LR
    subgraph actorCritic [Actor-Critic loop]
        S["state s"] --> Actor["Actor π(a|s;θ)"]
        S --> Critic["Critic V(s;w)"]
        Actor -->|"sample action a"| Env["Environment"]
        Env -->|"reward r, next state s'"| Adv["Advantage A = r + γV(s') − V(s)"]
        Adv -->|"weight on ∇log π"| Actor
        Adv -->|"TD target for V"| Critic
    end

This diagram is the destination. We start at the top-left (state in, action out) and build every arrow by the end of Section 2.

2. The math you need

2.1 The objective: expected reward

The goal is a single number: the expected reward under the current policy.

$J(\theta) = \mathbb{E}_{a \sim \pi_\theta}[R(a)] = \sum_a \pi_\theta(a) \cdot R(a)$

Read it aloud, symbol by symbol: ”$J$ of $\theta$ equals the expected value ($\mathbb{E}$) of the reward $R(a)$, when the action $a$ is drawn from the policy $\pi_\theta$ (that is what $a \sim \pi_\theta$ means); and that expected value is the same thing as the sum, over every action $a$, of $\pi_\theta(a)$ times $R(a)$.” The two halves are equal because an expectation is a probability-weighted average: you weight each action’s reward by how often the policy plays it.

$J$ is a function of the policy parameters $\theta$: change $\theta$, the probabilities shift, the weighted sum changes, and $J$ goes up or down. It goes up when we put more probability on higher-reward actions. The whole goal: gradient ascent on $J$.

$\theta \leftarrow \theta + \alpha \cdot \nabla_\theta J(\theta)$

We maximize reward, so we move along the gradient (ascent, the + sign), not against it. In code, optimizers minimize, so we descend on the loss $-J$: the minus sign you will see in every REINFORCE loop.

2.2 The obstacle: you cannot just differentiate the sum

We want $\nabla_\theta J$, the gradient of $J$ with respect to the parameters $\theta$. A “gradient” is just the collection of derivatives, one per parameter. It points in the direction that makes $J$ bigger, which is exactly the direction the update rule above wants to step. So we need to differentiate $J(\theta) = \sum_a \pi_\theta(a) \cdot R(a)$ with respect to $\theta$.

Two basic rules of derivatives are all we need:

1. The derivative of a sum is the sum of the derivatives. We are allowed to reach inside the $\sum_a$ and differentiate one term at a time, then add the results back up.
2. A constant multiplier stays put. Inside each term, $R(a)$ does not depend on $\theta$: the reward an action pays is fixed by the environment, while $\theta$ only changes how likely we are to pick that action. So $R(a)$ is just a constant here, and a constant in front of a derivative comes along unchanged (the same way $\frac{d}{dx}(5x) = 5$). We differentiate only the $\theta$-dependent piece, $\pi_\theta(a)$.

Applying both, term by term:

$\nabla_\theta J = \nabla_\theta \sum_a \pi_\theta(a) R(a) = \sum_a \nabla_\theta \big[\pi_\theta(a) R(a)\big] = \sum_a R(a) \cdot \nabla_\theta \pi_\theta(a)$

Read the final result aloud, symbol by symbol: “the gradient with respect to $\theta$ of $J$ ($\nabla_\theta J$) equals the sum, over every action $a$, of $R(a)$ times the gradient with respect to $\theta$ of $\pi_\theta(a)$ ($\nabla_\theta \pi_\theta(a)$).” In plain English: to raise $J$, push on each action’s probability in proportion to the reward that action earns, big rewards get a big push, small rewards a small one.

This is the naive gradient, and it has a fatal flaw. It needs $R(a)$ for every action, including the ones we never tried, and the sum is huge or infinite for large or continuous action sets. We cannot evaluate it. We need to turn it into something we can sample.

2.3 The score-function trick (the one clean derivation)

Why reach for the log at all? Go back to what broke in Section 2.2. The naive gradient is a sum over every action, and we cannot evaluate it. The one thing we can do in practice is sample actions from the policy and average what we see. A sum that we can estimate by sampling has a special name: an expectation. And an expectation under the policy always has the same shape, $\sum_a \pi_\theta(a) \cdot [\dots]$, with each term weighted by the probability of that action. That weight is what lets us replace “sum over all actions” with “average over the actions we actually drew.”

Now look at our sum, $\sum_a R(a) \cdot \nabla_\theta \pi_\theta(a)$. It has no $\pi_\theta(a)$ out front, so it is not yet an expectation we can sample. The whole job is to manufacture that missing $\pi_\theta(a)$ factor. The trick is to multiply and divide each term by $\pi_\theta(a)$, and the quantity $\nabla_\theta \pi_\theta(a) / \pi_\theta(a)$ is exactly the derivative of $\log \pi_\theta(a)$. That is the only reason the log appears: it is the identity that pulls a $\pi_\theta(a)$ to the front and leaves behind a sum we can sample.

Here is that identity, from the chain rule applied to $\log$:

$\nabla_\theta \log \pi_\theta(a) = \frac{\nabla_\theta \pi_\theta(a)}{\pi_\theta(a)}$

Rearrange:

$\nabla_\theta \pi_\theta(a) = \pi_\theta(a) \cdot \nabla_\theta \log \pi_\theta(a)$

Now substitute into the intractable sum:

$\nabla_\theta J = \sum_a R(a) \cdot \nabla_\theta \pi_\theta(a) = \sum_a \pi_\theta(a) \cdot R(a) \cdot \nabla_\theta \log \pi_\theta(a)$

The leading $\pi_\theta(a)$ is exactly the probability of sampling $a$, so the sum becomes an expectation:

$\boxed{\nabla_\theta J = \mathbb{E}_{a \sim \pi_\theta}\big[R(a) \cdot \nabla_\theta \log \pi_\theta(a)\big]}$

2.4 REINFORCE: sample, observe, push

The estimator from Section 2.3 is REINFORCE. Each episode: sample an action, see the reward, compute one gradient estimate, take a step.

$\hat{g} = R(a) \cdot \nabla_\theta \log \pi_\theta(a), \quad a \sim \pi_\theta$

Read it aloud, symbol by symbol: “the gradient estimate $\hat{g}$ equals the reward $R(a)$ times the gradient with respect to $\theta$ of $\log \pi_\theta(a)$, for an action $a$ sampled from the policy $\pi_\theta$ ($a \sim \pi_\theta$).”

$\nabla_\theta \log \pi_\theta(a)$ is a vector: the direction in parameter space that makes action $a$ more likely (the “nudge”). $R(a)$ is a single number sitting in front of it, so multiplying just stretches or shrinks that direction without rotating it. A big reward makes the nudge long (push hard toward $a$); a small reward makes it short; a negative reward flips the sign (push away from $a$). So in $\hat{g} = R(a) \cdot \nabla_\theta \log \pi_\theta(a)$, the "$\cdot$" is the scaling and $R(a)$ is the scale factor. One sample gives one noisy estimate $\hat{g}$ of the true gradient.

In the bandit there is no future, so the weight on the chosen angle is simply its reward $r$. No returns, no credit assignment. One shot, one reward, one push.

The Archer bandit, concretely. Before the code, here is the exact problem the environment encodes.

• State space. There is only one state, so there is nothing to observe. We still have to feed the network a fixed-size input, so the observation is a constant dummy: the length-1 vector $[0.]$ (declared as a Box of one float in $[0, 1]$). It never changes and carries no information.
• Action space. Nine discrete release angles $a_1 \dots a_9$ (a Discrete(9)), indexed $0 \dots 8$ in code. Each episode the agent picks exactly one.
• Reward. How close the arrow lands to the bullseye, shaped as a Gaussian bell over the angle:

$R(a) = \exp\!\left(-\frac{(a - a_\text{target})^2}{2\sigma^2}\right), \quad a_\text{target} = 4,\ \sigma = 1.5$

Read it aloud: the reward of angle $a$ is $e$ raised to minus the squared distance from the target angle $a_\text{target}$, divided by $2\sigma^2$. In plain English: shoot exactly at the bullseye and the exponent is $0$, so $R = e^0 = 1$ (a perfect score). Miss by a little and the reward dips a little; miss by a lot and it falls toward $0$. With $a_\text{target} = 4$ (angle $a_5$) and $\sigma = 1.5$, angle $a_5$ scores $1.00$, its neighbors $a_4$ and $a_6$ score about $0.80$, and the far edges $a_1$/$a_9$ score about $0.03$. A uniformly random angle averages about $0.42$, which is the baseline the policy has to beat.

The width $\sigma = 1.5$ is a deliberate choice, not a magic constant. It sets how fast reward decays as you leave the bullseye. Make it much smaller and the reward is essentially $0$ everywhere except the center, so the policy gets no “warmer or colder” signal and has to luck into angle $a_5$ before it can learn anything. Make it much larger and every angle scores near $1.0$, so the bullseye barely stands out and there is little worth learning. At $\sigma = 1.5$ the bell spans a couple of angles on each side: the bullseye clearly wins ($1.00$ versus $0.80$ for its neighbors), yet nearby angles still earn graded partial credit, which gives the policy a smooth slope to climb.

The episode ends after that single shot, so there is no next state and no future reward. That is exactly what the code below sets up: a Box observation, a Discrete action space, and this reward in step.

One training knob: the entropy bonus. The policy is a probability distribution over the 9 angles, and its entropy $H(\pi) = -\sum_a \pi(a) \log \pi(a)$ measures how spread out it is: high when the nine probabilities are even (the policy is still exploring), near zero when one angle dominates (the policy has committed). Each call to act returns this entropy next to the log-probability, straight from the Categorical distribution. We add a small $+\,\texttt{ent\_coef} \cdot H$ term to the objective so the policy pays a penalty for collapsing onto one angle too soon, which keeps it trying all nine until it has actually found the bullseye. It appears in the loss as the - ent_coef * ent piece, and a Check in Section 2.11 shows what happens without it.

Here is the full bandit training loop. The policy network maps the constant state to 9 logits, Categorical turns them into probabilities, and we sample one angle. Because the bandit always terminates after one shot, each episode is a single step, so there is no inner while-loop here (the MDP in Section 2.9 adds one):

import gymnasium as gym
from gymnasium import spaces
import numpy as np, torch, torch.nn as nn

class ArcherBandit(gym.Env):
    """One-state bandit: 9 angles, reward = Gaussian centered on bullseye."""
    # 9 discrete angles (0-8), bullseye at index 4, spread 1.5
    N_ANGLES, TARGET, SIGMA = 9, 4, 1.5

    def __init__(self) -> None:
        super().__init__()
        # observation_space declares what a state looks like. spaces.Box is the
        # space for continuous values: Box(low, high, shape, dtype), so this is a
        # length-1 float32 array bounded to [0, 1]. The bandit has no real state,
        # so the observation is always the dummy [0.]
        self.observation_space = spaces.Box(0., 1., (1,), np.float32)
        # spaces.Discrete(N) is the action space: one integer choice in 0..N-1
        # (here the 9 release angles a1..a9)
        self.action_space = spaces.Discrete(self.N_ANGLES)

    def reset(self, *, seed: int | None = None, options: dict | None = None
              ) -> tuple[np.ndarray, dict]:
        super().reset(seed=seed)
        # constant dummy state
        return np.array([0.], np.float32), {}

    def step(self, action: int) -> tuple[np.ndarray, float, bool, bool, dict]:
        # reward = Gaussian bell curve exp(-(a - target)^2 / (2σ^2))
        # it peaks at 1.0 when a == TARGET (dead center) and falls
        # smoothly toward 0 as a moves away from the bullseye
        r = float(np.exp(-((action - self.TARGET)**2) / (2*self.SIGMA**2)))
        # terminated=True: bandit episodes are always one step
        return np.array([0.], np.float32), r, True, False, {}

# ── Policy network ──────────────────────────────────────────────
class Policy(nn.Module):
    def __init__(self, obs_dim: int, n_act: int, h: int = 64) -> None:
        super().__init__()
        # two-layer MLP: obs -> hidden (tanh) -> one logit per action
        self.net = nn.Sequential(nn.Linear(obs_dim, h), nn.Tanh(), nn.Linear(h, n_act))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # raw logits; softmax is applied by Categorical
        return self.net(x)

# ── Action sampling ─────────────────────────────────────────────
# state_np is the observation straight from the env: a NumPy array (here [0.]).
# the _np suffix is a reminder it is NumPy and must be cast to a tensor first
def act(policy: nn.Module, state_np: np.ndarray
        ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    st = torch.tensor(state_np, dtype=torch.float32)
    # Categorical takes raw logits and applies softmax internally
    dist = torch.distributions.Categorical(logits=policy(st))
    # stochastic draw: exploration is built into the policy
    action = dist.sample()
    # log_prob(a) = log π(a) feeds the REINFORCE gradient;
    # entropy H(π) measures how spread out the policy is
    return action, dist.log_prob(action), dist.entropy()

# ── Training loop ───────────────────────────────────────────────
def train_bandit(episodes: int = 1500, lr: float = 0.01, ent_coef: float = 0.1
                 ) -> tuple[Policy, list[float]]:
    env = ArcherBandit()
    pol = Policy(1, env.N_ANGLES)
    opt = torch.optim.Adam(pol.parameters(), lr)
    hist = []
    for ep in range(episodes):
        # one episode = one shot: reset, sample an angle, step once, then update
        s, _ = env.reset()
        a, logp, ent = act(pol, s)
        _, r, *_ = env.step(int(a))
        # REINFORCE loss -r·log π(a), negated because optimizers minimize
        # but we want to maximize expected reward; the entropy bonus
        # (+ent_coef·H) delays collapse to one angle before we explore
        loss = -(r * logp) - ent_coef * ent
        opt.zero_grad(); loss.backward(); opt.step()
        hist.append(r)
    return pol, hist

# seed everything so the run below is reproducible
torch.manual_seed(2)
pol, hist = train_bandit()
probs = torch.softmax(pol(torch.zeros(1)), -1).detach().numpy()
print(f"mean reward (last 200): {np.mean(hist[-200:]):.3f}  (uniform ≈ 0.42)")
print(f"final π: {[f'{p:.3f}' for p in probs]}")
print(f"argmax = a{np.argmax(probs)+1}  (bullseye = a{ArcherBandit.TARGET+1})")

mean reward (last 200): 1.000  (uniform ≈ 0.42)
final π: ['0.000', '0.000', '0.000', '0.000', '1.000', '0.000', '0.000', '0.000', '0.000']
argmax = a5  (bullseye = a5)

The fan sharpens onto the bullseye: all probability concentrates on angle 5 (the target, since TARGET=4 is zero-indexed). The loss line -(r * logp) - ent_coef * ent is the entire algorithm: weight the log-probability of the taken action by its reward (the REINFORCE estimator) and add an entropy bonus to keep exploring.

The important subtlety: that loss is the gradient we derived, in disguise. What we worked out by hand was a gradient, not a loss. It was the estimator

$\hat{g} = R(a)\,\nabla_\theta \log \pi_\theta(a),$

which we wanted to climb with gradient ascent:

$\theta \leftarrow \theta + \alpha\,\hat{g}.$

But PyTorch optimizers only ever minimize; opt.step() does $\theta \leftarrow \theta - \alpha\,\nabla_\theta L$. So we define a surrogate loss whose gradient is exactly $-\hat{g}$:

$L(\theta) = -\,R(a)\,\log \pi_\theta(a) \quad\Rightarrow\quad \nabla_\theta L = -\,R(a)\,\nabla_\theta \log \pi_\theta(a) = -\hat{g}.$

Minimizing this $L$ is therefore the same as ascending $J$:

$\theta \leftarrow \theta - \alpha\,\nabla_\theta L = \theta + \alpha\,\hat{g}.$

That single minus sign is the “negated because optimizers minimize” in the comment. In the bandit $R(a)$ is just the immediate reward r, so -(r * logp) is literally $-R(a)\log\pi_\theta(a)$.

Two things make the match exact. First, r is a constant with respect to $\theta$ (it comes from the environment and carries no gradient), so the gradient flows only through $\log \pi_\theta(a)$, the same reason the hand derivation treated $R(a)$ as a constant multiplier. Second, the numeric value of loss is not meaningful, and you should not read it as “how wrong we are”: it is a stand-in whose only job is to make backward() produce $-\hat{g}$. The extra - ent_coef * ent term is not part of the derived gradient at all; it is a separate entropy regularizer bolted on to slow premature collapse.

2.5 The backward pass: all nine logits move

This section answers the question everyone asks: the policy gave nine probabilities, but we sampled only one angle. Whose gradient do we actually compute?

The loss uses only the sampled action $a_6$:

$L = -r \cdot \log \pi(a_6)$

A natural guess: “we only have a gradient for $a_6$.” That guess is wrong, and here is the derivation that shows why, one line at a time.

Step 1: rewrite $\log \pi(a_6)$ using the softmax. The policy turns logits into probabilities with a softmax, $\pi(a_6) = \dfrac{e^{z_6}}{\sum_k e^{z_k}}$. Take the log of that fraction, using the rule $\log \frac{A}{B} = \log A - \log B$:

$\log \pi(a_6) = \log e^{z_6} - \log \sum_k e^{z_k} = z_6 - \log \sum_k e^{z_k}$

The first term $z_6$ is just the taken logit. The second term, the log of the normalizer, contains every logit $z_1 \dots z_9$. That is the whole point: the instant you normalize, the loss is wired to all nine logits at once.

Step 2: differentiate with respect to one logit $z_k$. Substitute that expansion into $L = -r\,(z_6 - \log \sum_j e^{z_j})$ and differentiate:

$\frac{\partial L}{\partial z_k} = -r\left(\frac{\partial z_6}{\partial z_k} - \frac{\partial}{\partial z_k}\log \sum_j e^{z_j}\right)$

Take the two pieces inside the bracket separately. The first one:

$\frac{\partial z_6}{\partial z_k} = \mathbb{1}[k = a_6]$

The logits are independent inputs, so $z_6$ only changes when we wiggle $z_6$ itself. This equals $1$ when $k = a_6$ and $0$ otherwise, which is exactly what the indicator $\mathbb{1}[\cdot]$ means. The second one:

$\frac{\partial}{\partial z_k}\log \sum_j e^{z_j} = \frac{e^{z_k}}{\sum_j e^{z_j}} = \pi(a_k)$

The derivative of a log is $1/(\text{inside})$ times the derivative of the inside, and the derivative of $\sum_j e^{z_j}$ with respect to $z_k$ is just $e^{z_k}$. The ratio that pops out is exactly the softmax probability of angle $k$.

Step 3: put the pieces back together.

$\frac{\partial L}{\partial z_k} = -r\big(\mathbb{1}[k = a_6] - \pi(a_k)\big)$

This is nonzero for every $k$, because $\pi(a_k) > 0$ for all nine angles. One sampled action, nine gradients.

Step 4: read off the direction. The optimizer moves each logit opposite the loss gradient (it subtracts it), so the actual push is $-\partial L / \partial z_k = r\big(\mathbb{1}[k = a_6] - \pi(a_k)\big)$:

• Taken angle $a_6$ ($\mathbb{1} = 1$): push $\propto +r(1 - \pi_6) > 0$, so its logit goes up and $a_6$ becomes more likely.
• Every other angle ($\mathbb{1} = 0$): push $\propto -r\,\pi_k < 0$, so its logit goes down and that angle becomes less likely.

Step 5: the nine pushes sum to zero.

$\sum_k \big(\mathbb{1}[k = a_6] - \pi(a_k)\big) = \underbrace{\sum_k \mathbb{1}[k = a_6]}_{=\,1} - \underbrace{\sum_k \pi(a_k)}_{=\,1} = 0$

The one-hot contributes $1$, the probabilities sum to $1$, and they cancel. Probability is conserved: whatever $a_6$ gains, the other eight give up. The update redistributes the fan, it never inflates it.

The figure makes the formula concrete: one tall bar for the taken angle $a_6$ pushing it up, and eight short bars for the others pushing them down. The ups and downs are sized so they cancel exactly, which is the “sum to zero” we just derived. One sample, nine moves.

Here is a tiny snippet that reproduces the nine logit gradients with real numbers:

import numpy as np

# hypothetical logits (raw scores) the policy network outputs for 9 angles;
# these are NOT probabilities yet, softmax converts them next
logits = np.array([0.2, 0.6, 1.0, 1.6, 2.0, 1.6, 1.0, 0.6, 0.2])

# softmax π(aₖ) = exp(zₖ) / Σ exp(zⱼ) turns logits into a valid
# probability distribution that sums to 1
probs = np.exp(logits) / np.exp(logits).sum()

# a6 (0-indexed): the action we actually sampled
taken = 5
# reward received for that shot
r = 0.9

# per-logit gradient of the REINFORCE loss L = -r · log π(a_taken):
#   ∂L/∂zₖ = -r · (𝟙[k = taken] - π(aₖ))
#   for the taken action: gradient = -r·(1 - π(a₆)), negative so the logit rises
#   for every other action: gradient = +r·π(aₖ), positive so the logit falls
# np.eye(9)[taken] is a one-hot vector: 1 at position 'taken', 0 elsewhere
grads = -r * (np.eye(9)[taken] - probs)

print("Per-logit gradients (r=0.9, taken=a6):")
for i, g in enumerate(grads):
    tag = " ← taken" if i == taken else ""
    print(f"  a{i+1}: {g:+.3f}{tag}")
# all nine gradients sum to exactly zero: probability is conserved,
# only redistributed among angles, never created or destroyed
print(f"  sum: {grads.sum():.6f}")

# one REINFORCE step (this is the fan in the figure below): the optimizer
# subtracts the loss-gradient from the logits, so a6 rises and the other
# eight fall; lr is the step size, here a small illustrative value
lr = 0.16
new_logits = logits - lr * grads
after = np.exp(new_logits) / np.exp(new_logits).sum()
print("\npolicy before:", [f"{p:.3f}" for p in probs])
print("policy after :", [f"{p:.3f}" for p in after])

Per-logit gradients (r=0.9, taken=a6):
  a1: +0.038
  a2: +0.057
  a3: +0.085
  a4: +0.155
  a5: +0.231
  a6: -0.745 ← taken
  a7: +0.085
  a8: +0.057
  a9: +0.038
  sum: 0.000000

policy before: ['0.042', '0.063', '0.094', '0.172', '0.256', '0.172', '0.094', '0.063', '0.042']
policy after : ['0.042', '0.062', '0.093', '0.167', '0.246', '0.193', '0.093', '0.062', '0.042']

(Note: the signs are the gradient of the loss $L = -r \log\pi(a_6)$; the optimizer subtracts this from the logits, so a6’s logit rises and the others fall, exactly as expected.)

The nine logit-gradients form a single vector $\partial L / \partial \mathbf{z}$. We backprop it once (not nine times). Each weight’s gradient is the sum of all nine logits’ contributions via the chain rule:

$\frac{\partial L}{\partial \theta_j} = \sum_{k=1}^{9} \frac{\partial L}{\partial z_k} \cdot \frac{\partial z_k}{\partial \theta_j}$

Every logit $z_k$ is built from the same shared weights $\theta$, so the chain rule adds up all nine influences. The result is one gradient per weight and one Adam step, not nine.

One update, by hand. Take the same fan the snippet used, $\pi = [0.042, 0.063, 0.094, 0.172, 0.256, 0.172, 0.094, 0.063, 0.042]$ (peaked at $a_5$), and suppose we sample $a_6$ and score $r = 0.9$. The per-logit gradients are exactly the ones printed above: $-0.745$ on the taken angle $a_6$, small positives on the other eight, summing to zero. Now take one optimizer step at learning rate $\alpha \approx 0.16$. Each logit moves by $-\alpha \cdot \partial L / \partial z_k$, so $a_6$ rises and every other angle drops a little. The fan tilts toward the shot that scored well.

The figure shows the nine-angle fan before and after that single step: the bar on $a_6$ grows while the other eight shrink, and the total stays at 1. These are exactly the policy before and policy after rows the snippet prints (one step at lr ≈ 0.16). A single sample moves all nine logits, not just the one you tried. Probability is redistributed, never created.

2.6 The variance problem and the baseline

Without a baseline, every return is positive (in the Archer, reward is always > 0), so every sampled action is pushed up. Good and mediocre alike. The signal is a small difference riding on a big positive offset, and sample noise swamps it.

The fix. Subtract a baseline $b$ that does not depend on the action:

$\hat{g} = (R - b) \cdot \nabla_\theta \log \pi_\theta(a)$

We call the centered weight the advantage $A = R - b$. Better than baseline pushes up ($A > 0$); worse pushes down ($A < 0$). The estimator now takes both signs and mostly cancels, dramatically cutting variance.

Zero-bias proof. We subtracted $b$, so we added a term $-b \cdot \nabla \log \pi(a)$. Does that bias the gradient? We show its expectation is zero, one equality at a time.

Start from the term we added and turn the expectation into a sum. An expectation under $\pi$ is just “sum over all actions, each weighted by its probability $\pi(a)$,” and the constant $b$ pulls out front:

$\mathbb{E}_{a \sim \pi}[b \cdot \nabla \log \pi(a)] = b \cdot \sum_a \pi(a)\, \nabla \log \pi(a)$

Now undo the log-derivative trick. From Section 2.3 we have the identity $\pi(a)\, \nabla \log \pi(a) = \nabla \pi(a)$ (the same trick, just run in reverse), so the $\pi(a)$ and the $\nabla \log$ collapse into a single $\nabla \pi(a)$:

$b \cdot \sum_a \pi(a)\, \nabla \log \pi(a) = b \cdot \sum_a \nabla \pi(a)$

Swap the sum and the gradient. Both are linear operations, so adding up then differentiating equals differentiating then adding up:

$b \cdot \sum_a \nabla \pi(a) = b \cdot \nabla \sum_a \pi(a)$

Finally, the sum of all action probabilities is $1$ by definition (it is a probability distribution), and the gradient of the constant $1$ is $0$, no matter how you change $\theta$:

$b \cdot \nabla \underbrace{\sum_a \pi(a)}_{=\,1} = b \cdot \nabla(1) = 0$

The punchline: the baseline term vanishes in expectation. Subtracting $b$ changes the variance of the estimator but not its mean. We still climb the same $J$, just with far less noise.

The push magnitude. We know REINFORCE pushes good actions up and bad actions down. But how hard does it push? That depends on how surprised the policy is by its own choice.

For a softmax policy, we can derive the gradient of $\log \pi(a)$ with respect to the logit $z_a$ of the chosen action directly. Start from the same softmax expansion we used in Section 2.5:

$\log \pi(a) = z_a - \log \sum_j e^{z_j}$

Differentiate with respect to $z_a$, one term at a time. The first term is simply $\partial z_a / \partial z_a = 1$. The second term is the log-normalizer again, whose derivative with respect to $z_a$ is $e^{z_a} / \sum_j e^{z_j} = \pi(a)$:

$\frac{\partial \log \pi(a)}{\partial z_a} = 1 - \pi(a)$

This is just the $k = a$ case of the per-logit gradient from Section 2.5, where the indicator $\mathbb{1}[k = a]$ equals $1$. Read it as: “one minus the probability the policy already assigned to that action.” This is the score function for the taken action, and it controls how big the parameter update is. Two concrete cases make the intuition click:

• Surprising action pays off. The archer tries angle 7, which the policy thought was unlikely ($\pi(a_7) = 0.05$). It scores well. The push magnitude is $1 - 0.05 = 0.95$: a large update. The policy had a lot to learn from this surprise.
• Confident action pays off. The archer tries angle 5, which the policy already favored ($\pi(a_5) = 0.80$). It scores well. The push magnitude is $1 - 0.80 = 0.20$: a small update. The policy already knew this was good; confirming it again does not warrant a big change.

The policy automatically spends its learning budget where it matters most: on actions it was wrong about, not on ones it already had right. This is built into the math of $\nabla \log \pi$; you get it for free.

The curve is just $1 - \pi(a)$. On the left, an action the policy thought unlikely (small $\pi$) gets a near-1 push when it pays off. On the right, an action the policy was already sure of (large $\pi$) barely moves. The downward slope is the policy spending its learning budget on surprises, exactly the two cases above.

Back to the baseline itself: what value should we actually subtract? Two answers, one better than the other.

The simplest choice is a single constant, the same number for every state and action. The one that cuts variance the most is approximately the average return, $b^* \approx \mathbb{E}[G_t]$. Subtracting it centers the weights: shots that beat the average get a positive weight (push up), shots below it get a negative weight (push down), and the pushes now take both signs and largely cancel instead of all piling up.

A learned, per-state value $V(s)$ does even better. A single constant judges every state by the same bar, so it unfairly marks every action from a hard state (low returns) as “bad” and every action from an easy state (high returns) as “good,” even the best move in a hard state. Using $V(s)$, the expected return from this particular state, the centered weight becomes the advantage $A = G_t - V(s)$: “how much better than typical-from-here did this action do?” That is exactly the right question, and it removes even more variance. That learned $V(s)$ is the critic we build in Section 2.8.

2.7 Adding states: from bandit to MDP

So far the Archer stood at one fixed spot and took one shot. Now we let it walk. The instant an action changes what happens next, the bandit becomes a full sequential problem.

Two things change in the mechanics:

1. The policy reads the state. $\pi_\theta(a) \to \pi_\theta(a \mid s)$. The Archer decides differently at 40 m than at 10 m. Same weights $\theta$, one distribution per distance.

2. The weight becomes the return. The single reward $r$ becomes $G_t$, the discounted sum of rewards that follows a step. A quiet “step closer” that costs $-0.2$ can now be credited for the +10 shot it set up. This is credit assignment.

The return $G_t = r_t + \gamma r_{t+1} + \gamma^2 r_{t+2} + \dots$ was defined in the MDPs & Bellman post. Discounting and returns were explored in the DP, MC & TD post. REINFORCE waits for the full episode to compute $G_t$, exactly like Monte Carlo prediction from that same post.

Deriving the policy gradient with states. With those two changes, the objective $J(\theta)$ is the average return over whole episodes, and we want its gradient $\nabla_\theta J(\theta)$ so we can climb it. It is the same derivation as the bandit, with trajectories in place of single shots. Five small steps.

Step 0: name the pieces. A trajectory $\tau$ is one full episode written out, $\tau = (s_0, a_0, r_0, s_1, a_1, r_1, \dots)$, and its return is $R(\tau) = \sum_t \gamma^t r_t$ (all the discounted rewards in that episode). The objective is the average return over every episode the policy could produce:

$J(\theta) = \mathbb{E}_\tau[R(\tau)] = \sum_\tau P(\tau; \theta)\, R(\tau)$

Read it aloud: ”$J$ of $\theta$ is the expected return, which is the sum over every possible trajectory $\tau$ of how likely that trajectory is, $P(\tau; \theta)$, times its return $R(\tau)$.” This is the exact same shape as the bandit’s $J(\theta) = \sum_a \pi_\theta(a)\, R(a)$, with whole trajectories $\tau$ now playing the role single actions $a$ played before.

Step 1: write how likely a trajectory is. A trajectory unrolls one step at a time: you start somewhere, then over and over the policy picks an action and the environment picks the next state. Multiply all those chances together:

$P(\tau; \theta) = \rho(s_0) \prod_t \pi_\theta(a_t \mid s_t)\, P(s_{t+1} \mid s_t, a_t)$

Symbol by symbol: $\rho(s_0)$ is the chance of the starting state; $\prod_t$ means “multiply over every timestep”; $\pi_\theta(a_t \mid s_t)$ is our policy’s chance of the action it took (the only factor with $\theta$ in it, the only part we control); and $P(s_{t+1} \mid s_t, a_t)$ is the environment’s chance of the next state, which we do not control and which carries no $\theta$. Remember that last point, it is the whole trick in Step 3.

Dropping this $P(\tau; \theta)$ back into Step 0, our full objective now reads:

$J(\theta) = \sum_\tau \Big[\rho(s_0) \prod_t \pi_\theta(a_t \mid s_t)\, P(s_{t+1} \mid s_t, a_t)\Big] R(\tau)$

This is the thing we will differentiate next. It looks busy, but only the $\pi_\theta$ factor depends on $\theta$, so everything else will fall away once we take the gradient.

Step 2: differentiate, and reuse the log trick. We want $\nabla_\theta J$ of the objective we just wrote. We will keep writing the bracketed product compactly as $P(\tau; \theta)$, so the thing to differentiate is $J(\theta) = \sum_\tau P(\tau; \theta)\, R(\tau)$. Three small derivative rules get us there, one at a time.

Rule 1, the sum rule. The derivative of a sum is the sum of the derivatives, so we can push the gradient inside the sum and deal with one trajectory at a time:

$\nabla_\theta J = \sum_\tau \nabla_\theta \big[P(\tau; \theta)\, R(\tau)\big]$

Rule 2, the constant-multiplier rule. Inside each term, $R(\tau)$ does not depend on $\theta$: the return of a fixed trajectory is a fixed number, and $\theta$ only changes how likely that trajectory is, not the reward it pays. A constant in front of a derivative comes along unchanged (the same way $\frac{d}{dx}(5x) = 5$), so $R(\tau)$ slides out and the gradient only has to act on $P(\tau; \theta)$:

$\nabla_\theta J = \sum_\tau R(\tau)\, \nabla_\theta P(\tau; \theta)$

Rule 3, the log-derivative trick. We are now stuck on $\nabla_\theta P(\tau; \theta)$, the gradient of a probability, which is not an average over anything we can sample. This is the exact wall we hit in the bandit, and the same trick breaks it. For any positive function, the chain rule gives $\nabla_\theta \log P = \frac{\nabla_\theta P}{P}$; multiply both sides by $P$ and you get

$\nabla_\theta P(\tau; \theta) = P(\tau; \theta)\, \nabla_\theta \log P(\tau; \theta)$

Substitute that in. The $P(\tau; \theta)$ we just created sits back in front of $R(\tau)$, which turns the bare sum into a probability-weighted average, that is, an expectation we can estimate by rolling out episodes:

$\nabla_\theta J = \sum_\tau P(\tau; \theta)\, \big[R(\tau)\, \nabla_\theta \log P(\tau; \theta)\big] = \mathbb{E}_\tau\big[R(\tau)\, \nabla_\theta \log P(\tau; \theta)\big]$

Step 3: take the log, and watch the environment cancel. We still owe one piece: $\nabla_\theta \log P(\tau; \theta)$. Start from the trajectory probability in Step 1 and use one more rule, the log of a product is the sum of the logs, which turns that long product into a tidy sum:

$\log P(\tau; \theta) = \log \rho(s_0) + \sum_t \log \pi_\theta(a_t \mid s_t) + \sum_t \log P(s_{t+1} \mid s_t, a_t)$

Now differentiate term by term (sum rule once more). A term with no $\theta$ in it has derivative $0$, exactly like $\frac{d}{dx}(5) = 0$. The start-state term $\log \rho(s_0)$ and every transition term $\log P(s_{t+1} \mid s_t, a_t)$ are the environment’s, with no $\theta$ anywhere in them, so they drop to zero:

$\nabla_\theta \log P(\tau; \theta) = \underbrace{\nabla_\theta \log \rho(s_0)}_{=\,0} + \sum_t \nabla_\theta \log \pi_\theta(a_t \mid s_t) + \underbrace{\sum_t \nabla_\theta \log P(s_{t+1} \mid s_t, a_t)}_{=\,0}$

Only the policy terms survive:

$\nabla_\theta \log P(\tau; \theta) = \sum_t \nabla_\theta \log \pi_\theta(a_t \mid s_t)$

This is the payoff: the unknown environment dynamics vanish from the gradient. That single fact is why REINFORCE is model-free, you never need to know how the world transitions. Put it back into the expectation from Step 2:

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

Step 4: only count rewards an action could cause (reward-to-go). Right now every action $a_t$ is weighted by the full-episode return $R(\tau)$, including rewards that arrived before it. But an action cannot change the past. Those earlier rewards are uncorrelated with the push $\nabla_\theta \log \pi_\theta(a_t \mid s_t)$, so by the same zero-expectation argument as the baseline proof, dropping them does not change the average, it only trims variance. Keep just the return from step $t$ onward, the reward-to-go $G_t = \sum_{t' \ge t} \gamma^{t'-t} r_{t'}$, and you arrive at the complete REINFORCE gradient with states:

$\nabla_\theta J(\theta) = \mathbb{E}_\tau \Big[\sum_t G_t \cdot \nabla_\theta \log \pi_\theta(a_t \mid s_t)\Big]$

Read it against the bandit: the bandit was the one-step case, a single term with $G = r$. Adding states just sums the same per-action push over the whole trajectory, each weighted by its return.

The figure traces one trajectory from far (40 m) to the final shot. Each bar is that step’s return $G_t$, and you can see the big +10 at the shot flowing backwards: the discount shrinks it a little at every earlier step, but even the first quiet move inherits most of it. That backward flow is credit assignment.

“Step closer at 40 m” has a return of +5.87, not -0.2, because the discounted future includes the +10 shot it set up. That is how an end-of-episode reward teaches the quiet moves that set it up.

2.8 The advantage: “better than typical here”

REINFORCE with states is painfully noisy because the return $G_t$ is a big, lurching number. The fix from Section 2.6 generalizes: subtract a per-state baseline $V(s)$.

$A_t = G_t - V(s_t)$

The same return can mean very different things in different states. In the Archer MDP, $V(40\text{m}) \approx 6$ (far, hard) and $V(10\text{m}) \approx 9$ (close, easy). The same return $G = 7$ gives:

• at 40 m: $A = 7 - 6 = +1$ (above expectation, push up)
• at 10 m: $A = 7 - 9 = -2$ (below expectation, push down)

Same outcome, opposite lesson, because “good” is judged relative to what is typical at that distance. A single global baseline could never do this.

The advantage can also be written $A(s,a) = Q(s,a) - V(s)$: how much better is taking $a$ than acting typically from $s$? Its sign sets the direction; its magnitude sets how hard to push. That is all the policy needs.

The policy-gradient theorem becomes:

$\nabla_\theta J(\theta) = \mathbb{E}_{s,a}\big[A(s,a) \cdot \nabla_\theta \log \pi_\theta(a \mid s)\big]$

Where did the $\sum_t$ go? It got folded into the expectation. In Section 2.7 we had $\nabla_\theta J = \mathbb{E}_\tau\big[\sum_t G_t \cdot \nabla_\theta \log \pi_\theta(a_t \mid s_t)\big]$: average over trajectories, and inside each trajectory sum the per-step push over every timestep. Summing one quantity over every timestep of every trajectory is the same as averaging that quantity over all the state-action pairs $(s, a)$ the policy actually visits, so the trajectory average $\mathbb{E}_\tau$ and the timestep sum $\sum_t$ merge into the single expectation $\mathbb{E}_{s,a}$ above. We also swapped the weight $G_t$ for the lower-variance advantage $A(s, a)$.

This is the form that every method in the rest of reinforcement learning shares. They differ only in how they estimate $A$.

2.9 Actor-Critic: two heads, one loop

Give the per-state baseline $V(s)$ its own network and let it learn alongside the policy. That single move is the foundation the rest of modern RL is built on.

The Actor is the policy $\pi_\theta(a \mid s)$. It acts: it chooses the move and takes the shot. The Critic is the value $V_w(s)$. It judges: it never acts. Having watched thousands of episodes, it knows how good each position is. Its verdict is the advantage, which is the only feedback the actor needs, and it is far steadier than a raw return.

The advantage is the TD error from the DP, MC & TD post:

$A = r + \gamma V_w(s') - V_w(s)$

Read it as: “the reward I just got ($r$), plus what the critic thinks the next state is worth ($\gamma V_w(s')$), minus what the critic thought this state was worth ($V_w(s)$).”

Why the minus? Split the formula into two halves. The first half, $r + \gamma V_w(s')$, is a fresh, better estimate of how good $s$ was, because it is built from one real reward we actually observed plus the critic’s value for wherever we landed. The second half, $V_w(s)$, is the critic’s old guess for $s$, made before we took the step. Subtracting the old guess from the updated estimate leaves only the surprise: how much better or worse the step turned out than the critic had predicted.

That subtraction is exactly the baseline idea from Sections 2.6 and 2.8. Without the $-V_w(s)$ term we would weight every action by its raw value, a big, almost-always-positive number that barely separates good moves from bad. Subtracting $V_w(s)$ re-centers everything around zero, so only moves that beat the local expectation get pushed up. If $A > 0$, reality beat the critic, push the action up; if $A < 0$, reality was worse, push it down; if $A \approx 0$, the step was about as expected, so barely move.

Two losses, one pass:

• Critic: make its prediction $V(s)$ match what actually happened. The TD target $r + \gamma V(s')$ is the better estimate from above, so we minimize the squared gap $(r + \gamma V(s') - V(s))^2$. The minus here is plain prediction error, how far the guess $V(s)$ sits from the target; squaring makes it positive and punishes big misses harder, so gradient descent slides $V(s)$ toward the target. The target is treated as a fixed number (via detach), so $V(s)$ chases it instead of both sides drifting.
• Actor: push each action by its advantage. We want to maximize $A \cdot \log \pi(a \mid s)$, raising the log-probability of actions whose advantage is positive. But PyTorch optimizers only minimize, so we flip the sign and minimize $-(A \cdot \log \pi(a \mid s))$. This minus is the same gradient-ascent-as-loss bookkeeping from the bandit in Section 2.4, not new math: minimizing the negative is identical to maximizing the original.

We will build the Actor-Critic in five pieces, each right after the idea it implements: the environment, the two networks, the action helper, the training loop, and a greedy evaluation. Read them top to bottom and they concatenate into the full runnable program (also linked from the assignment).

Start with the environment. The Archer MDP adds a distance to the bandit: the archer can step closer, step back, or shoot. Shooting ends the episode with a reward that is largest at 10 m and shrinks with distance, and every walk step costs a small $-0.2$, which pressures the policy to close in before firing.

import gymnasium as gym
from gymnasium import spaces
import numpy as np, torch, torch.nn as nn

# discount factor: how much we value future rewards vs. immediate
GAMMA = 0.99

# ── Environment: Archer MDP ─────────────────────────────────────
# unlike the bandit (one state, one shot), the MDP has a DISTANCE state;
# the archer can walk closer, step back, or shoot from the current distance
class ArcherMDP(gym.Env):
    """MDP: distance-to-target as state, 3 actions (closer/back/shoot)."""
    # three discrete actions
    CLOSER, BACK, SHOOT = 0, 1, 2
    # distance range in meters
    MIN_D, MAX_D = 10., 50.
    def __init__(self, max_steps=25):
        super().__init__()
        # observation = normalized distance d/50, so it lies in [0, 1]
        self.observation_space = spaces.Box(0., 1., (1,), np.float32)
        self.action_space = spaces.Discrete(3)
        self.max_steps = max_steps
    def _obs(self) -> np.ndarray:
        return np.array([self.d / self.MAX_D], np.float32)
    def reset(self, *, seed=None, options=None) -> tuple[np.ndarray, dict]:
        super().reset(seed=seed)
        # start at a random distance from {10, 20, 30, 40, 50}
        self.d = float(self.np_random.choice([10.,20.,30.,40.,50.]))
        self.t = 0
        return self._obs(), {}
    def shoot_reward(self, d: float) -> float:
        # reward for shooting: 10 at 10 m (closest), drops linearly with distance;
        # optimal strategy is to walk to 10 m, then shoot for max reward
        return 10. - 0.28*(d - self.MIN_D)
    def step(self, a: int) -> tuple[np.ndarray, float, bool, bool, dict]:
        self.t += 1
        if a == self.SHOOT:
            # shooting ends the episode (terminated=True) with a distance-based reward
            return self._obs(), float(self.shoot_reward(self.d)), True, False, {}
        # walking costs -0.2 per step (a small penalty to encourage efficiency)
        self.d = float(np.clip(
            self.d + (-10. if a == self.CLOSER else 10.), self.MIN_D, self.MAX_D))
        return self._obs(), -0.2, False, self.t >= self.max_steps, {}

Next, the two heads. The actor is the policy: state in, one logit per action out, and a softmax (applied inside Categorical later) turns those logits into $\pi(a \mid s)$. The critic is the value: state in, a single scalar $V(s)$ out, its estimate of the return from that state. Same input, two different jobs, two separate small networks.

# ── Actor (Policy network) ──────────────────────────────────────
# maps state -> one logit per action; softmax (inside Categorical) converts
# logits to π(a|s). This IS the policy we are optimizing
class Policy(nn.Module):
    def __init__(self, obs_dim, n_act, h=64):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(obs_dim, h), nn.Tanh(), nn.Linear(h, n_act))
    # x: (batch, obs_dim) -> logits (batch, n_act)
    def forward(self, x: torch.Tensor) -> torch.Tensor: return self.net(x)

# ── Critic (Value network) ──────────────────────────────────────
# maps state -> single scalar V(s), the expected return from state s
# under the current policy; used as a baseline to compute the advantage
class Value(nn.Module):
    def __init__(self, obs_dim, h=64):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(obs_dim, h), nn.Tanh(), nn.Linear(h, 1))
    # net gives (batch, 1); squeeze(-1) drops the trailing dim -> (batch,)
    # so V(s) lines up elementwise with the (batch,) rewards and advantages
    def forward(self, x: torch.Tensor) -> torch.Tensor: return self.net(x).squeeze(-1)

The act helper samples an action from the current policy and returns the three things the update needs: the action itself, its log-probability $\log \pi(a \mid s)$ (the “score” the gradient pushes on), and the entropy $H(\pi)$ (how spread out the policy is, used later for the exploration bonus).

# ── Action sampling ─────────────────────────────────────────────
def act(policy: nn.Module, s: np.ndarray) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    st = torch.tensor(s, dtype=torch.float32)
    # Categorical applies softmax to logits -> π(a|s), then we sample
    dist = torch.distributions.Categorical(logits=policy(st))
    # stochastic action: exploration is built into the policy
    a = dist.sample()
    # log_prob = log π(a|s) is the "score" in the REINFORCE gradient;
    # entropy = H(π) measures exploration breadth
    return a, dist.log_prob(a), dist.entropy()

Now the training loop, the heart of the algorithm. Each episode runs in five beats: roll out a full episode, build the TD target, turn it into the advantage, update the critic, then update the actor. The walkthrough right after the block unpacks each beat.

# ── Actor-Critic training loop ──────────────────────────────────
def actor_critic(env: gym.Env, episodes: int = 2000, lr_a: float = 0.004, lr_c: float = 0.03, ent_coef: float = 0.01, warmup: int = 150):
    # actor: learns WHAT to do
    pol = Policy(1, env.action_space.n)
    # critic: learns HOW GOOD each state is
    val = Value(1)
    # actor optimizer (slower lr)
    popt = torch.optim.Adam(pol.parameters(), lr_a)
    # critic optimizer (faster lr)
    vopt = torch.optim.Adam(val.parameters(), lr_c)
    curve = []
    for ep in range(episodes):
        s, _ = env.reset()
        # collect one full episode: states, log-probs, entropies, rewards,
        # next-states, and done flags for every transition
        S, logps, ents, R, S2, Dn = [], [], [], [], [], []
        done = False
        while not done:
            a, logp, ent = act(pol, s)
            s2, r, term, trunc, _ = env.step(int(a))
            done = term or trunc
            S.append(s); logps.append(logp); ents.append(ent)
            R.append(r); S2.append(s2); Dn.append(float(term))
            s = s2

        # convert episode data to tensors for batch computation
        states = torch.tensor(np.array(S), dtype=torch.float32)
        snext  = torch.tensor(np.array(S2), dtype=torch.float32)
        rew    = torch.tensor(R, dtype=torch.float32)
        # 1.0 at the terminal step, 0.0 otherwise
        dn     = torch.tensor(Dn)

        # TD target r + γ·V(s')·(1 - done); the (1-dn) factor zeroes out the
        # bootstrap at terminal states, since there is no future after the
        # episode ends, so V(s') should not contribute there
        with torch.no_grad():
            target = rew + GAMMA * val(snext) * (1 - dn)

        v   = val(states)
        # advantage A = TD_target - V(s) = r + γV(s') - V(s):
        # positive means the action beat expectation, push probability UP;
        # negative means it was worse, push probability DOWN;
        # .detach() stops actor gradients from flowing into the critic
        adv = (target - v).detach()

        # ── Critic update: minimize (TD_target - V(s))² ────────────
        # trains V(s) to predict the discounted return from each state
        vopt.zero_grad()
        (target - v).pow(2).mean().backward()
        vopt.step()

        # ── Actor update (only after warmup) ────────────────────────
        # warmup lets the critic learn a reasonable V(s) before the actor
        # uses the advantage; early random advantages would just be noise
        if ep >= warmup:
            logp = torch.stack(logps)
            ent  = torch.stack(ents)
            popt.zero_grad()
            # actor loss = -𝔼[A · log π(a|s)] - ent_coef · H(π):
            # first term is REINFORCE weighted by advantage (negated for descent),
            # second term is the entropy bonus that keeps the policy exploring
            (-(adv * logp).mean() - ent_coef * ent.mean()).backward()
            popt.step()
        curve.append(sum(R))
    return pol, val, curve