Module 1. Introduction: Reinforcement Learning for LLMs

# Uncomment the installation line when needed in Colab.
# %pip -q install numpy matplotlib
from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Iterable

import matplotlib.pyplot as plt
import numpy as np

OUTPUT_DIR = Path("outputs/module-01/en")
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

np.set_printoptions(precision=6, suppress=True)
SEED = 7
rng = np.random.default_rng(SEED)

def discounted_return(rewards: list[float], gamma: float = 1.0) -> float:
    """Compute G_0 by backward recursion for a non-empty finite sequence."""
    reward_array = np.asarray(rewards, dtype=float)
    if reward_array.ndim != 1 or reward_array.size == 0:
        raise ValueError("rewards must be a non-empty 1D sequence")
    if not np.isfinite(reward_array).all():
        raise ValueError("rewards must be finite")
    if not np.isfinite(gamma) or not 0.0 <= gamma <= 1.0:
        raise ValueError("gamma must be finite and lie in [0, 1]")
    total = 0.0
    for reward in reward_array[::-1]:
        total = float(reward) + gamma * total
    return float(total)


assert discounted_return([0, 0, 1], gamma=1.0) == 1.0
assert np.isclose(discounted_return([1, 1, 1], gamma=0.9), 1 + 0.9 + 0.81)

discounted_return([0, 0, 0, 1], gamma=0.95)

def softmax_with_temperature(logits: np.ndarray, temperature: float = 1.0) -> np.ndarray:
    """Numerically stable softmax(logits / temperature)."""
    logits = np.asarray(logits, dtype=float)
    if logits.ndim != 1 or logits.size == 0:
        raise ValueError("logits must be a non-empty one-dimensional array")
    if not np.all(np.isfinite(logits)):
        raise ValueError("logits must be finite")
    if not np.isfinite(temperature) or temperature <= 0:
        raise ValueError("temperature must be positive and finite")
    scaled = logits / float(temperature)
    scaled = scaled - np.max(scaled)
    weights = np.exp(scaled)
    return weights / np.sum(weights)


def categorical_entropy(probs: np.ndarray) -> float:
    """Entropy in nats for a categorical distribution."""
    probs = np.asarray(probs, dtype=float)
    if probs.ndim != 1 or probs.size == 0:
        raise ValueError("probs must be a non-empty one-dimensional array")
    if np.any(probs < 0) or not np.all(np.isfinite(probs)):
        raise ValueError("probs must be finite and non-negative")
    if not np.isclose(probs.sum(), 1.0, atol=1e-8):
        raise ValueError("probs must sum to 1")
    positive = probs > 0
    return float(-np.sum(probs[positive] * np.log(probs[positive])))


logits = np.array([2.0, 1.0, 0.0, -1.0])
for temperature in [0.25, 0.5, 1.0, 2.0, 4.0]:
    probs = softmax_with_temperature(logits, temperature)
    print(f"T={temperature:>4}: probabilities={probs}, H={categorical_entropy(probs):.4f}")

# An implementation without max subtraction would overflow on this example.
p = softmax_with_temperature(np.array([1000.0, 999.0, 998.0]), temperature=1.0)
assert np.all(np.isfinite(p)) and np.isclose(p.sum(), 1.0)

temperatures = np.geomspace(0.1, 12.0, 400)
entropies = np.array([
    categorical_entropy(softmax_with_temperature(logits, t))
    for t in temperatures
])

fig, ax = plt.subplots(figsize=(8.5, 4.8))
ax.plot(temperatures, entropies, linewidth=2)
ax.axhline(np.log(len(logits)), linestyle="--", linewidth=1.2,
           label="maximum entropy ln K")
ax.set_xscale("log")
ax.set_xlabel("temperature T")
ax.set_ylabel("entropy H(pi_T), nats")
ax.set_title("Entropy of a pure softmax policy for fixed logits")
ax.legend()
fig.tight_layout()
fig.savefig(OUTPUT_DIR / "temperature_entropy.png", dpi=160)
plt.show()

@dataclass(frozen=True)
class CorridorConfig:
    n_states: int = 5
    start: int = 2
    left_terminal: int = 0
    right_terminal: int = 4
    max_steps: int = 20


def corridor_step(
    state: int,
    action: int,
    cfg: CorridorConfig = CorridorConfig(),
) -> tuple[int, float, bool]:
    """Execute one transition in the five-cell corridor."""
    if action not in (0, 1):
        raise ValueError("action must be 0 (left) or 1 (right)")
    if state in (cfg.left_terminal, cfg.right_terminal):
        return state, 0.0, True
    next_state = state - 1 if action == 0 else state + 1
    done = next_state in (cfg.left_terminal, cfg.right_terminal)
    reward = 1.0 if next_state == cfg.right_terminal else 0.0
    return next_state, reward, done


def random_policy(_: int, rng: np.random.Generator) -> int:
    """Select left or right with equal probability."""
    return int(rng.integers(0, 2))


def run_episode(
    rng: np.random.Generator,
    cfg: CorridorConfig = CorridorConfig(),
) -> dict:
    """Collect one trajectory, including its stopping rule and return."""
    state = cfg.start
    states = [state]
    actions: list[int] = []
    rewards: list[float] = []
    done = False
    for _ in range(cfg.max_steps):
        action = random_policy(state, rng)
        next_state, reward, done = corridor_step(state, action, cfg)
        actions.append(action)
        rewards.append(reward)
        states.append(next_state)
        state = next_state
        if done:
            break
    return {
        "states": states,
        "actions": actions,
        "rewards": rewards,
        "done": done,
        "return": discounted_return(rewards, gamma=1.0),
    }


for i in range(5):
    episode = run_episode(np.random.default_rng(SEED + i))
    print(
        f"episode {i}: states={episode['states']}, "
        f"actions={episode['actions']}, return={episode['return']}"
    )

def plot_corridor_episode(
    episode: dict,
    cfg: CorridorConfig = CorridorConfig(),
) -> None:
    """Display the cells and the first transitions of one trajectory."""
    fig, ax = plt.subplots(figsize=(8.8, 2.4))
    ax.set_xlim(-0.5, cfg.n_states - 0.5)
    ax.set_ylim(-0.35, 1.2)
    ax.axis("off")

    for state in range(cfg.n_states):
        if state == cfg.left_terminal:
            face, edge, label = "#F8E2D2", "#D55E00", "terminal\n0"
        elif state == cfg.right_terminal:
            face, edge, label = "#FFF0C8", "#E69F00", "terminal\n+1"
        elif state == cfg.start:
            face, edge, label = "#D9EDF8", "#0072B2", "start\n2"
        else:
            face, edge, label = "#E6F4EF", "#009E73", str(state)
        box = plt.Rectangle(
            (state - 0.42, 0.22), 0.84, 0.52,
            facecolor=face, edgecolor=edge, linewidth=2,
        )
        ax.add_patch(box)
        ax.text(state, 0.48, label, ha="center", va="center", fontsize=11)

    states = episode["states"]
    for t, (start, end) in enumerate(zip(states[:-1], states[1:])):
        y = 0.96 + 0.04 * (t % 2)
        ax.annotate(
            "", xy=(end, y), xytext=(start, y),
            arrowprops={"arrowstyle": "->", "linewidth": 1.6},
        )
    ax.set_title(
        f"sample episode: return={episode['return']}, "
        f"steps={len(episode['actions'])}"
    )
    fig.tight_layout()
    fig.savefig(OUTPUT_DIR / "corridor_episode.png", dpi=160)
    plt.show()


plot_corridor_episode(run_episode(np.random.default_rng(SEED)))

def exact_corridor_value(max_steps: int = 20) -> float:
    """Compute the exact probability of reaching the right terminal within a horizon."""
    if not isinstance(max_steps, int) or isinstance(max_steps, bool) or max_steps <= 0:
        raise ValueError("max_steps must be a positive integer")
    values = np.zeros(5, dtype=float)
    values[4] = 1.0
    for _ in range(max_steps):
        updated = values.copy()
        for state in (1, 2, 3):
            updated[state] = 0.5 * values[state - 1] + 0.5 * values[state + 1]
        updated[0] = 0.0
        updated[4] = 1.0
        values = updated
    return float(values[2])


def monte_carlo_value(
    n_episodes: int,
    seed: int = SEED,
    cfg: CorridorConfig = CorridorConfig(),
) -> tuple[float, np.ndarray]:
    """Estimate the value by averaging returns from independent episodes."""
    if not isinstance(n_episodes, int) or isinstance(n_episodes, bool) or n_episodes <= 0:
        raise ValueError("n_episodes must be a positive integer")
    local_rng = np.random.default_rng(seed)
    returns = np.array(
        [run_episode(local_rng, cfg)["return"] for _ in range(n_episodes)],
        dtype=float,
    )
    return float(returns.mean()), returns


estimate_1000, returns_1000 = monte_carlo_value(1000, seed=SEED)
exact = exact_corridor_value(20)
print(f"Monte Carlo estimate at N=1000: {estimate_1000:.3f}")
print(f"Exact value at horizon 20: {exact:.12f}")
print(f"Approximate standard error: {np.sqrt(exact * (1 - exact) / 1000):.4f}")

# The Monte Carlo estimate changes with SEED; test its domain, not one draw.
assert np.isfinite(estimate_1000) and 0.0 <= estimate_1000 <= 1.0
assert np.isclose(exact, 1023 / 2048)

n = 5000
_, returns = monte_carlo_value(n, seed=SEED)
ns = np.arange(1, n + 1)
running = np.cumsum(returns) / ns
exact = exact_corridor_value(20)
standard_error = np.sqrt(exact * (1 - exact) / ns)

fig, ax = plt.subplots(figsize=(8.8, 4.8))
ax.fill_between(
    ns,
    exact - 2 * standard_error,
    exact + 2 * standard_error,
    alpha=0.35,
    label="approximate ±2 standard errors",
)
ax.plot(ns, running, linewidth=1.5, label="running Monte Carlo estimate")
ax.axhline(exact, linestyle="--", linewidth=1.8,
           label="exact finite-horizon value")
ax.set_xscale("log")
ax.set_ylim(0.42, 0.58)
ax.set_xlabel("number of episodes N")
ax.set_ylabel("value estimate")
ax.set_title("Monte Carlo evaluation of the corridor policy")
ax.legend()
fig.tight_layout()
fig.savefig(OUTPUT_DIR / "corridor_mc_convergence.png", dpi=160)
plt.show()