RL for LLM4 / 13

Module 4. Policy gradients: from REINFORCE to PPO

Read the explanations and code in the existing notebooks, or download the files. Code does not run on this site.

Download course files (.zip) 382.1 KB

The archive preserves the source folder structure, both languages, and supporting files.

Module 4 lab — policy gradients, GAE, and PPO

Main module notebook

Download .ipynb module-04-policy-gradients-ppo.en.ipynb · 18.4 KB

Show notebook text

Module 4 lab — policy gradients, GAE, and PPO

This notebook assembles the complete computational route of the module: baseline behavior, a realistic PPO loop, and response-token masking for autoregressive policies. CartPole is implemented inside the course, so the notebook is independent of Gymnasium version changes. The code keeps task termination, external truncation, and rollout-buffer cuts distinct.

The smoke profile is a quick end-to-end CPU check. The full profile increases both updates and random seeds. The numbers and plots are course-code results under the stated protocol, not external benchmarks.

Python

from __future__ import annotations

from dataclasses import replace
from pathlib import Path
import os
import sys
import time

import matplotlib.pyplot as plt
import numpy as np

ROOT = Path.cwd()
if not (ROOT / "modules").is_dir():
    for parent in [ROOT, *ROOT.parents]:
        if (parent / "modules" / "module-04").is_dir():
            ROOT = parent
            break
MODULE_DIR = ROOT / "modules" / "module-04"
if str(MODULE_DIR) not in sys.path:
    sys.path.insert(0, str(MODULE_DIR))

from module_4_reference import (
    gae_advantages,
    masked_token_ppo_objective,
    masked_token_ppo_terms,
    response_token_mask,
    sequence_log_probability_ratio,
    sequence_probability_ratio,
    simulate_group_baseline_bias,
)
from module_4_ppo_training import (
    aggregate_training_runs,
    profile_config,
    train_ppo,
)

RUN_MODE = os.environ.get("RL_LLM_M4_PROFILE", "smoke").strip().lower()
if RUN_MODE not in {"smoke", "full"}:
    raise ValueError("RL_LLM_M4_PROFILE must be 'smoke' or 'full'")

SEEDS = [11, 37] if RUN_MODE == "smoke" else [11, 37, 71, 101, 131]
print("repository root:", ROOT)
print("profile:", RUN_MODE, "| seeds:", SEEDS)

from module_4_llm_sequence import (
    decode_sequence,
    evaluate_sequence_policy,
    sequence_profile_config,
    train_tiny_sequence_policy,
)

1. Baselines: variance reduction and finite-group scale

Compare no baseline, centering on the same group mean, and RLOO. The simulation does not replace the algebra; it checks the predicted $(G-1)/G$ factor and reveals the spread of group gradient estimates.

Python

baseline_result = simulate_group_baseline_bias(
    group_size=8,
    groups=80_000 if RUN_MODE == "smoke" else 400_000,
    seed=2026,
)
for key, value in baseline_result.items():
    print(f"{key:>24s}: {value: .6f}")

labels = ["no baseline", "same-group mean", "RLOO"]
stds = [
    baseline_result["plain_std"],
    baseline_result["same_group_std"],
    baseline_result["rloo_std"],
]
fig, ax = plt.subplots(figsize=(8.2, 4.8))
ax.bar(labels, stds)
ax.set_ylabel("standard deviation of group gradient estimates")
ax.set_title("A baseline changes variance; same-group centering also changes scale")
ax.grid(axis="y", alpha=0.25)
plt.xticks(rotation=8)
plt.show()

2. GAE: two masks answer two different questions

terminated decides whether the next-state value may be used for bootstrapping. trace_cuts decides whether the GAE recursion may continue into the next array record. A rollout-buffer boundary can permit bootstrap while still forbidding recursion into an unrelated record.

Python

rewards = np.array([1.0, 100.0])
values = np.array([2.0, 0.0])
next_values = np.array([10.0, 0.0])
terminated = np.array([False, True])

# Record 0 is cut by the rollout buffer, not by a task terminal. It bootstraps
# from V(next_state)=10, but its GAE recursion must not leak into record 1.
trace_cuts = np.array([True, True])
advantages = gae_advantages(
    rewards,
    values,
    next_values,
    terminated,
    trace_cuts,
    gamma=0.9,
    gae_lambda=1.0,
)
print("advantages with separate bootstrap and trace masks:", advantages)
print("first record target:", values[0] + advantages[0])
assert np.allclose(advantages, [8.0, 100.0])

3. Complete PPO on CartPole

The implementation includes an actor–critic, fixed rollouts, stored old log-probabilities and values, GAE, repeated shuffled minibatch epochs, clipped policy and value objectives, entropy regularization, gradient-norm clipping, KL diagnostics, and separate deterministic evaluation.

Python

runs = []
started = time.perf_counter()
for seed in SEEDS:
    config = profile_config(RUN_MODE, seed=seed)
    # In smoke mode, evaluate every second update to keep the notebook quick.
    if RUN_MODE == "smoke":
        config = replace(config, evaluation_every=2)
    result = train_ppo(config)
    runs.append(result)
    print(
        f"seed={seed:3d} | final mean={result['final_evaluation_mean']:7.2f} "
        f"± {result['final_evaluation_std']:6.2f}"
    )
elapsed = time.perf_counter() - started
curve = aggregate_training_runs(runs)
print(f"training wall time: {elapsed:.2f} s")

Python

steps = curve["environment_steps"]
finite = np.isfinite(curve["median"])
fig, ax = plt.subplots(figsize=(9.0, 5.0))
ax.fill_between(
    steps[finite], curve["q25"][finite], curve["q75"][finite],
    alpha=0.25, label="interquartile range",
)
ax.plot(steps[finite], curve["median"][finite], marker="o", label="median evaluation return")
for run in curve["all_curves"]:
    ax.plot(steps[finite], run[finite], alpha=0.25, linewidth=0.9)
ax.axhline(500.0, linestyle="--", linewidth=1.0, label="time-limit ceiling")
ax.set_xlabel("environment transitions")
ax.set_ylabel("deterministic evaluation return")
ax.set_title(f"Self-contained PPO on CartPole — {RUN_MODE} profile")
ax.grid(alpha=0.25)
ax.legend()
plt.show()

A rising curve from one convenient seed would not establish robustness, and CartPole performance does not transfer automatically to LLMs. The plot verifies the end-to-end implementation contract: collection, targets, repeated optimization, and evaluation work together. Meaningful method comparisons require aligned budgets, several seeds, and uncertainty summaries.

Python

# Diagnostics from the first seed. They are batch-local quantities, not universal thresholds.
history = runs[0]["history"]
updates = np.asarray(history["update"])
fig, ax = plt.subplots(figsize=(9.0, 4.8))
ax.plot(updates, history["approx_kl"], marker="o", label="non-negative approximate KL")
ax.plot(updates, history["active_clip_fraction"], marker="s", label="active clip fraction")
ax.plot(updates, history["ratio_outside_fraction"], marker="^", label="ratio outside fraction")
ax.set_xlabel("PPO update")
ax.set_title("Local optimization diagnostics for one seed")
ax.grid(alpha=0.25)
ax.legend()
plt.show()

4. Autoregressive responses: shift, mask, and reduction unit

The next example deliberately places extreme values in prompt and padding columns. A correct response mask must remove them completely. It then compares local token ratios with the exact response-level product and shows three reductions of the tokenwise PPO surrogate.

For broader tokenization context, see Modern LLMs, Module 2 “Tokenization and Embeddings” (modern_llms.module_02_tokenization). The post-training architecture is mapped in Module 14 “Learning from Feedback: RLHF, DPO, RLOO, and GRPO” (modern_llms.module_12_preference_optimization). This notebook remains self-contained.

Python

# Original token positions: BOS, prompt_1, prompt_2, response_1, response_2, EOS, PAD.
input_ids = np.array([
    [1, 10, 11, 20, 21, 2, 0],
    [1, 12, 30, 2, 0, 0, 0],
])
prompt_lengths = np.array([3, 2])
mask = response_token_mask(
    input_ids,
    prompt_lengths,
    pad_token_id=0,
    eos_token_id=2,
    include_eos=True,
)

# Columns align with labels input_ids[:, 1:]. Prompt and padding values below are
# deliberately extreme so that the mask contract is easy to audit.
old_log_probs = np.array([
    [-9.0, -8.0, -1.0, -1.2, -0.8, -20.0],
    [-7.0, -0.7, -0.9, -30.0, -40.0, -50.0],
])
new_log_probs = np.array([
    [ 9.0,  8.0, -0.7, -1.5, -0.9,  20.0],
    [ 7.0, -0.5, -1.2,  30.0,  40.0,  50.0],
])
advantages = np.array([1.5, -0.8])

ratio, unclipped, clipped, pessimistic = masked_token_ppo_terms(
    new_log_probs,
    old_log_probs,
    advantages,
    mask,
    clip_epsilon=0.2,
)
print("response mask aligned with input_ids[:, 1:]:\n", mask)
print("local ratios (masked positions are multiplicative identity 1):\n", np.round(ratio, 4))
print("exact sequence log-ratios:", np.round(sequence_log_probability_ratio(new_log_probs, old_log_probs, mask), 4))
print("exact sequence ratios:", np.round(sequence_probability_ratio(new_log_probs, old_log_probs, mask), 4))
for reduction in ["token_mean", "sequence_mean", "sequence_sum_mean"]:
    value = masked_token_ppo_objective(
        new_log_probs, old_log_probs, advantages, mask,
        clip_epsilon=0.2, reduction=reduction,
    )
    print(f"{reduction:>18s}: {value: .6f}")

assert np.all(pessimistic[mask == 0] == 0.0)

5. A minimal autoregressive RLOO/PPO training run

The previous autoregressive section exposed the tensor contract without training a policy. We now run the complete path. A small GRU receives one of three prompts and must generate the corresponding two-symbol answer followed by a sampled EOS. A deterministic verifier supplies dense but non-differentiable reward. For each prompt, the code samples a response group, constructs leave-one-out advantages, freezes old token log-probabilities, and performs several PPO epochs.

This is neither a language-quality model nor a miniature claim about production RLHF. Its purpose is narrower and more useful: one executable path contains variable response lengths, the causal shift, prompt/padding masks, generated EOS, a group baseline, and token-level clipping.

Python

sequence_config = sequence_profile_config(RUN_MODE)
print("Autoregressive experiment profile:", sequence_config)
sequence_run = train_tiny_sequence_policy(seed=5, **sequence_config)
first_batch = sequence_run["first_batch"]
assert first_batch is not None

print("First rows of the initial rollout:")
for i in range(min(6, len(first_batch.prompt_classes))):
    print(
        f"  prompt={first_batch.prompt_classes[i]} | "
        f"{decode_sequence(first_batch.input_ids[i].cpu().numpy())} | "
        f"reward={first_batch.rewards[i]:.2f} | "
        f"advantage={first_batch.advantages[i]:+.3f} | "
        f"mask={first_batch.response_mask[i].cpu().numpy().astype(int).tolist()}"
    )
for prompt_class in range(3):
    idx = np.flatnonzero(first_batch.prompt_classes == prompt_class)
    np.testing.assert_allclose(first_batch.advantages[idx].sum(), 0.0, atol=1e-12)
print("RLOO advantages sum to zero within every prompt group.")

Python

sequence_history = sequence_run["history"]
iterations = np.asarray(sequence_history["iteration"])
mean_reward = np.asarray(sequence_history["evaluation_mean_reward"])
plt.figure(figsize=(9, 5))
plt.plot(iterations, mean_reward)
plt.xlabel("Iteration")
plt.ylabel("Mean evaluation reward")
plt.title("A tiny autoregressive policy learns from verifiable reward")
plt.tight_layout()
plt.show()
print("initial evaluation reward:", float(mean_reward[0]))
print("final evaluation reward:", sequence_run["final_evaluation"]["mean_reward"])

Python

success = np.asarray(sequence_history["success_rate"])
plt.figure(figsize=(9, 5))
plt.plot(iterations, success)
plt.ylim(-0.02, 1.02)
plt.xlabel("Iteration")
plt.ylabel("Exact-response success rate")
plt.title("Whole-response success during training")
plt.tight_layout()
plt.show()

print("Greedy responses after training:")
greedy = evaluate_sequence_policy(
    sequence_run["model"], samples_per_prompt=1, deterministic=True
)
for prompt_class, row, reward in zip(
    greedy["prompt_classes"], greedy["input_ids"], greedy["rewards"]
):
    print(f"  prompt={prompt_class} | {decode_sequence(row)} | reward={reward:.2f}")

What the run establishes—and what it does not

The curve checks implementation agreement: reward changes the probability of the desired response, while prompt tokens, padding, and positions after EOS do not enter policy loss. It does not establish PPO robustness for large language models. Production post-training adds a reference policy, distributed rollout collection, a much richer reward model or verifier, length control, mixed objectives, off-policy delay, and substantial systems constraints.

For the broader architecture, see Modern LLMs, Module 14 “Learning from Feedback: RLHF, DPO, RLOO, and GRPO” (modern_llms.module_12_preference_optimization). The variational logic behind entropy and KL regularization is developed in Information Theory for ML, Module 8 “Maximum entropy, exponential models, and KL regularization” (it_ml.module_08_maxent_kl_policy). The notebook itself remains self-contained; both links are optional deepening.

6. Practical takeaways

REINFORCE does not require differentiable reward, but its variance depends on parameterization and baseline.
Advantages used by the actor loss need an explicit detach contract.
GAE needs at least two logical masks: one for bootstrap and one for backward recursion.
PPO clipping limits local surrogate gain; it does not guarantee a trust region.
A complete PPO implementation includes rollout collection, a value function, repeated minibatch updates, diagnostics, and a multi-seed protocol.
Transferring the method to LLMs requires careful causal shift, response masking, EOS handling, and an explicit token-reduction rule.
Group leave-one-out advantages fit into a real autoregressive loop, but one response-level coefficient still does not solve token-level credit assignment.
Every instructional curve needs a boundary of inference: a successful smoke run validates the code contract, not the method’s scaling behavior.