SARSA

"""Tabular SARSA - the on-policy update, transcribed.

One thing to watch above all others: a' is chosen BEFORE the update and then
executed. Drawing it twice - once to form the target, once to act - is neither
SARSA nor Q-learning, and nothing in any metric will tell you.
"""

import random


def epsilon_greedy(q_row, n_actions, epsilon):
    """The behaviour policy. It is also the policy being evaluated - that
    identity is the whole method."""
    if random.random() < epsilon:
        return random.randrange(n_actions)
    return max(range(n_actions), key=lambda a: q_row[a])


def train(env, n_states, n_actions, episodes=5_000, alpha=0.1, gamma=0.99, epsilon=0.1):
    # Q[s][a] = expected discounted return from taking a in s and then
    # continuing to behave epsilon-greedily. Not optimally. Epsilon-greedily.
    Q = [[0.0] * n_actions for _ in range(n_states)]

    for _ in range(episodes):
        state = env.reset()
        # The first action is drawn before the loop: every iteration needs an
        # action already in hand, because the update it forms depends on it.
        action = epsilon_greedy(Q[state], n_actions, epsilon)
        done = False

        while not done:
            next_state, reward, done = env.step(action)

            # Draw the next action NOW, from the same policy. This value is the
            # a' inside the target AND the action executed on the next pass of
            # the loop. One draw, used for both. Two draws is the classic bug.
            next_action = 0 if done else epsilon_greedy(Q[next_state], n_actions, epsilon)

            # Q[s'][a'], not max over a of Q[s'][a]. That single substitution
            # is the entire difference from Q-learning, and it is why this
            # learns the value of a policy that sometimes acts at random.
            bootstrap = 0.0 if done else Q[next_state][next_action]

            td_target = reward + gamma * bootstrap
            Q[state][action] += alpha * (td_target - Q[state][action])

            state, action = next_state, next_action

    return Q


def expected_sarsa_target(q_row, n_actions, epsilon, gamma, reward):
    """Expected SARSA: compute the expectation instead of sampling it.

    Epsilon-greedy puts epsilon/|A| on every action plus 1 - epsilon on the
    greedy one. Summing that out by hand removes the variance of the draw and
    leaves the bias untouched. Set epsilon to 0 here and the sum collapses to a
    max - which is Q-learning. The two methods are endpoints of one family.
    """
    greedy = max(range(n_actions), key=lambda a: q_row[a])

    expectation = 0.0
    for a in range(n_actions):
        prob = epsilon / n_actions
        if a == greedy:
            prob += 1.0 - epsilon
        expectation += prob * q_row[a]

    return reward + gamma * expectation

"""Tabular SARSA - typed, annealed, with the on-policy invariant asserted."""

from dataclasses import dataclass, field
from typing import Protocol

import numpy as np
from numpy.typing import NDArray

QTable = NDArray[np.float64]


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


@dataclass(frozen=True)
class SarsaConfig:
    n_states: int
    n_actions: int
    alpha: float = 0.1
    gamma: float = 0.99
    epsilon_start: float = 1.0
    epsilon_end: float = 0.01
    epsilon_decay_episodes: int = 1_000
    # Sampled target vs. the exact expectation. Expected SARSA is strictly
    # less noisy for the cost of a sum over actions, so it is the default
    # whenever the action set is small enough to sum over.
    expected: 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 not 0.0 < self.alpha <= 1.0:
            raise ValueError(f"alpha must be in (0, 1], got {self.alpha}")
        if self.n_states < 1 or self.n_actions < 1:
            raise ValueError("state and action spaces must be non-empty")


@dataclass
class TrainingResult:
    q_table: QTable
    # SARSA's advantage is online performance WHILE exploring, so the return
    # under the behaviour policy is the quantity that matters. Recording only
    # the greedy policy's return measures the one thing this does not optimize.
    behaviour_returns: list[float] = field(default_factory=list)
    greedy_returns: list[float] = field(default_factory=list)
    mean_abs_td_error: list[float] = field(default_factory=list)

    @property
    def exploration_cost(self) -> float:
        """Gap between what the greedy policy would earn and what the agent
        actually earns while still exploring. It should shrink as epsilon
        anneals; if it does not, the anneal has stalled."""
        if not self.behaviour_returns or not self.greedy_returns:
            return 0.0
        window = min(50, len(self.behaviour_returns), len(self.greedy_returns))
        return float(
            np.mean(self.greedy_returns[-window:]) - np.mean(self.behaviour_returns[-window:])
        )


def _policy_probabilities(q_row: NDArray[np.float64], epsilon: float) -> NDArray[np.float64]:
    """Epsilon-greedy as an explicit distribution.

    Writing the policy down as probabilities rather than as a sampling
    procedure is what makes Expected SARSA one line, and it makes the
    on-policy claim auditable: this is the distribution the target integrates
    over and it must be the one the agent samples from.
    """
    n_actions = q_row.shape[0]
    probabilities = np.full(n_actions, epsilon / n_actions, dtype=np.float64)
    probabilities[int(np.argmax(q_row))] += 1.0 - epsilon
    return probabilities


def _anneal(episode: int, config: SarsaConfig) -> float:
    progress = min(1.0, episode / config.epsilon_decay_episodes)
    return config.epsilon_start + progress * (config.epsilon_end - config.epsilon_start)


def train(env: Environment, config: SarsaConfig, episodes: int = 5_000) -> TrainingResult:
    q = np.zeros((config.n_states, config.n_actions), dtype=np.float64)
    rng = np.random.default_rng(0)
    result = TrainingResult(q_table=q)

    for episode in range(episodes):
        # Fixed epsilon converges to the value of the epsilon-soft policy, not
        # the optimal one. Both are correct answers; only one of them is the
        # answer people usually report.
        epsilon = _anneal(episode, config)

        state = env.reset()
        probabilities = _policy_probabilities(q[state], epsilon)
        action = int(rng.choice(config.n_actions, p=probabilities))

        done = False
        behaviour_return = 0.0
        td_errors: list[float] = []

        while not done:
            next_state, reward, done = env.step(action)

            if done:
                # No future past a terminal. Bootstrapping past one is the
                # single most common bug in any TD implementation.
                bootstrap = 0.0
                next_action = 0
            else:
                next_probabilities = _policy_probabilities(q[next_state], epsilon)
                # Draw once. This action is both the a' in the target and the
                # action executed next iteration - the on-policy invariant.
                next_action = int(rng.choice(config.n_actions, p=next_probabilities))
                bootstrap = (
                    float(next_probabilities @ q[next_state])
                    if config.expected
                    else float(q[next_state, next_action])
                )

            td_error = reward + config.gamma * bootstrap - q[state, action]
            q[state, action] += config.alpha * td_error

            behaviour_return += reward
            td_errors.append(abs(td_error))
            state, action = next_state, next_action

        result.behaviour_returns.append(behaviour_return)
        result.mean_abs_td_error.append(float(np.mean(td_errors)) if td_errors else 0.0)

    return result

"""Expected SARSA - vectorized across environments stepping in lockstep.

Note what is absent: a replay buffer. On-policy learning cannot reuse
experience collected under a previous policy, so the only parallelism
available is breadth - many environments running the CURRENT policy at once -
rather than depth into stored history. That constraint is the method, not an
oversight, and it is the reason deep RL was built on off-policy methods.
"""

import numpy as np
from numpy.typing import NDArray

QTable = NDArray[np.float64]


def expected_bootstrap(
    q: QTable,
    next_states: NDArray[np.intp],
    epsilon: float,
    out: NDArray[np.float64],
) -> NDArray[np.float64]:
    """Expectation under epsilon-greedy, in closed form, for a whole batch.

    The sum over actions never needs to be built explicitly. Epsilon-greedy is
    a uniform mixture plus a point mass on the greedy action, so its
    expectation is exactly epsilon * mean(row) + (1 - epsilon) * max(row) -
    two reductions instead of a weighted sum, and no probability vector is
    materialized. Setting epsilon to 0 leaves max(row), which is the
    Q-learning target; the family really is one parameterization.
    """
    rows = q[next_states]                       # one gather, (batch, |A|)
    np.max(rows, axis=1, out=out)               # greedy component, in place
    out *= 1.0 - epsilon
    out += epsilon * rows.mean(axis=1)          # uniform component
    return out


def run_vectorized(
    envs,
    n_states: int,
    n_actions: int,
    steps: int = 100_000,
    alpha: float = 0.1,
    gamma: float = 0.99,
    epsilon: float = 0.1,
) -> QTable:
    rng = np.random.default_rng(0)
    n_envs = len(envs)

    q = np.zeros((n_states, n_actions), dtype=np.float64)

    # Scratch buffers allocated once and reused every step, so the hot loop
    # allocates nothing at all.
    bootstrap = np.empty(n_envs, dtype=np.float64)
    actions = np.empty(n_envs, dtype=np.intp)

    states = np.array([env.reset() for env in envs], dtype=np.intp)
    actions[:] = select_actions(q, states, epsilon, n_actions, rng)

    for _ in range(steps // n_envs):
        results = [env.step(int(a)) for env, a in zip(envs, actions)]
        next_states = np.fromiter((r[0] for r in results), dtype=np.intp, count=n_envs)
        rewards = np.fromiter((r[1] for r in results), dtype=np.float64, count=n_envs)
        dones = np.fromiter((r[2] for r in results), dtype=np.bool_, count=n_envs)

        expected_bootstrap(q, next_states, epsilon, bootstrap)
        bootstrap[dones] = 0.0

        targets = rewards + gamma * bootstrap
        deltas = alpha * (targets - q[states, actions])

        # np.add.at, not q[states, actions] += deltas: when two environments
        # visit the same (s, a) in one batch, plain fancy-index assignment
        # keeps only the last write and silently drops the others.
        np.add.at(q, (states, actions), deltas)

        # Choosing the next actions AFTER the update rather than before is
        # deliberate: it keeps the behaviour policy synchronized with the table
        # the agent is about to act on, which is what on-policy means once the
        # updates are batched.
        states = np.where(dones, [env.reset() for env in envs], next_states)
        actions[:] = select_actions(q, states, epsilon, n_actions, rng)

    return q


def select_actions(
    q: QTable,
    states: NDArray[np.intp],
    epsilon: float,
    n_actions: int,
    rng: np.random.Generator,
) -> NDArray[np.intp]:
    """Epsilon-greedy for every environment in one pass - one RNG draw for the
    explore mask, one argmax over the gathered rows, one select."""
    explore = rng.random(states.shape[0]) < epsilon
    greedy = np.argmax(q[states], axis=1)
    return np.where(explore, rng.integers(n_actions, size=states.shape[0]), greedy).astype(np.intp)