TRPO & PPO

"""PPO - the clipped surrogate, transcribed.

A linear softmax policy again, so the score has a closed form and the only new
machinery is visible: the probability RATIO against the policy that collected
the batch, and the flat region where the gradient is exactly zero.

The structure to notice is the epoch loop. A plain policy gradient takes one
step per batch; this takes several, and the clip is what makes that safe.
"""

import math
import random


def softmax(scores):
    largest = max(scores)
    exponentials = [math.exp(score - largest) for score in scores]
    total = sum(exponentials)
    return [value / total for value in exponentials]


def policy_probabilities(theta, state):
    scores = []
    for action_weights in theta:
        score = 0.0
        for weight, feature in zip(action_weights, state):
            score += weight * feature
        scores.append(score)
    return softmax(scores)


def sample_action(probabilities, uniform_draw):
    cumulative = 0.0
    for action, probability in enumerate(probabilities):
        cumulative += probability
        if uniform_draw < cumulative:
            return action
    return len(probabilities) - 1


def returns_to_go(rewards, gamma):
    out = [0.0] * len(rewards)
    running = 0.0
    for step in range(len(rewards) - 1, -1, -1):
        running = rewards[step] + gamma * running
        out[step] = running
    return out


def clipped_coefficient(ratio, advantage, epsilon):
    """The scalar that multiplies the score, for min(r*A, clip(r)*A).

    The gradient of r with respect to theta is r * grad log pi, so the whole
    per-sample gradient is (A * r) * grad log pi - EXCEPT inside the flat
    region, where the objective is constant in theta and the gradient is
    exactly zero.

    Read the two cases. A positive advantage wants the action's probability
    raised, and past 1 + epsilon there is no further reward for raising it. A
    negative advantage wants it lowered, and past 1 - epsilon there is no
    further reward for lowering it. The min is what makes this one-sided: it
    declines to reward an overlarge step, and still penalizes a wrong one.
    """
    if advantage >= 0.0 and ratio > 1.0 + epsilon:
        return 0.0
    if advantage < 0.0 and ratio < 1.0 - epsilon:
        return 0.0
    return advantage * ratio


def train(env, n_features, n_actions, iterations=500, episodes_per_batch=16,
          epochs=4, alpha=0.01, gamma=0.99, epsilon=0.2):
    theta = [[0.0] * n_features for _ in range(n_actions)]

    for _ in range(iterations):
        batch = []
        every_return = []

        for _ in range(episodes_per_batch):
            states, actions, rewards, old_log_probs = [], [], [], []
            state = env.reset()
            done = False

            while not done:
                probabilities = policy_probabilities(theta, state)
                action = sample_action(probabilities, random.random())
                next_state, reward, done = env.step(action)

                states.append(state)
                actions.append(action)
                rewards.append(reward)
                # The log-probability under the policy that COLLECTED this
                # sample, frozen here. Every epoch below compares against it,
                # and it must not be recomputed as the policy moves.
                old_log_probs.append(math.log(probabilities[action] + 1e-12))

                state = next_state

            discounted = returns_to_go(rewards, gamma)
            every_return.extend(discounted)
            batch.append((states, actions, discounted, old_log_probs))

        baseline = sum(every_return) / len(every_return)

        # SEVERAL passes over the same data. This is where the sample
        # efficiency over a plain policy gradient comes from, and it is only
        # safe because the clip removes the incentive to keep moving.
        for _ in range(epochs):
            gradient = [[0.0] * n_features for _ in range(n_actions)]
            clipped_count = 0
            total_count = 0

            for states, actions, discounted, old_log_probs in batch:
                for state, action, total_return, old_log_prob in zip(
                    states, actions, discounted, old_log_probs
                ):
                    advantage = total_return - baseline
                    probabilities = policy_probabilities(theta, state)

                    # The ratio starts at exactly 1 on the first epoch, because
                    # the policy has not moved yet, and drifts from there.
                    log_ratio = math.log(probabilities[action] + 1e-12) - old_log_prob
                    ratio = math.exp(log_ratio)

                    coefficient = clipped_coefficient(ratio, advantage, epsilon)
                    total_count += 1
                    if coefficient == 0.0:
                        clipped_count += 1
                        continue        # zero gradient; this sample is inert

                    for candidate in range(n_actions):
                        indicator = 1.0 if candidate == action else 0.0
                        scaled = coefficient * (indicator - probabilities[candidate])
                        for index, feature in enumerate(state):
                            gradient[candidate][index] += scaled * feature

            # Ascent. A clip fraction climbing toward one means most of this
            # epoch contributed nothing at all and the compute was spent anyway.
            clip_fraction = clipped_count / total_count
            for candidate in range(n_actions):
                for index in range(n_features):
                    theta[candidate][index] += alpha * gradient[candidate][index] / total_count

    return theta

"""PPO - GAE, minibatch epochs, KL early stopping, TRPO shown for contrast."""

from __future__ import annotations

from dataclasses import dataclass

import torch
from torch import Tensor, nn
from torch.distributions import Categorical


@dataclass(frozen=True)
class PpoConfig:
    n_features: int
    n_actions: int
    learning_rate: float = 3e-4
    gamma: float = 0.99
    gae_lambda: float = 0.95
    clip_epsilon: float = 0.2
    epochs: int = 4
    minibatches: int = 4
    value_coefficient: float = 0.5
    entropy_coefficient: float = 0.01
    grad_clip_norm: float = 0.5
    # The admission that the clip does NOT bound the KL. It bounds per-sample
    # ratios, and a policy can move a long way with every ratio inside the
    # interval - so the trust region is actually enforced here, by stopping.
    target_kl: float = 0.015

    def __post_init__(self) -> None:
        if not 0.0 < self.clip_epsilon < 1.0:
            raise ValueError(f"clip epsilon must be in (0, 1), got {self.clip_epsilon}")
        if self.epochs < 1 or self.minibatches < 1:
            raise ValueError("epochs and minibatch count must be at least one")
        if not 0.0 <= self.gae_lambda <= 1.0:
            raise ValueError(f"lambda must be in [0, 1], got {self.gae_lambda}")


class RunningNormalizer:
    """Welford statistics over observations.

    Nominally a detail. Measurably a large part of this method's advantage -
    careful ablation found much of PPO's edge over TRPO came from normalization,
    initialization and annealing rather than from the clipped objective. It is
    therefore part of the MODEL, and it must be checkpointed with the weights:
    restoring the network without these statistics feeds it differently scaled
    inputs and produces a different policy, silently.
    """

    def __init__(self, size: int) -> None:
        self.mean = torch.zeros(size)
        self.var = torch.ones(size)
        self.count = 1e-4

    def update(self, batch: Tensor) -> None:
        batch_mean, batch_var = batch.mean(0), batch.var(0, unbiased=False)
        batch_count = batch.shape[0]
        delta = batch_mean - self.mean
        total = self.count + batch_count

        self.mean = self.mean + delta * batch_count / total
        m_a = self.var * self.count
        m_b = batch_var * batch_count
        self.var = (m_a + m_b + delta.pow(2) * self.count * batch_count / total) / total
        self.count = total

    def __call__(self, x: Tensor) -> Tensor:
        return (x - self.mean) / torch.sqrt(self.var + 1e-8)


def orthogonal_init(module: nn.Module, gain: float) -> nn.Module:
    """Orthogonal weights with a SMALL gain on the policy head.

    A near-uniform initial policy matters more here than it looks: the policy
    is the only source of exploration, and one that starts sharp has already
    lost options it will never sample again.
    """
    if isinstance(module, nn.Linear):
        nn.init.orthogonal_(module.weight, gain=gain)
        nn.init.constant_(module.bias, 0.0)
    return module


def ppo_epoch(
    model: nn.Module,
    optimizer: torch.optim.Optimizer,
    states: Tensor,
    actions: Tensor,
    old_log_probs: Tensor,
    advantages: Tensor,
    returns: Tensor,
    config: PpoConfig,
) -> dict[str, float]:
    batch_size = states.shape[0]
    minibatch_size = batch_size // config.minibatches
    indices = torch.randperm(batch_size)

    stats = {"clip_fraction": 0.0, "approx_kl": 0.0, "policy_loss": 0.0, "value_loss": 0.0}

    for start in range(0, batch_size, minibatch_size):
        batch = indices[start : start + minibatch_size]

        logits, values = model(states[batch])
        distribution = Categorical(logits=logits)
        log_probs = distribution.log_prob(actions[batch])

        log_ratio = log_probs - old_log_probs[batch]
        ratio = log_ratio.exp()

        # Advantages normalized per MINIBATCH, not per batch - one of the
        # code-level details that turned out to matter as much as the objective.
        minibatch_advantages = advantages[batch]
        minibatch_advantages = (minibatch_advantages - minibatch_advantages.mean()) / (
            minibatch_advantages.std() + 1e-8
        )

        unclipped = ratio * minibatch_advantages
        clipped = torch.clamp(ratio, 1 - config.clip_epsilon, 1 + config.clip_epsilon)
        # torch.min takes the pessimistic branch, which is what makes the bound
        # one-sided: it declines to reward an overlarge step without ever
        # rewarding one.
        policy_loss = -torch.min(unclipped, clipped * minibatch_advantages).mean()

        value_loss = nn.functional.mse_loss(values, returns[batch])
        entropy = distribution.entropy().mean()

        loss = (
            policy_loss
            + config.value_coefficient * value_loss
            - config.entropy_coefficient * entropy
        )

        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        nn.utils.clip_grad_norm_(model.parameters(), config.grad_clip_norm)
        optimizer.step()

        with torch.no_grad():
            # The low-variance KL estimator, not -log_ratio.mean(): the naive
            # one is unbiased but noisy enough to trip early stopping at random.
            stats["approx_kl"] = float(((ratio - 1) - log_ratio).mean())
            stats["clip_fraction"] = float(
                ((ratio - 1).abs() > config.clip_epsilon).float().mean()
            )
        stats["policy_loss"] = float(policy_loss)
        stats["value_loss"] = float(value_loss)

    return stats


def fisher_vector_product(
    model: nn.Module, states: Tensor, vector: Tensor, damping: float = 0.1
) -> Tensor:
    """TRPO's inner operation, for contrast: F @ v without ever forming F.

    The Fisher is the Hessian of the KL at zero divergence, and the Hessian
    times a vector is the gradient of (gradient . vector) - so one extra
    backward pass gives the product. Conjugate gradients then solve F x = g in
    a handful of iterations, and a backtracking line search enforces the KL
    budget exactly.

    This is the principled version of what the clip approximates. It costs an
    order of magnitude more per update, resists minibatching, and needs a
    second-order routine inside a first-order training stack - which is why the
    cruder method is the one that gets used.
    """
    logits, _ = model(states)
    distribution = Categorical(logits=logits)
    kl = torch.distributions.kl_divergence(
        Categorical(logits=logits.detach()), distribution
    ).mean()

    grads = torch.autograd.grad(kl, list(model.parameters()), create_graph=True)
    flat_grad = torch.cat([g.reshape(-1) for g in grads])

    grad_vector_product = (flat_grad * vector).sum()
    second = torch.autograd.grad(grad_vector_product, list(model.parameters()))
    flat_second = torch.cat([g.reshape(-1) for g in second])

    # Damping keeps the system well-conditioned; without it conjugate gradients
    # stall on the Fisher's near-null directions.
    return flat_second + damping * vector

"""PPO - flattened rollouts, fused ratio statistics, epochs that stop early.

Two of the three changes here save compute. The third saves the run: stopping
the epoch loop on measured KL is both the cheapest optimization available and
the mechanism that actually enforces the trust region the clip only gestures at.
"""

from __future__ import annotations

import numpy as np
import torch
from numpy.typing import NDArray
from torch import Tensor


class FlatRollout:
    """(T, N) collected, (T*N,) consumed - one reshape, no copy.

    The gradient is a sum over independent (state, action, advantage, old
    log-prob) tuples, so time-major or env-major is irrelevant to the estimator
    and the layout can be chosen purely to keep minibatch gathers contiguous.
    """

    def __init__(self, horizon: int, n_envs: int, n_features: int) -> None:
        self.states = np.zeros((horizon, n_envs, n_features), dtype=np.float32)
        self.actions = np.zeros((horizon, n_envs), dtype=np.int64)
        self.rewards = np.zeros((horizon, n_envs), dtype=np.float32)
        self.values = np.zeros((horizon, n_envs), dtype=np.float32)
        # Frozen at collection time and never recomputed as the policy moves.
        self.old_log_probs = np.zeros((horizon, n_envs), dtype=np.float32)
        self.keeps_future = np.zeros((horizon, n_envs), dtype=np.float32)
        self.continues = np.zeros((horizon, n_envs), dtype=np.float32)
        self.advantages = np.zeros((horizon, n_envs), dtype=np.float32)
        # Sized once so the per-epoch shuffle never allocates.
        self.permutation = np.arange(horizon * n_envs, dtype=np.int64)

    def compute_advantages(self, last_value: NDArray[np.float32], gamma: float, lam: float) -> None:
        """GAE fused into one reversed sweep, writing into the existing buffer.

        Delta, the decay and the accumulation are one expression per step, so
        the block is traversed once and no intermediate delta array is written.
        """
        running = np.zeros_like(last_value)
        next_value = last_value

        for step in range(self.rewards.shape[0] - 1, -1, -1):
            delta = (
                self.rewards[step]
                + gamma * next_value * self.keeps_future[step]
                - self.values[step]
            )
            running = delta + gamma * lam * self.continues[step] * running
            self.advantages[step] = running
            next_value = self.values[step]

    def flatten(self, n_features: int):
        return (
            torch.from_numpy(self.states.reshape(-1, n_features)),
            torch.from_numpy(self.actions.reshape(-1)),
            torch.from_numpy(self.old_log_probs.reshape(-1)),
            torch.from_numpy(self.advantages.reshape(-1)),
            torch.from_numpy((self.advantages + self.values).reshape(-1)),
        )


def ratio_statistics(log_ratio: Tensor, clip_epsilon: float) -> tuple[float, float]:
    """Approximate KL and clip fraction from one tensor, in one pass.

    Both diagnostics are functions of the same log-ratio, so computing them
    together avoids a second exponential over the whole minibatch. The KL
    estimator is (r - 1) - log r rather than -log r: the naive form is unbiased
    but noisy enough to trip early stopping at random, and this one has
    dramatically lower variance at the same cost.
    """
    with torch.no_grad():
        ratio = log_ratio.exp()
        approx_kl = ((ratio - 1) - log_ratio).mean()
        clip_fraction = ((ratio - 1).abs() > clip_epsilon).float().mean()
    return float(approx_kl), float(clip_fraction)


def train_on_rollout(
    model: torch.nn.Module,
    optimizer: torch.optim.Optimizer,
    rollout: FlatRollout,
    n_features: int,
    epochs: int,
    minibatches: int,
    clip_epsilon: float,
    target_kl: float,
) -> dict[str, float]:
    states, actions, old_log_probs, advantages, returns = rollout.flatten(n_features)
    batch_size = states.shape[0]
    minibatch_size = batch_size // minibatches

    stats = {"epochs_run": 0.0, "approx_kl": 0.0, "clip_fraction": 0.0}

    for epoch in range(epochs):
        # Shuffled in place, reusing the permutation array allocated once.
        np.random.shuffle(rollout.permutation)
        order = torch.from_numpy(rollout.permutation)

        epoch_kl = 0.0
        for start in range(0, batch_size, minibatch_size):
            batch = order[start : start + minibatch_size]

            logits, values = model(states[batch])          # one shared forward, both heads
            log_probs = torch.log_softmax(logits, dim=1).gather(
                1, actions[batch].unsqueeze(1)
            ).squeeze(1)

            log_ratio = log_probs - old_log_probs[batch]
            ratio = log_ratio.exp()

            normalized = advantages[batch]
            normalized = (normalized - normalized.mean()) / (normalized.std() + 1e-8)

            clipped = torch.clamp(ratio, 1 - clip_epsilon, 1 + clip_epsilon)
            policy_loss = -torch.min(ratio * normalized, clipped * normalized).mean()
            value_loss = torch.nn.functional.mse_loss(values, returns[batch])

            optimizer.zero_grad(set_to_none=True)
            (policy_loss + 0.5 * value_loss).backward()
            optimizer.step()

            epoch_kl, clip_fraction = ratio_statistics(log_ratio, clip_epsilon)
            stats["clip_fraction"] = clip_fraction

        stats["epochs_run"] = float(epoch + 1)
        stats["approx_kl"] = epoch_kl

        # Early stop. The cheapest optimization in the method - a saturated
        # epoch whose samples are mostly clipped contributes nearly nothing
        # anyway - and simultaneously the mechanism that actually bounds the
        # policy movement, which the clip does not.
        if epoch_kl > target_kl:
            break

    return stats