REINFORCE

"""REINFORCE - the policy gradient, transcribed.

A linear softmax policy, chosen because its score has a closed form you can
read straight off the code:

    d/dtheta_b  log pi(a | s)  =  (1[b == a] - pi(b | s)) * s

No autodiff, so the identity at the centre of the method is visible rather
than delegated.
"""

import math
import random


def softmax(scores):
    # Subtract the max before exponentiating. Without it a large score
    # overflows, and the overflow is silent right up until it is a nan.
    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):
    """theta[a][i] is the weight on feature i for action a."""
    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 run_episode(env, theta, max_steps=1_000):
    """The return is not available until the episode ends, which is why this
    method cannot be applied to a continuing task without truncating it."""
    states, actions, rewards = [], [], []
    state = env.reset()

    for _ in range(max_steps):
        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)

        state = next_state
        if done:
            break

    return states, actions, rewards


def returns_to_go(rewards, gamma):
    """G_t = r_t + gamma * G_(t+1), computed backwards in a single pass.

    Reward-to-GO, not the whole episode return. An action cannot have caused a
    reward that arrived before it, so including earlier rewards adds variance
    and contributes no signal - the first variance reduction, and free.
    """
    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 train(env, n_features, n_actions, episodes=2_000, batch_episodes=10,
          alpha=0.01, gamma=0.99):
    theta = [[0.0] * n_features for _ in range(n_actions)]

    for _ in range(episodes // batch_episodes):
        # Gradients accumulate over a BATCH of episodes. A single-episode
        # gradient is noisy enough that the policy random-walks, so batching
        # here is what makes the estimator usable rather than a throughput
        # optimization.
        gradient = [[0.0] * n_features for _ in range(n_actions)]
        trajectories = []
        every_return = []

        for _ in range(batch_episodes):
            states, actions, rewards = run_episode(env, theta)
            discounted = returns_to_go(rewards, gamma)
            trajectories.append((states, actions, discounted))
            every_return.extend(discounted)

        # The baseline. Subtracting ANY function of state - a constant
        # included - leaves the gradient's mean untouched, because the
        # expected score is zero. It is free, and it is the single largest
        # variance reduction available.
        baseline = sum(every_return) / len(every_return)

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

                # The score, for every action's weights at once: taking action
                # a pushes a's weights up and every other action's down, in
                # proportion to how likely they were.
                for candidate in range(n_actions):
                    indicator = 1.0 if candidate == action else 0.0
                    coefficient = advantage * (indicator - probabilities[candidate])
                    for index, feature in enumerate(state):
                        gradient[candidate][index] += coefficient * feature

        # ASCENT, not descent: expected return is being maximized, and the sign
        # error here is the easiest one in the method to make.
        for candidate in range(n_actions):
            for index in range(n_features):
                theta[candidate][index] += alpha * gradient[candidate][index] / batch_episodes

    return theta

"""REINFORCE - typed, batched, entropy-regularized, advantages normalized."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import NamedTuple, Protocol

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


class Trajectory(NamedTuple):
    states: Tensor        # (T, features)
    actions: Tensor       # (T,)
    rewards: Tensor       # (T,)


class Environment(Protocol):
    def reset(self) -> Tensor: ...
    def step(self, action: int) -> tuple[Tensor, float, bool]: ...


@dataclass(frozen=True)
class ReinforceConfig:
    n_features: int
    n_actions: int
    hidden: tuple[int, ...] = (64, 64)
    learning_rate: float = 3e-4
    gamma: float = 0.99
    episodes_per_batch: int = 16
    # The only thing standing between this policy and premature determinism.
    # The policy is the sole source of exploration, so a distribution that
    # sharpens before the evidence supports it stops sampling the alternatives
    # and never computes their gradients again.
    entropy_coefficient: float = 0.01
    grad_clip_norm: float = 0.5
    normalize_advantages: bool = True

    def __post_init__(self) -> None:
        if not 0.0 <= self.gamma < 1.0:
            raise ValueError(f"gamma must be in [0, 1), got {self.gamma}")
        if self.episodes_per_batch < 2:
            raise ValueError("a single-episode gradient random-walks; batch at least two")
        if self.entropy_coefficient < 0.0:
            raise ValueError("a negative entropy coefficient rewards collapse")


@dataclass
class BatchDiagnostics:
    mean_return: float
    # Entropy is the leading indicator of the characteristic failure. A return
    # curve that plateaus is diagnosed by the entropy curve that collapsed a
    # few hundred updates earlier.
    policy_entropy: float
    grad_norm: float
    history: list[float] = field(default_factory=list)


def build_policy(config: ReinforceConfig) -> nn.Module:
    layers: list[nn.Module] = []
    width = config.n_features
    for size in config.hidden:
        layers += [nn.Linear(width, size), nn.Tanh()]
        width = size
    layers.append(nn.Linear(width, config.n_actions))
    return nn.Sequential(*layers)


def discounted_returns_to_go(rewards: Tensor, gamma: float) -> Tensor:
    """Reward-to-go: an action is credited only with what followed it."""
    out = torch.zeros_like(rewards)
    running = torch.zeros((), dtype=rewards.dtype)
    for step in range(rewards.shape[0] - 1, -1, -1):
        running = rewards[step] + gamma * running
        out[step] = running
    return out


def collect_batch(env: Environment, policy: nn.Module, config: ReinforceConfig) -> list[Trajectory]:
    """Rollouts run under no_grad: the graph is not needed here.

    Storing a per-step autograd graph across a whole batch of episodes holds
    memory proportional to total steps for no reason - the log-probabilities
    are recomputed in one batched pass at update time instead.
    """
    trajectories: list[Trajectory] = []

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

            while not done:
                distribution = Categorical(logits=policy(state))
                action = distribution.sample()
                next_state, reward, done = env.step(int(action))

                states.append(state)
                actions.append(action)
                rewards.append(reward)
                state = next_state

            trajectories.append(
                Trajectory(
                    states=torch.stack(states),
                    actions=torch.stack(actions),
                    rewards=torch.tensor(rewards, dtype=torch.float32),
                )
            )

    return trajectories


def update(
    policy: nn.Module,
    optimizer: torch.optim.Optimizer,
    trajectories: list[Trajectory],
    config: ReinforceConfig,
) -> BatchDiagnostics:
    states = torch.cat([t.states for t in trajectories])
    actions = torch.cat([t.actions for t in trajectories])
    advantages = torch.cat(
        [discounted_returns_to_go(t.rewards, config.gamma) for t in trajectories]
    )

    if config.normalize_advantages:
        # Standardizing within the batch. Technically this biases the
        # estimator, because the statistics depend on the sample - and it is
        # near-universal anyway, because the stabilization is worth more than
        # the bias costs.
        advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)

    distribution = Categorical(logits=policy(states))
    log_probabilities = distribution.log_prob(actions)

    # This quantity's VALUE means nothing. It is constructed so that its
    # gradient equals the policy gradient, and nothing about its magnitude is
    # interpretable. Logging it as a training loss - or tuning against it - is
    # the most common misreading in applied policy gradients.
    surrogate = -(log_probabilities * advantages).mean()
    entropy = distribution.entropy().mean()
    objective = surrogate - config.entropy_coefficient * entropy

    optimizer.zero_grad(set_to_none=True)
    objective.backward()
    grad_norm = nn.utils.clip_grad_norm_(policy.parameters(), config.grad_clip_norm)
    optimizer.step()

    return BatchDiagnostics(
        mean_return=float(torch.stack([t.rewards.sum() for t in trajectories]).mean()),
        policy_entropy=float(entropy),
        grad_norm=float(grad_norm),
    )

"""REINFORCE - parallel rollouts, returns computed without a Python loop.

The structural fact that makes this method cheap to accelerate: the episodes
in a batch are INDEPENDENT. Unlike an off-policy method whose updates are
sequentially dependent, on-policy data collection within one batch is
embarrassingly parallel, and collection is where nearly all the wall-clock is.
"""

from __future__ import annotations

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


def batched_returns_to_go(
    rewards: NDArray[np.float32],
    episode_starts: NDArray[np.bool_],
    gamma: float,
    out: NDArray[np.float32],
) -> NDArray[np.float32]:
    """Reward-to-go for a whole flattened batch, one reversed pass.

    The recursion G_t = r_t + gamma * G_(t+1) is genuinely sequential in time,
    so it cannot be vectorized along the time axis. What it can do is run once
    over the concatenated batch instead of once per episode, with the running
    total reset at each episode boundary - which turns B separate Python loops
    into one, and writes into a buffer that was allocated once.
    """
    running = np.float32(0.0)
    for position in range(len(rewards) - 1, -1, -1):
        # A boundary resets the accumulator, so returns never leak across
        # episodes. Getting this wrong credits one episode's actions with the
        # next episode's rewards, and the result still trains - just worse.
        if episode_starts[position]:
            running = np.float32(0.0)
        running = rewards[position] + gamma * running
        out[position] = running
    return out


class RolloutBuffer:
    """Pre-allocated storage for one batch. Sized once, written in place.

    Appending to Python lists and stacking at the end allocates on every step
    and copies the whole batch at the boundary; a fixed buffer does neither.
    """

    def __init__(self, capacity: int, n_envs: int, n_features: int) -> None:
        self.states = np.zeros((capacity, n_envs, n_features), dtype=np.float32)
        self.actions = np.zeros((capacity, n_envs), dtype=np.int64)
        self.rewards = np.zeros((capacity, n_envs), dtype=np.float32)
        self.starts = np.zeros((capacity, n_envs), dtype=np.bool_)
        self._cursor = 0

    def add(
        self,
        states: NDArray[np.float32],
        actions: NDArray[np.int64],
        rewards: NDArray[np.float32],
        starts: NDArray[np.bool_],
    ) -> None:
        index = self._cursor
        self.states[index] = states      # in place; nothing allocated per step
        self.actions[index] = actions
        self.rewards[index] = rewards
        self.starts[index] = starts
        self._cursor += 1

    def flatten(self) -> tuple[Tensor, Tensor, NDArray[np.float32], NDArray[np.bool_]]:
        """Collapse (T, N) into (T*N,) with one reshape and no copy.

        Environment-major or time-major does not matter to the estimator: the
        gradient is a sum over independent (state, action, advantage) triples,
        so the batch can be flattened however keeps memory contiguous.
        """
        steps = self._cursor
        states = self.states[:steps].reshape(-1, self.states.shape[-1])
        return (
            torch.from_numpy(states),
            torch.from_numpy(self.actions[:steps].reshape(-1)),
            self.rewards[:steps].reshape(-1),
            self.starts[:steps].reshape(-1),
        )


@torch.no_grad()
def act_vectorized(policy: torch.nn.Module, states: Tensor, generator: torch.Generator) -> Tensor:
    """Sample an action for every environment from one forward pass.

    Gumbel-max rather than a Categorical object: sampling from a softmax is
    argmax(logits + gumbel noise), which needs no normalization, no per-env
    distribution object, and no exponentials at all.
    """
    logits = policy(states)
    uniform = torch.rand(logits.shape, generator=generator, device=logits.device)
    gumbel = -torch.log(-torch.log(uniform + 1e-20) + 1e-20)
    return (logits + gumbel).argmax(dim=1)


def surrogate_gradient_step(
    policy: torch.nn.Module,
    optimizer: torch.optim.Optimizer,
    states: Tensor,
    actions: Tensor,
    advantages: Tensor,
    entropy_coefficient: float,
) -> None:
    """One forward over the entire flattened batch.

    Log-probabilities are recomputed here rather than retained from rollout,
    which is both faster and dramatically cheaper in memory: no autograd graph
    is held across thousands of environment steps.
    """
    logits = policy(states)
    log_probabilities = torch.log_softmax(logits, dim=1)
    # gather, not a Python loop over samples: one indexed read per row.
    chosen = log_probabilities.gather(1, actions.unsqueeze(1)).squeeze(1)

    entropy = -(log_probabilities.exp() * log_probabilities).sum(dim=1).mean()
    objective = -(chosen * advantages).mean() - entropy_coefficient * entropy

    optimizer.zero_grad(set_to_none=True)
    objective.backward()
    optimizer.step()