PPO for Language Models › Implementation: PPO Training Step

from dataclasses import dataclass
from typing import Tuple

import torch
import torch.nn.functional as F


@dataclass
class PPOConfig:
    """Configuration for PPO training."""

    clip_epsilon: float = 0.2
    kl_coef: float = 0.1
    value_coef: float = 0.5
    gamma: float = 1.0
    lam: float = 0.95


class LanguageModelPPO:
    """
    PPO trainer for language models.

    This simplified implementation demonstrates the key components
    without the full complexity of production systems.
    """

    def __init__(self, policy_model, ref_model, value_model, config: PPOConfig):
        """
        Args:
            policy_model: The LLM being optimized
            ref_model: Frozen reference model for KL computation
            value_model: Value function (often shares backbone with policy)
            config: PPO hyperparameters
        """
        self.policy = policy_model
        self.ref = ref_model
        self.value = value_model
        self.config = config

        # Freeze reference model
        for param in self.ref.parameters():
            param.requires_grad = False

    def compute_log_probs(self, model, input_ids, response_ids):
        """
        Compute log probabilities for response tokens.

        Args:
            model: Language model
            input_ids: Full sequence (prompt + response)
            response_ids: Just the response tokens

        Returns:
            Log probabilities for each response token
        """
        with torch.no_grad() if model == self.ref else torch.enable_grad():
            outputs = model(input_ids)
            logits = outputs.logits if hasattr(outputs, "logits") else outputs

            # Get logits for positions that predict response tokens
            # (shifted by 1 for autoregressive prediction)
            response_start = input_ids.shape[1] - response_ids.shape[1]
            response_logits = logits[:, response_start - 1 : -1, :]

            # Compute log probabilities
            log_probs = F.log_softmax(response_logits, dim=-1)

            # Gather log probs for actual tokens
            token_log_probs = torch.gather(
                log_probs, dim=-1, index=response_ids.unsqueeze(-1)
            ).squeeze(-1)

        return token_log_probs

    def compute_rewards_and_advantages(
        self,
        reward_scores: torch.Tensor,
        log_probs_policy: torch.Tensor,
        log_probs_ref: torch.Tensor,
        values: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Compute KL-shaped rewards and GAE advantages.
        """
        batch_size, seq_len = log_probs_policy.shape

        # KL penalty per token
        kl = log_probs_policy - log_probs_ref

        # Initialize rewards (sparse: only at final token)
        rewards = torch.zeros_like(log_probs_policy)
        rewards[:, -1] = reward_scores

        # Apply KL shaping
        shaped_rewards = rewards - self.config.kl_coef * kl

        # Compute GAE
        advantages = torch.zeros_like(shaped_rewards)
        last_gae = torch.zeros(batch_size, device=shaped_rewards.device)

        for t in reversed(range(seq_len)):
            next_value = (
                values[:, t + 1]
                if t < seq_len - 1
                else torch.zeros_like(last_gae)
            )
            delta = (
                shaped_rewards[:, t]
                + self.config.gamma * next_value
                - values[:, t]
            )
            advantages[:, t] = last_gae = (
                delta + self.config.gamma * self.config.lam * last_gae
            )

        returns = advantages + values

        return advantages, returns, kl

    def ppo_loss(
        self,
        log_probs_new: torch.Tensor,
        log_probs_old: torch.Tensor,
        advantages: torch.Tensor,
    ) -> torch.Tensor:
        """
        Compute clipped PPO policy loss.
        """
        # Normalize advantages
        advantages = (advantages - advantages.mean()) / (
            advantages.std() + 1e-8
        )

        # Probability ratio
        ratio = torch.exp(log_probs_new - log_probs_old)

        # Clipped surrogate
        clipped_ratio = torch.clamp(
            ratio, 1 - self.config.clip_epsilon, 1 + self.config.clip_epsilon
        )

        # Loss (negative because we minimize)
        policy_loss = -torch.min(
            ratio * advantages, clipped_ratio * advantages
        ).mean()

        return policy_loss

    def value_loss(
        self, values: torch.Tensor, returns: torch.Tensor
    ) -> torch.Tensor:
        """
        Compute value function loss.
        """
        return F.mse_loss(values, returns)

# Initialize configuration
config = PPOConfig()