Policy Gradient Methods › Complete Implementation

import torch.nn as nn


class ValueNetwork(nn.Module):
    """Baseline network that estimates expected reward given a prompt."""

    def __init__(self, vocab_size, hidden_size=256):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, hidden_size)
        self.lstm = nn.LSTM(hidden_size, hidden_size, batch_first=True)
        self.output = nn.Linear(hidden_size, 1)

    def forward(self, input_ids):
        embeds = self.embedding(input_ids)
        lstm_out, (h_n, _) = self.lstm(embeds)
        # Use final hidden state to predict expected reward for this context
        value = self.output(h_n[-1])
        return value.squeeze(-1)

import torch
import torch.nn.functional as F


def reinforce_training_step(
    policy,
    value_net,
    prompt_ids,
    reward_fn,
    policy_optimizer,
    value_optimizer,
    max_length=20,
    vocab_size=1000,
):
    """
    Single REINFORCE training step with baseline.

    Args:
        policy: Language model policy (the model being trained)
        value_net: Baseline value network (the critic)
        prompt_ids: Input prompt token IDs
        reward_fn: Function that scores complete sequences
        policy_optimizer: Optimizer for the policy network
        value_optimizer: Optimizer for the value network
        max_length: Maximum number of tokens to generate
        vocab_size: Size of the vocabulary

    Returns:
        Dictionary with loss values, reward, and advantage statistics
    """
    policy.train()
    value_net.train()

    # Step 1: Generate a complete sequence and collect per-token log-probs
    current_ids = prompt_ids.clone()
    log_probs = []

    for _ in range(max_length):
        logits = policy(current_ids)
        probs = F.softmax(logits[:, -1, :], dim=-1)

        # Sample next token from the policy distribution
        action = torch.multinomial(probs, 1)

        # Record the log-probability of the selected token
        log_prob = torch.log(probs.gather(1, action) + 1e-10)
        log_probs.append(log_prob)

        current_ids = torch.cat([current_ids, action], dim=1)

    # Sum log-probs across the sequence to get log pi(y|x)
    log_probs = torch.cat(log_probs, dim=1).sum(dim=1)

    # Step 2: Score the complete sequence with the reward function
    reward = reward_fn(current_ids)

    # Step 3: Compute baseline prediction from the value network
    baseline = value_net(prompt_ids)

    # Step 4: Compute advantage (reward above baseline)
    advantage = reward - baseline.detach()

    # Step 5: Policy loss (REINFORCE with baseline)
    # Detach advantage so policy gradients do not flow through value network
    policy_loss = -(log_probs * advantage).mean()

    # Step 6: Value loss (mean squared error)
    value_loss = F.mse_loss(baseline, reward)

    # Step 7: Update policy parameters
    policy_optimizer.zero_grad()
    policy_loss.backward()
    policy_optimizer.step()

    # Step 8: Update value network parameters
    value_optimizer.zero_grad()
    value_loss.backward()
    value_optimizer.step()

    return {
        "policy_loss": policy_loss.item(),
        "value_loss": value_loss.item(),
        "reward": reward.mean().item(),
        "advantage": advantage.mean().item(),
    }

# Setup
import torch

torch.manual_seed(42)
vocab_size = 100
hidden_size = 64

policy = LanguageModelPolicy(vocab_size, hidden_size, num_layers=1)
value_net = ValueNetwork(vocab_size, hidden_size)

policy_optimizer = torch.optim.Adam(policy.parameters(), lr=1e-3)
value_optimizer = torch.optim.Adam(value_net.parameters(), lr=1e-3)


# Simple reward: bonus for token diversity (unique tokens / total tokens) minus a fixed cost
def length_penalty_reward(sequence_ids):
    """Reward function based on token diversity in the generated portion."""
    batch_size = sequence_ids.shape[0]
    rewards = []
    for i in range(batch_size):
        generated = sequence_ids[i, 5:]  # Strip the 5-token prompt
        n_unique = len(generated.unique())
        n_total = len(generated)
        diversity = n_unique / max(n_total, 1)  # 0 to 1
        rewards.append(
            diversity - 0.5
        )  # Centered: positive if diverse, negative if repetitive
    return torch.tensor(rewards, dtype=torch.float32)


# Training loop
prompt = torch.randint(0, vocab_size, (1, 5))  # Batch of 1, prompt length 5
training_history = []

for step in range(100):
    metrics = reinforce_training_step(
        policy,
        value_net,
        prompt,
        length_penalty_reward,
        policy_optimizer,
        value_optimizer,
        max_length=15,
        vocab_size=vocab_size,
    )
    training_history.append(metrics)

import matplotlib.pyplot as plt

use_book_style(fixed_canvas=True)
plt.rcParams["figure.figsize"] = (2.0, 2.4)

steps = range(len(training_history))
policy_losses = [h["policy_loss"] for h in training_history]
value_losses = [h["value_loss"] for h in training_history]
rewards_hist = [h["reward"] for h in training_history]

# First plot
fig = plt.figure()
fig.subplots_adjust(left=0.29, right=0.96, bottom=0.20, top=0.78)
plt.plot(
    steps,
    policy_losses,
    alpha=0.7,
    color=theme_color("steelblue"),
    linewidth=1.1,
)
plt.xlabel("Training Step")
plt.ylabel("Policy Loss")
plt.title("Policy Loss\n(High Variance)")
plt.show()

# Second plot
fig = plt.figure()
fig.subplots_adjust(left=0.29, right=0.96, bottom=0.20, top=0.78)
plt.plot(
    steps, value_losses, alpha=0.7, color=theme_color("coral"), linewidth=1.1
)
plt.xlabel("Training Step")
plt.ylabel("Value Loss")
plt.title("Value Loss\n(Baseline Learning)")
plt.show()

# Third plot
fig = plt.figure()
fig.subplots_adjust(left=0.29, right=0.96, bottom=0.20, top=0.78)
plt.plot(
    steps,
    rewards_hist,
    alpha=0.7,
    color=theme_color("forestgreen"),
    linewidth=1.1,
)
plt.xlabel("Training Step")
plt.ylabel("Reward")
plt.title("Average Reward")
plt.show()