Demo: Comparing Three REINFORCE Estimators

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. Subtracting a baseline from the return before it multiplies the gradient preserves unbiasedness and 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)\). With the same seeds and settings, this run should reproduce \(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.

Adding a constant reward bonus

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 uses it: the largest change to any weight is at floating-point precision, and the run finishes at the same return as the no-bonus run. Mean-centering makes the update invariant to the location of zero on the reward scale.