Demo: Three REINFORCEs

Last lecture REINFORCE learned CartPole from scratch (\(23.7 \to 290.5\)), but jaggedly: one noisy scalar \(G(\tau)\) multiplies the entire gradient, and CartPole’s all-positive returns push probability toward everything the agent did. Today we subtract a baseline before the return touches the gradient. It is provably unbiased, and it removes most of the variance. Three training runs, identical in every way (same network seed, same episode seeds, same optimizer) except the number that multiplies the score.

import gymnasium as gym
import numpy as np
import matplotlib.pyplot as plt
import scienceplots
import torch
import torch.nn as nn

plt.style.use(["science", "no-latex"])
TEAL, CARDINAL, GRAY = "#009090", "#9c1b33", "#c9c9c9"

env = gym.make("CartPole-v1")
GAMMA, LR, N_EPISODES = 0.99, 3e-3, 1200

def run_episode(policy, seed):                        # reused by every run below
    state, _ = env.reset(seed=seed)
    log_probs, rewards, done = [], [], False
    while not done:
        logits = policy(torch.tensor(state, dtype=torch.float32))
        dist = torch.distributions.Categorical(logits=logits)
        action = dist.sample()                        # a_t ~ pi(.|s_t)
        log_probs.append(dist.log_prob(action))       # log pi(a_t|s_t)
        state, reward, terminated, truncated, _ = env.step(int(action))
        rewards.append(reward)
        done = terminated or truncated
    G = 0.0
    for r in reversed(rewards):                       # G_t = r_t + gamma * G_{t+1}
        G = r + GAMMA * G
    return torch.stack(log_probs).sum(), G, sum(rewards)

Run 1: REINFORCE, verbatim

Last lecture’s exact run: weight the score by the raw return \(G(\tau)\). Same story, same numbers, so expect \(23.7 \to 290.5\).

torch.manual_seed(5)
policy = nn.Sequential(nn.Linear(4, 32), nn.ReLU(), nn.Linear(32, 2))
optimizer = torch.optim.Adam(policy.parameters(), lr=LR)

returns_vanilla = []
for episode in range(N_EPISODES):
    log_prob_sum, G, steps = run_episode(policy, seed=episode)
    loss = -G * log_prob_sum                          # weight: the raw return
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    returns_vanilla.append(steps)

returns_vanilla = np.array(returns_vanilla)
print(f"REINFORCE: first-50 mean {returns_vanilla[:50].mean():.1f}, "
      f"last-50 mean {returns_vanilla[-50:].mean():.1f}, "
      f"{(returns_vanilla == 500).sum()} episodes at the 500 cap")

REINFORCE: first-50 mean 23.7, last-50 mean 290.5, 45 episodes at the 500 cap

Run 2: subtract a baseline

One change: weight the score by \(G(\tau) - b\), where \(b\) is a running average of past returns. (The baseline must not depend on the current episode’s actions, which is exactly where the unbiasedness proof would break.) Same seeds, same learning rate.

torch.manual_seed(5)
policy_base = nn.Sequential(nn.Linear(4, 32), nn.ReLU(), nn.Linear(32, 2))
optimizer = torch.optim.Adam(policy_base.parameters(), lr=LR)

b = 0.0
returns_baseline = []
for episode in range(N_EPISODES):
    log_prob_sum, G, steps = run_episode(policy_base, seed=episode)
    loss = -(G - b) * log_prob_sum                    # weight: return minus baseline
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    b = 0.9 * b + 0.1 * G                             # update b only AFTER using it
    returns_baseline.append(steps)

returns_baseline = np.array(returns_baseline)
print(f"With baseline: first-50 mean {returns_baseline[:50].mean():.1f}, "
      f"last-50 mean {returns_baseline[-50:].mean():.1f}, "
      f"{(returns_baseline == 500).sum()} episodes at the 500 cap")

With baseline: first-50 mean 21.4, last-50 mean 500.0, 727 episodes at the 500 cap

Run 3: mean-center across a batch

Collect ten episodes and weight each by its centered return \(G^{(i)} - \bar{G}\), so the batch estimates its own baseline at zero extra cost. With the noise this low, we can raise the learning rate tenfold, to \(3 \times 10^{-2}\): the exact rate that collapsed vanilla REINFORCE to a return of \(9\) in last lecture’s what-if.

torch.manual_seed(5)
policy_batch = nn.Sequential(nn.Linear(4, 32), nn.ReLU(), nn.Linear(32, 2))
optimizer = torch.optim.Adam(policy_batch.parameters(), lr=10 * LR)

BATCH = 10
returns_batch = []
for update in range(N_EPISODES // BATCH):
    batch_lp, batch_G = [], []
    for j in range(BATCH):
        log_prob_sum, G, steps = run_episode(policy_batch, seed=update * BATCH + j)
        batch_lp.append(log_prob_sum)
        batch_G.append(G)
        returns_batch.append(steps)
    Gs = torch.tensor(batch_G)
    weights = Gs - Gs.mean()                          # the batch estimates its own baseline
    loss = -(weights * torch.stack(batch_lp)).mean()
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

returns_batch = np.array(returns_batch)
print(f"Batch mean-centering: first-50 mean {returns_batch[:50].mean():.1f}, "
      f"last-50 mean {returns_batch[-50:].mean():.1f}, "
      f"{(returns_batch == 500).sum()} episodes at the 500 cap")

Batch mean-centering: first-50 mean 31.7, last-50 mean 500.0, 919 episodes at the 500 cap

Same estimator mean, three different variances

We proved the baseline cannot move the expected gradient, so every difference between these curves is variance reduction and nothing else.

window = 50
fig, ax = plt.subplots(figsize=(7, 3))
for returns, color, label in [
    (returns_vanilla, GRAY, "REINFORCE (last lecture)"),
    (returns_baseline, CARDINAL, "With running-mean baseline"),
    (returns_batch, TEAL, "With batch mean-centering"),
]:
    moving = np.convolve(returns, np.ones(window) / window, mode="valid")
    ax.plot(np.arange(window, len(returns) + 1), moving, color=color,
            linewidth=1.8, label=label)
ax.axhline(481.5, color="black", linewidth=1.0, linestyle="--",
           label="Hand-coded policy (481.5)")
ax.set_xlabel("Episode")
ax.set_ylabel("Return (steps survived)")
ax.set_ylim(-20, 560)
ax.legend(frameon=False, loc="upper center", bbox_to_anchor=(0.5, 1.30), ncol=2)
plt.show()

Vanilla REINFORCE (gray) never touches the hand-coded policy’s \(481.5\). The running-mean baseline (cardinal) passes it around episode \(546\); batch mean-centering (teal) passes it around episode \(311\) and hovers at the cap from then on.

What if? (change me!)

Problem 24’s generous author returns: add BONUS = 100 to every episode’s return. The true gradient does not move, since we proved a constant cannot move it; only the noise does. Predict both runs before you execute: what happens to vanilla REINFORCE, and what happens to the mean-centered one?

BONUS = 100     # change me: 0, 100, 1000 -- added to every episode's return
N_WHATIF = 400

torch.manual_seed(5)
policy_v = nn.Sequential(nn.Linear(4, 32), nn.ReLU(), nn.Linear(32, 2))
optimizer = torch.optim.Adam(policy_v.parameters(), lr=LR)
returns_v = []
for episode in range(N_WHATIF):                       # vanilla + bonus
    log_prob_sum, G, steps = run_episode(policy_v, seed=episode)
    loss = -(G + BONUS) * log_prob_sum
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    returns_v.append(steps)

torch.manual_seed(5)
policy_b = nn.Sequential(nn.Linear(4, 32), nn.ReLU(), nn.Linear(32, 2))
optimizer = torch.optim.Adam(policy_b.parameters(), lr=10 * LR)
returns_b, largest_weight_change = [], 0.0
for update in range(N_WHATIF // BATCH):               # batch mean-centering + bonus
    batch_lp, batch_G = [], []
    for j in range(BATCH):
        log_prob_sum, G, steps = run_episode(policy_b, seed=update * BATCH + j)
        batch_lp.append(log_prob_sum)
        batch_G.append(G + BONUS)
        returns_b.append(steps)
    Gs = torch.tensor(batch_G)
    weights = Gs - Gs.mean()
    honest = Gs - BONUS                               # what the weights would have been
    weights_honest = honest - honest.mean()
    largest_weight_change = max(largest_weight_change,
                                (weights - weights_honest).abs().max().item())
    loss = -(weights * torch.stack(batch_lp)).mean()
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

returns_v, returns_b = np.array(returns_v), np.array(returns_b)
print(f"Vanilla + bonus:       last-50 mean {returns_v[-50:].mean():.1f}   "
      f"(with honest rewards it had reached {returns_vanilla[:N_WHATIF][-50:].mean():.1f})")
print(f"Mean-centered + bonus: last-50 mean {returns_b[-50:].mean():.1f}   "
      f"(with honest rewards: {returns_batch[:N_WHATIF][-50:].mean():.1f})")
print(f"Largest change the bonus made to any weight the optimizer saw: "
      f"{largest_weight_change:.1e}")

Vanilla + bonus:       last-50 mean 20.8   (with honest rewards it had reached 83.5)
Mean-centered + bonus: last-50 mean 483.9   (with honest rewards: 483.9)
Largest change the bonus made to any weight the optimizer saw: 3.1e-05

The bonus stalls vanilla REINFORCE at the random policy’s level, with the signal buried under the bonus’s common bulk. The mean-centered run subtracts the bonus before the optimizer ever sees it: the largest change the bonus makes to any weight is floating-point dust, and the run finishes exactly where the no-bonus run did. Where you put zero no longer matters, so Problem 24’s scandal is resolved, and Problem 25 resolves it in closed form.

Punchline: none of today’s changes moved the gradient’s mean, and we proved they could not, yet they turned last lecture’s jagged, unfinished curve into two curves that walk calmly past the hand-coded policy. The best baseline is exactly Problem 2’s control-variate coefficient, the same variance reduction the course opened with.