ppo.py

"""Proximal Policy Optimization learner and training entry point."""

from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING

import torch
import torch.nn as nn
from gymnasium.vector import VectorEnv
from torch.distributions import Distribution

from deeprl.algorithms.on_policy import (
    OnPolicyConfig,
    OnPolicyLearner,
    train_on_policy,
)
from deeprl.buffers import RolloutBuffer
from deeprl.logger import Logger, MetricHistory
from deeprl.stats import TrainingStats
from deeprl.utils import explained_variance

if TYPE_CHECKING:
    from torch.utils.tensorboard import SummaryWriter

    from deeprl.evaluate import Evaluator


@dataclass(frozen=True)
class PPOConfig(OnPolicyConfig):
    """Optimization hyperparameters specific to PPO-Clip."""

    # Optimization
    lr_critic           :   float        = 3e-4
    lr_actor            :   float        = 1e-3
    batch_size          :   int          = 64
    update_epochs       :   int          = 10
    max_grad_norm       :   float        = 0.5

    # PPO clipped objective
    clip_range          :   float        = 0.2
    value_clip_range    :   float | None = None
    ent_coef            :   float        = 0.0
    normalize_advantage :   bool         = True

    # Optional policy-KL early stopping
    target_kl           :   float | None = None


class PPOLearner(OnPolicyLearner):
    """PPO-Clip optimizers and multi-epoch minibatch update."""

    def __init__(
        self,
        policy      :   nn.Module,
        value_fn    :   nn.Module,
        cfg         :   PPOConfig,
    ) -> None:
        super().__init__(policy, value_fn, cfg)
        self.cfg: PPOConfig = cfg

        self.optim_actor = torch.optim.SGD(
            self.policy.parameters(),
            lr=self.cfg.lr_actor,
        )
        self.optim_critic = torch.optim.SGD(
            self.value_fn.parameters(),
            lr=self.cfg.lr_critic,
        )

    def update(
        self,
        buffer          :   RolloutBuffer,
        advantages      :   torch.Tensor,
        value_targets   :   torch.Tensor,
    ) -> TrainingStats:
        """Optimize the policy and critic over one completed rollout."""

        # ------------------------------------------------------------------
        # 0. Setup
        # ------------------------------------------------------------------

        # Restore training behavior after eval-mode rollout collection. The
        # forward passes below build fresh autograd graphs for every minibatch.
        self.policy.train()
        self.value_fn.train()

        # Accumulate sample-weighted metrics over all accepted minibatch updates.
        metric_sums: dict[str, float] = {
            "losses/policy_loss": 0.0,
            "losses/value_loss": 0.0,
            "diagnostics/entropy": 0.0,
            "diagnostics/approximate_kl": 0.0,
            "diagnostics/clip_fraction": 0.0,
            "grads/policy_norm": 0.0,
            "grads/value_norm": 0.0,
        }
        num_samples = 0
        gradient_steps = 0

        # Normalize once over the complete rollout so every epoch uses the same
        # fixed advantage weights and normalization statistics.
        if self.cfg.normalize_advantage and advantages.numel() > 1:
            advantages = (
                advantages - advantages.mean()
            ) / (advantages.std(unbiased=False) + 1e-8)

        # This flag propagates a target-KL stop from the minibatch loop to the
        # surrounding epoch loop.
        stop_early = False

        for _ in range(self.cfg.update_epochs):

            # Each call reshuffles the complete rollout before yielding flat
            # minibatches of M transitions.
            for batch in buffer.get(
                advantages=advantages,
                value_targets=value_targets,
                batch_size=self.cfg.batch_size,
            ):

                # --------------------------------------------------------------
                # 1. Minibatch re-evaluation
                # --------------------------------------------------------------

                # Re-evaluate pi_theta(.|S), log pi_theta(A|S), and V_phi(S)
                # under the latest parameters. Each scalar output has shape (M,).
                action_distributions: Distribution = self.policy(batch.observations)
                log_probs: torch.Tensor = action_distributions.log_prob(batch.actions)
                values: torch.Tensor = self.value_fn(batch.observations)

                # --------------------------------------------------------------
                # 2. Clipped value objective
                # --------------------------------------------------------------

                # value_targets are the fixed TD(lambda) regression labels. With
                # clipping disabled, use the ordinary critic semi-gradient loss.
                if self.cfg.value_clip_range is None:
                    value_loss = 0.5 * (
                        batch.value_targets - values
                    ).square().mean()
                else:
                    # Restrict the candidate prediction relative to V_old(S):
                    # clipped_value = V_old + clip(V - V_old, -epsilon_v, epsilon_v)
                    clipped_values = batch.old_values + (
                        values - batch.old_values
                    ).clamp(
                        min=-self.cfg.value_clip_range,
                        max=self.cfg.value_clip_range,
                    )
                    unclipped_value_loss = (batch.value_targets - values).square()
                    clipped_value_loss = (
                        batch.value_targets - clipped_values
                    ).square()

                    # The pessimistic maximum prevents the critic from benefiting
                    # merely by moving farther than the configured value range.
                    value_loss = 0.5 * torch.maximum(
                        unclipped_value_loss,
                        clipped_value_loss,
                    ).mean()

                # --------------------------------------------------------------
                # 3. PPO clipped policy objective
                # --------------------------------------------------------------

                # Entropy remains differentiable because it contributes to the
                # policy objective when ent_coef is nonzero.
                entropy = action_distributions.entropy().mean()

                # ratio = pi_theta(A|S) / pi_old(A|S). The denominator remains
                # fixed because old_log_probs were stored during collection.
                log_ratios = log_probs - batch.old_log_probs
                ratios = torch.exp(log_ratios)

                # Maximize the smaller of the ordinary and ratio-clipped
                # surrogate terms for each sampled transition.
                unclipped_objective = ratios * batch.advantages
                clipped_objective = ratios.clip(
                    min=1.0 - self.cfg.clip_range,
                    max=1.0 + self.cfg.clip_range,
                ) * batch.advantages
                policy_loss = -(
                    torch.min(unclipped_objective, clipped_objective).mean()
                    + self.cfg.ent_coef * entropy
                )

                # --------------------------------------------------------------
                # 4. Policy diagnostics and KL early stopping
                # --------------------------------------------------------------

                # These diagnostics use the rollout actions and do not participate
                # in either gradient computation.
                with torch.no_grad():
                    approximate_kl = ((ratios - 1.0) - log_ratios).mean()
                    clip_fraction = (
                        (ratios - 1.0).abs() > self.cfg.clip_range
                    ).float().mean()

                # Stop before another optimizer step when the current policy has
                # already moved substantially beyond the requested target KL.
                if (
                    self.cfg.target_kl is not None
                    and approximate_kl > 1.5 * self.cfg.target_kl
                ):
                    stop_early = True
                    break

                # --------------------------------------------------------------
                # 5. Parameter updates
                # --------------------------------------------------------------

                # Critic semi-gradient update with global L2 gradient clipping.
                self.optim_critic.zero_grad(set_to_none=True)
                value_loss.backward()
                value_grad_norm = torch.nn.utils.clip_grad_norm_(
                    self.value_fn.parameters(),
                    max_norm=self.cfg.max_grad_norm,
                )
                self.optim_critic.step()

                # PPO policy update with global L2 gradient clipping.
                self.optim_actor.zero_grad(set_to_none=True)
                policy_loss.backward()
                policy_grad_norm = torch.nn.utils.clip_grad_norm_(
                    self.policy.parameters(),
                    max_norm=self.cfg.max_grad_norm,
                )
                self.optim_actor.step()

                # --------------------------------------------------------------
                # 6. Minibatch metric accumulation
                # --------------------------------------------------------------

                # Weight by samples so a smaller final minibatch receives its
                # proportional influence. Count only completed optimizer updates.
                num_batch_samples = batch.observations.shape[0]
                num_samples += num_batch_samples
                gradient_steps += 1

                for name, value in (
                    ("losses/policy_loss", policy_loss),
                    ("losses/value_loss", value_loss),
                    ("diagnostics/entropy", entropy),
                    ("diagnostics/approximate_kl", approximate_kl),
                    ("diagnostics/clip_fraction", clip_fraction),
                    ("grads/policy_norm", policy_grad_norm),
                    ("grads/value_norm", value_grad_norm),
                ):
                    metric_sums[name] += value.detach().item() * num_batch_samples

            if stop_early:
                break

        # ------------------------------------------------------------------
        # 7. Update metric reduction
        # ------------------------------------------------------------------

        # Produce one sample-weighted mean for the complete PPO update phase.
        metrics = {
            name: total / num_samples
            for name, total in metric_sums.items()
        }

        # Explained variance uses collection-time values and fixed TD(lambda)
        # targets over the complete rollout, matching the other on-policy learners.
        metrics.update(
            {
                "diagnostics/explained_variance": explained_variance(
                    buffer.values.flatten(),
                    value_targets.flatten(),
                ),
                "charts/policy_learning_rate": self.optim_actor.param_groups[0]["lr"],
                "charts/value_learning_rate": self.optim_critic.param_groups[0]["lr"],
            }
        )

        return TrainingStats(
            metrics=metrics,
            gradient_steps=gradient_steps,
        )


def ppo(
    envs        :   VectorEnv,
    policy      :   nn.Module,
    value_fn    :   nn.Module,
    cfg         :   PPOConfig,
    writer      :   SummaryWriter | None = None,
    logger      :   Logger | None = None,
    evaluator   :   Evaluator | None = None,
) -> MetricHistory:
    """Construct a PPO learner and run the shared on-policy training loop."""

    learner = PPOLearner(policy, value_fn, cfg)

    return train_on_policy(
        envs=envs,
        learner=learner,
        writer=writer,
        logger=logger,
        evaluator=evaluator,
    )