Temporal-Difference Learning: TD(0) & TD(lambda)

"""Temporal-difference learning, transcribed from the objective.

No library. The one substitution that defines the method:

    target = R + gamma * V(S')       instead of     target = G_t

Four things to read for.

1. td_zero() and monte_carlo() differ in exactly one line. Everything
   said about bias, variance, online learning and the Markov property
   follows from that line and nothing else.

2. lambda_return() is the forward view and td_lambda() is the backward
   view. forward_backward_equivalence() checks numerically that they
   produce the same total update offline -- which is the whole reason
   the trace mechanism exists, since the forward view refers to the
   future and cannot be computed online.

3. batch_td_versus_batch_mc() is the sharpest statement of what
   bootstrapping does. Given identical data, Monte Carlo converges to
   the least-squares fit of the observed returns and TD converges to
   the value function of the maximum-likelihood Markov model. Different
   answers, and TD's is better exactly when the process is Markov.

4. aliasing_demo() is the same comparison with the Markov property
   removed, where the sign of the argument flips.
"""

import random

Transition = tuple[int, float, int]   # (state, reward, next_state)
Episode = list[Transition]

TERMINAL = -1


def td_zero(
    episodes: list[Episode], n_states: int, gamma: float, alpha: float, initial: float = 0.0
) -> list[float]:
    """One-step bootstrapping. The whole method is the target line.

    Note what does NOT appear: any reference to the end of the episode.
    The update is available at every step, which is what makes this
    usable on a process that never terminates -- the property Monte
    Carlo cannot match at any sample size.
    """
    values = [initial] * n_states

    for episode in episodes:
        for state, reward, next_state in episode:
            # The target contains V. That is bootstrapping, and every
            # property of this method comes from it: the variance
            # collapses because the target is one reward plus a lookup,
            # and the bias appears because the lookup is wrong early
            # and its error propagates backwards.
            bootstrapped = 0.0 if next_state == TERMINAL else values[next_state]
            td_error = reward + gamma * bootstrapped - values[state]
            values[state] += alpha * td_error

    return values


def monte_carlo(
    episodes: list[Episode], n_states: int, gamma: float, alpha: float, initial: float = 0.0
) -> list[float]:
    """The same loop with the realized return as the target.

    Read against td_zero() above: one line differs. This one waits for
    the episode to end, which is both why it is unbiased and why it
    cannot run on a continuing task.
    """
    values = [initial] * n_states

    for episode in episodes:
        returns = [0.0] * len(episode)
        running = 0.0
        for index in reversed(range(len(episode))):
            running = episode[index][1] + gamma * running
            returns[index] = running

        for index, (state, _, _) in enumerate(episode):
            values[state] += alpha * (returns[index] - values[state])

    return values


def n_step_return(
    episode: Episode, start: int, n: int, values: list[float], gamma: float
) -> float:
    """n real rewards, then bootstrap.

    The discrete version of the dial. n = 1 is TD(0), n at or beyond the
    episode length is Monte Carlo, and the useful values are in
    between -- which is the whole observation that lambda then makes
    continuous.
    """
    total = 0.0
    discount = 1.0

    for offset in range(n):
        index = start + offset
        if index >= len(episode):
            return total  # ran off the end: this IS the full return
        total += discount * episode[index][1]
        discount *= gamma

        if episode[index][2] == TERMINAL:
            return total

    return total + discount * values[episode[start + n - 1][2]]


def lambda_return(
    episode: Episode, start: int, values: list[float], gamma: float, lam: float
) -> float:
    """The forward view: a geometric average over every n-step return.

        G^lambda = (1 - lambda) * sum_n lambda^(n-1) * G^(n)

    This is the target the method is actually trying to hit, and it
    cannot be computed online because it refers to the whole remaining
    episode. The trace mechanism below exists entirely to produce the
    same updates without ever forming this quantity.
    """
    horizon = len(episode) - start
    total = 0.0
    weight = 1.0 - lam

    for n in range(1, horizon):
        total += weight * n_step_return(episode, start, n, values, gamma)
        weight *= lam

    # The remaining weight, lambda^(horizon-1), goes to the full return.
    return total + (lam ** (horizon - 1)) * n_step_return(episode, start, horizon, values, gamma)


def td_lambda(
    episodes: list[Episode],
    n_states: int,
    gamma: float,
    alpha: float,
    lam: float,
    replacing: bool = True,
    initial: float = 0.0,
) -> list[float]:
    """The backward view: one TD error, distributed by eligibility.

    The trace records how recently and how often each state was
    visited, decaying at gamma*lambda per step. Applying today's error
    in proportion to it credits the states that led here, without any
    lookahead, buffer or episode boundary.

    Accumulating traces add 1 on each visit and can grow without bound
    when a state is revisited quickly, which multiplies the effective
    step size and destabilizes runs that look like learning-rate
    problems. Replacing traces cap at 1 and are the safer default.
    """
    values = [initial] * n_states

    for episode in episodes:
        traces = [0.0] * n_states

        for state, reward, next_state in episode:
            bootstrapped = 0.0 if next_state == TERMINAL else values[next_state]
            td_error = reward + gamma * bootstrapped - values[state]

            traces[state] = 1.0 if replacing else traces[state] + 1.0

            # Every eligible state is updated by the same error. This
            # is O(|S|) per step as written, which is why real
            # implementations prune traces below a cutoff -- the decay
            # is geometric, so most of these contribute nothing.
            for index in range(n_states):
                if traces[index] == 0.0:
                    continue
                values[index] += alpha * td_error * traces[index]
                traces[index] *= gamma * lam

    return values


def forward_backward_equivalence(
    episode: Episode, n_states: int, gamma: float, alpha: float, lam: float
) -> dict[str, float]:
    """The equivalence that justifies the trace mechanism.

    Offline -- with the value function held fixed for the whole episode
    -- the total update from the backward view equals the total update
    from the forward view. That is why a decaying memory of the past
    can implement an objective defined over the future.

    The equality is exact offline and approximate online, because the
    online version bootstraps off values that changed during the
    episode. True-online TD(lambda) with dutch traces restores exactness
    and is what a serious implementation uses.
    """
    frozen = [0.0] * n_states

    forward = [0.0] * n_states
    for index, (state, _, _) in enumerate(episode):
        target = lambda_return(episode, index, frozen, gamma, lam)
        forward[state] += alpha * (target - frozen[state])

    backward = [0.0] * n_states
    traces = [0.0] * n_states
    for state, reward, next_state in episode:
        bootstrapped = 0.0 if next_state == TERMINAL else frozen[next_state]
        td_error = reward + gamma * bootstrapped - frozen[state]
        traces[state] += 1.0

        for index in range(n_states):
            if traces[index] == 0.0:
                continue
            backward[index] += alpha * td_error * traces[index]
            traces[index] *= gamma * lam

    largest_gap = max(abs(f - b) for f, b in zip(forward, backward))
    return {
        "largest_gap": largest_gap,
        "equivalent": largest_gap < 1e-9,
    }


def batch_td_versus_batch_mc(gamma: float = 1.0, sweeps: int = 500) -> dict[str, object]:
    """The clearest statement of what bootstrapping actually does.

    The classic eight-episode example. State A appears once, followed
    immediately by state B with reward 0; state B appears eight times,
    with reward 1 in six of them.

    Monte Carlo fits the observed returns: A was followed by a total
    reward of 0 exactly once, so V(A) = 0. It is the least-squares
    answer and it is unarguable given the data as a flat collection of
    outcomes.

    TD uses the transition structure: A always goes to B, and B is worth
    0.75, so V(A) = 0.75. It is the value function of the
    maximum-likelihood Markov model of the data.

    Both are correct answers to different questions, and TD's is the one
    you want exactly when the process really is Markov -- which is the
    entire argument for bootstrapping, stated in two numbers.
    """
    episodes: list[Episode] = [[(0, 0.0, 1), (1, 1.0, TERMINAL)]]
    episodes += [[(1, 1.0, TERMINAL)] for _ in range(5)]
    episodes += [[(1, 0.0, TERMINAL)] for _ in range(2)]

    # Batch: repeatedly sweep the same episodes with a small step until
    # each method reaches its own fixed point.
    td_values = [0.0, 0.0]
    for _ in range(sweeps):
        for episode in episodes:
            for state, reward, next_state in episode:
                bootstrapped = 0.0 if next_state == TERMINAL else td_values[next_state]
                td_values[state] += 0.01 * (reward + gamma * bootstrapped - td_values[state])

    mc_values = [0.0, 0.0]
    for _ in range(sweeps):
        for episode in episodes:
            running = 0.0
            returns = [0.0] * len(episode)
            for index in reversed(range(len(episode))):
                running = episode[index][1] + gamma * running
                returns[index] = running
            for index, (state, _, _) in enumerate(episode):
                mc_values[state] += 0.01 * (returns[index] - mc_values[state])

    return {
        "td_value_of_A": td_values[0],
        "mc_value_of_A": mc_values[0],
        "value_of_B": td_values[1],
        "note": (
            "same data, different answers: MC fits the observed returns, TD fits the "
            "maximum-likelihood Markov model"
        ),
    }


def aliasing_demo(episodes_count: int = 40_000, gamma: float = 0.9, seed: int = 0):
    """Where the argument for bootstrapping reverses.

    One observation hides two underlying states whose continuations
    differ. Monte Carlo averages the returns that actually occurred and
    is correct for the process as experienced. TD bootstraps off the
    aliased successor's estimate, which summarizes nothing, and
    converges to a value that is not the value of anything.

    The gap between the two estimators is therefore a direct
    measurement of how non-Markov the state representation is -- a
    diagnostic worth running whenever a state was constructed rather
    than given.
    """
    rng = random.Random(seed)
    episodes: list[Episode] = []

    for _ in range(episodes_count):
        # Observation 0 leads to observation 1, which is really two
        # different situations paying +1 and -1 with equal probability.
        payoff = 1.0 if rng.random() < 0.5 else -1.0
        episodes.append([(0, 0.0, 1), (1, payoff, TERMINAL)])

    return {
        "td_value_of_0": td_zero(episodes, 2, gamma, 0.01)[0],
        "mc_value_of_0": monte_carlo(episodes, 2, gamma, 0.01)[0],
        "note": "the gap between them measures how much the state representation hides",
    }

"""TD learning with the choices that change the answer made explicit.

Four things the literal version decided silently:

  * whether the step size decays, which decides whether the run
    converges or merely tracks -- and almost every implementation calls
    the second one "converged";
  * which trace type is in use, where accumulating traces can diverge
    on their own through fast revisits;
  * whether all three deadly-triad ingredients are present, which is
    the difference between slow learning and exponential divergence;
  * where to stop decaying traces, which is an approximation rather
    than an optimization and should be stated.

Each becomes a named object or a raised exception. The diagnostic that
matters most -- the gap between the bootstrapped and Monte Carlo
estimates, which measures how non-Markov the state is -- is a function
rather than a paragraph.
"""

from __future__ import annotations

from collections.abc import Sequence
from dataclasses import dataclass, field
from enum import Enum
from typing import NamedTuple

import numpy as np
import numpy.typing as npt

FloatArray = npt.NDArray[np.float64]

TERMINAL = -1


class TemporalDifferenceError(Exception):
    """Base for every recoverable failure in these learners."""


class InvalidSchedule(TemporalDifferenceError):
    """A step size, discount or lambda outside its meaningful range."""


class Divergence(TemporalDifferenceError):
    """Values left the range any bounded return could produce.

    Raised rather than allowed to continue because deadly-triad failures
    grow geometrically: by the time a value function is visibly absurd
    it has usually been diverging for thousands of steps, and no step
    size recovers it.
    """


class DeadlyTriad(TemporalDifferenceError):
    """Bootstrapping, function approximation and off-policy updates.

    Individually standard; together they can diverge rather than
    converge slowly. Any two are safe, which is exactly why this
    configuration is reached by accident and then debugged as a
    learning-rate problem.
    """


class TraceType(Enum):
    """How a revisit affects eligibility.

    ACCUMULATING adds one and can grow without bound on fast revisits,
    which multiplies the effective step size. REPLACING caps at one and
    is the safe default. DUTCH is the form that makes the online update
    exactly equal the offline one -- true-online TD(lambda).
    """

    ACCUMULATING = "accumulating"
    REPLACING = "replacing"
    DUTCH = "dutch"


@dataclass(frozen=True)
class StepSizeSchedule:
    """Alpha over time, with the guarantee it does or does not carry.

    Decaying as one over visit count to a power in (0.5, 1] satisfies
    the Robbins-Monro conditions and converges. A constant alpha tracks
    a changing environment and converges to nothing -- frequently the
    right choice, and a different claim.
    """

    initial: float
    decay_exponent: float = 0.0

    def __post_init__(self) -> None:
        if not 0.0 < self.initial <= 1.0:
            raise InvalidSchedule(f"initial step size must lie in (0, 1]; got {self.initial}")
        if not 0.0 <= self.decay_exponent <= 1.0:
            raise InvalidSchedule(f"decay exponent must lie in [0, 1]; got {self.decay_exponent}")

    def at(self, visits: int) -> float:
        if self.decay_exponent == 0.0:
            return self.initial
        return self.initial / max(visits, 1) ** self.decay_exponent

    @property
    def satisfies_robbins_monro(self) -> bool:
        """Whether this schedule converges, or only tracks.

        Exposed rather than documented, because a run reported as
        converged under a constant step size is describing something
        that has no fixed point.
        """
        return 0.5 < self.decay_exponent <= 1.0


@dataclass(frozen=True)
class LearnerConfig:
    """Everything that changes what the run computes, validated once."""

    gamma: float
    lam: float
    schedule: StepSizeSchedule
    trace_type: TraceType = TraceType.REPLACING
    # Traces decay geometrically, so below this they contribute nothing
    # measurable. Pruning is what makes a large state space affordable,
    # and it is an approximation rather than a free optimization.
    trace_cutoff: float = 1e-4
    uses_function_approximation: bool = False
    is_off_policy: bool = False

    def __post_init__(self) -> None:
        if not 0.0 <= self.gamma <= 1.0:
            raise InvalidSchedule(f"gamma must lie in [0, 1]; got {self.gamma}")
        if not 0.0 <= self.lam <= 1.0:
            raise InvalidSchedule(f"lambda must lie in [0, 1]; got {self.lam}")
        if self.trace_cutoff <= 0.0:
            raise InvalidSchedule(f"trace cutoff must be positive; got {self.trace_cutoff}")

        # Bootstrapping is present whenever lambda < 1: at lambda = 1
        # the method is Monte Carlo and the triad does not apply.
        bootstraps = self.lam < 1.0
        if bootstraps and self.uses_function_approximation and self.is_off_policy:
            raise DeadlyTriad(
                "bootstrapping, function approximation and off-policy updates together "
                "can diverge exponentially. Remove one: raise lambda to 1, use a table, "
                "or sample on-policy -- no step size fixes this combination"
            )

    @property
    def trace_decay(self) -> float:
        """gamma * lambda. Worth naming, because lowering the discount
        silently shortens credit assignment as well as the horizon."""
        return self.gamma * self.lam

    @property
    def effective_credit_horizon(self) -> float:
        """Roughly how many steps back a TD error is still felt."""
        decay = self.trace_decay
        return float("inf") if decay >= 1.0 else 1.0 / (1.0 - decay)


class LearningReport(NamedTuple):
    """The values, plus whether to believe them."""

    values: FloatArray
    visits: npt.NDArray[np.int64]
    mean_td_error: float
    converges: bool
    active_traces_mean: float


@dataclass
class TdLambdaLearner:
    """Tabular TD(lambda) with pruned traces.

    Traces are held as a dict of the states whose eligibility is still
    above the cutoff, rather than as a dense array swept every step.
    The dense form is O(|S|) per step for contributions that are
    numerically zero; this is O(k) for a k that stays small because the
    decay is geometric.
    """

    n_states: int
    config: LearnerConfig
    values: FloatArray = field(init=False)
    visits: npt.NDArray[np.int64] = field(init=False)
    _traces: dict[int, float] = field(init=False, default_factory=dict)
    _td_error_sum: float = field(init=False, default=0.0)
    _steps: int = field(init=False, default=0)

    def __post_init__(self) -> None:
        if self.n_states <= 0:
            raise TemporalDifferenceError(f"n_states must be positive; got {self.n_states}")
        self.values = np.zeros(self.n_states, dtype=np.float64)
        self.visits = np.zeros(self.n_states, dtype=np.int64)

    def begin_episode(self) -> None:
        """Traces do not survive a terminal state.

        Carrying them across an episode boundary credits the new
        episode's first reward to the previous episode's states, which
        is a silent and surprisingly common bug.
        """
        self._traces.clear()

    def step(self, state: int, reward: float, next_state: int) -> float:
        """One transition. Returns the TD error, which is the diagnostic.

        Guard clauses first: the interesting path is the update and
        every other path is an early exit.
        """
        if not 0 <= state < self.n_states:
            raise TemporalDifferenceError(f"state {state} outside [0, {self.n_states})")
        if next_state != TERMINAL and not 0 <= next_state < self.n_states:
            raise TemporalDifferenceError(f"next state {next_state} outside the table")

        bootstrapped = 0.0 if next_state == TERMINAL else float(self.values[next_state])
        td_error = reward + self.config.gamma * bootstrapped - float(self.values[state])

        self.visits[state] += 1
        alpha = self.config.schedule.at(int(self.visits[state]))

        current = self._traces.get(state, 0.0)
        if self.config.trace_type is TraceType.ACCUMULATING:
            self._traces[state] = current + 1.0
        elif self.config.trace_type is TraceType.REPLACING:
            self._traces[state] = 1.0
        else:
            # Dutch: the form that makes the online update exactly
            # equal the offline lambda-return update.
            self._traces[state] = (1.0 - alpha) * current + 1.0

        decay = self.config.trace_decay
        surviving: dict[int, float] = {}

        for eligible, trace in self._traces.items():
            self.values[eligible] += alpha * td_error * trace
            decayed = trace * decay
            if decayed >= self.config.trace_cutoff:
                surviving[eligible] = decayed
        self._traces = surviving

        self._td_error_sum += td_error
        self._steps += 1

        # Divergence grows geometrically, so catching it early is the
        # difference between losing a step and losing a run.
        largest = float(np.max(np.abs(self.values)))
        if not np.isfinite(largest) or largest > 1e12:
            raise Divergence(
                f"value magnitude reached {largest:.3e}; with bootstrapping this grows "
                "geometrically, so check the trace type and the triad before the step size"
            )
        return td_error

    def report(self) -> LearningReport:
        return LearningReport(
            values=self.values.copy(),
            visits=self.visits.copy(),
            # Should sit near zero once converged. A sustained
            # departure means the dynamics moved, not that learning is
            # incomplete.
            mean_td_error=self._td_error_sum / max(self._steps, 1),
            converges=self.config.schedule.satisfies_robbins_monro,
            active_traces_mean=float(len(self._traces)),
        )


class Transition(NamedTuple):
    state: int
    reward: float
    next_state: int


class MarkovGap(NamedTuple):
    """How much the state representation is hiding.

    TD bootstraps off the successor's estimate and is biased when that
    estimate summarizes nothing; Monte Carlo averages realized returns
    and is not. The gap between them on the same episodes is therefore
    a direct measurement of the aliasing -- the one diagnostic that
    catches the failure mode bootstrapping has and Monte Carlo does not.
    """

    td_values: FloatArray
    mc_values: FloatArray
    largest_gap: float
    worst_state: int


def markov_gap(
    episodes: Sequence[Sequence[Transition]], n_states: int, gamma: float, alpha: float
) -> MarkovGap:
    """Run both estimators on the same data and compare.

    Worth running whenever the state was constructed rather than given,
    which is almost always. A large gap does not say which estimator is
    right; it says the Markov assumption is doing work it may not be
    able to support.
    """
    if not episodes:
        raise TemporalDifferenceError("no episodes")

    td_values = np.zeros(n_states, dtype=np.float64)
    mc_values = np.zeros(n_states, dtype=np.float64)

    for episode in episodes:
        for transition in episode:
            bootstrapped = (
                0.0 if transition.next_state == TERMINAL else float(td_values[transition.next_state])
            )
            td_values[transition.state] += alpha * (
                transition.reward + gamma * bootstrapped - td_values[transition.state]
            )

        rewards = np.fromiter((t.reward for t in episode), dtype=np.float64, count=len(episode))
        running = 0.0
        returns = np.empty_like(rewards)
        for index in range(len(rewards) - 1, -1, -1):
            running = rewards[index] + gamma * running
            returns[index] = running

        for index, transition in enumerate(episode):
            mc_values[transition.state] += alpha * (returns[index] - mc_values[transition.state])

    differences = np.abs(td_values - mc_values)
    worst = int(np.argmax(differences))
    return MarkovGap(td_values, mc_values, float(differences[worst]), worst)

"""TD(lambda) with linear function approximation and sparse traces.

This is where the method is actually used, and it is also where two of
the three deadly-triad ingredients arrive together -- so the off-policy
question stops being theoretical.

Three changes carry the weight:

  1. traces live over FEATURES, not states, and are kept as an active
     set rather than a dense vector, because tile-coded and hashed
     features activate a handful of components and the trace decays
     geometrically;
  2. the update fuses the weight change and the trace decay into one
     pass over that active set;
  3. true-online TD(lambda) with dutch traces makes the online update
     exactly equal the offline one, for one extra scalar per step.

What none of this changes: the bias from bootstrapping, and the
requirement that the state be Markov. A faster wrong answer is still
wrong.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np
import numpy.typing as npt

FloatArray = npt.NDArray[np.float64]
IndexArray = npt.NDArray[np.int32]


@dataclass(frozen=True)
class LinearConfig:
    n_features: int
    gamma: float
    lam: float
    alpha: float
    trace_cutoff: float = 1e-4


class SparseTrace:
    """Eligibility over features, as an active set.

    A dense trace is O(d) per step to decay, and almost all of those
    components are numerically zero because the decay is geometric --
    at gamma*lambda of 0.9 a component falls below 1e-4 in under a
    hundred steps regardless of how large d is.

    Holding the active components as parallel index and value arrays
    keeps both the decay and the update O(k). The arrays are sized once
    at a capacity the geometric decay bounds, so the hot path allocates
    nothing.
    """

    def __init__(self, capacity: int) -> None:
        self.indices = np.zeros(capacity, dtype=np.int32)
        self.values = np.zeros(capacity, dtype=np.float64)
        self.size = 0
        # Position of each feature in the active arrays, or -1. One
        # lookup instead of a scan when a feature is re-activated.
        self._slot: dict[int, int] = {}

    def clear(self) -> None:
        """Traces do not survive a terminal state.

        Carrying them across an episode boundary credits the next
        episode's first reward to the previous episode's features --
        silent, and a surprisingly common bug.
        """
        self.size = 0
        self._slot.clear()

    def bump(self, feature: int, increment: float, dutch_factor: float) -> None:
        existing = self._slot.get(feature, -1)
        if existing >= 0:
            # Dutch trace: (1 - alpha*x'e) * e + x, the form that makes
            # the online update exactly equal the offline one.
            self.values[existing] = dutch_factor * self.values[existing] + increment
            return

        if self.size >= len(self.indices):
            raise RuntimeError("trace capacity exceeded; raise the cutoff or the capacity")

        self.indices[self.size] = feature
        self.values[self.size] = increment
        self._slot[feature] = self.size
        self.size += 1

    def active(self) -> tuple[IndexArray, FloatArray]:
        return self.indices[: self.size], self.values[: self.size]

    def decay_and_prune(self, decay: float, cutoff: float) -> None:
        """One pass: scale, drop the negligible, compact.

        Compacting in place rather than rebuilding keeps the arrays
        contiguous, which is what makes the gather in the next step
        stream rather than stride.
        """
        write = 0
        for read in range(self.size):
            scaled = self.values[read] * decay
            if scaled < cutoff:
                self._slot.pop(int(self.indices[read]), None)
                continue
            self.indices[write] = self.indices[read]
            self.values[write] = scaled
            self._slot[int(self.indices[write])] = write
            write += 1
        self.size = write


class TrueOnlineTdLambda:
    """Linear TD(lambda), dutch traces, sparse features.

    True-online rather than conventional because the correction costs
    one scalar per step and makes the online sequence of weights
    exactly match what the offline lambda-return update would produce.
    The conventional form only matches in the limit of a vanishing step
    size, which is not the regime anyone runs in.
    """

    def __init__(self, config: LinearConfig, trace_capacity: int = 4096) -> None:
        if not 0.0 <= config.lam <= 1.0 or not 0.0 <= config.gamma <= 1.0:
            raise ValueError("gamma and lambda must lie in [0, 1]")

        self.config = config
        self.weights = np.zeros(config.n_features, dtype=np.float64)
        self.trace = SparseTrace(trace_capacity)
        self._previous_value = 0.0

    def value(self, features: IndexArray, activations: FloatArray) -> float:
        """One sparse dot product: sum of w[i] * x[i] over active i."""
        return float(np.dot(self.weights[features], activations))

    def begin_episode(self) -> None:
        self.trace.clear()
        self._previous_value = 0.0

    def step(
        self,
        features: IndexArray,
        activations: FloatArray,
        reward: float,
        next_features: IndexArray,
        next_activations: FloatArray,
        terminal: bool,
    ) -> float:
        """One transition. Everything is O(k) in the active set.

        The fused pass at the end does the weight update and the trace
        decay together, so no dense trace vector is ever written --
        which at a feature dimension in the millions is the difference
        between feasible and not.
        """
        alpha = self.config.alpha
        gamma = self.config.gamma

        current_value = self.value(features, activations)
        next_value = 0.0 if terminal else self.value(next_features, next_activations)
        td_error = reward + gamma * next_value - current_value

        # Dutch-trace factor: 1 - alpha * (x . e), computed from the
        # active set only.
        active_indices, active_values = self.trace.active()
        overlap = 0.0
        if active_values.size:
            # Align the feature activations onto the active positions.
            lookup = {int(index): position for position, index in enumerate(active_indices)}
            for position, feature in enumerate(features):
                slot = lookup.get(int(feature), -1)
                if slot >= 0:
                    overlap += activations[position] * active_values[slot]
        dutch_factor = 1.0 - alpha * overlap

        for position, feature in enumerate(features):
            self.trace.bump(int(feature), float(activations[position]), dutch_factor)

        active_indices, active_values = self.trace.active()

        # The fused update. np.add.at handles a feature appearing more
        # than once in the active set, which a plain fancy-index add
        # would silently collapse.
        np.add.at(self.weights, active_indices, alpha * td_error * active_values)

        # True-online correction: the term that makes this exact rather
        # than asymptotic. One scalar, applied to the current features.
        correction = alpha * (current_value - self._previous_value)
        np.add.at(self.weights, features, -correction * activations)

        self._previous_value = next_value
        self.trace.decay_and_prune(gamma * self.config.lam, self.config.trace_cutoff)
        return td_error


@dataclass
class LeastSquaresTd:
    """Batch TD as a single linear solve.

    The TD fixed point under linear approximation is defined by
    A w = b with

        A = sum_t x_t (x_t - gamma x_{t+1})^T,     b = sum_t r_t x_t

    Accumulating those and solving once uses every transition maximally,
    where incremental TD extracts one gradient step from each and
    discards it. On a fixed batch this reaches the same fixed point in
    a single pass rather than in thousands.

    The cost is why it is not the default: O(d^2) memory and an O(d^3)
    solve, which rules it out at exactly the feature dimensions that
    make the sparse incremental form attractive. It is the right tool
    for offline evaluation with a few thousand features and the wrong
    one for online learning with a few million.
    """

    n_features: int
    gamma: float
    regularization: float = 1e-6
    _a: FloatArray = None  # type: ignore[assignment]
    _b: FloatArray = None  # type: ignore[assignment]

    def __post_init__(self) -> None:
        self._a = np.eye(self.n_features, dtype=np.float64) * self.regularization
        self._b = np.zeros(self.n_features, dtype=np.float64)

    def accumulate(
        self, features: FloatArray, rewards: FloatArray, next_features: FloatArray
    ) -> None:
        """A whole batch of transitions in two matrix products.

        Dense here on purpose: this method is chosen when d is small
        enough for a d-by-d matrix, and at that size the dense form is
        faster than any sparse bookkeeping.
        """
        if features.shape != next_features.shape:
            raise ValueError("feature and successor matrices must have the same shape")
        if features.shape[0] != rewards.shape[0]:
            raise ValueError("feature rows and reward count differ")

        difference = features - self.gamma * next_features
        self._a += features.T @ difference
        self._b += features.T @ rewards

    def solve(self) -> FloatArray:
        """One solve rather than thousands of gradient steps.

        solve() rather than inv(): the A matrix here is not symmetric,
        it can be badly conditioned when features are correlated, and
        forming an explicit inverse makes both worse for no benefit.
        """
        try:
            return np.linalg.solve(self._a, self._b)
        except np.linalg.LinAlgError as error:
            raise ValueError(
                "the LSTD system is singular; correlated or redundant features make A "
                "rank-deficient, so either regularize more or drop the duplicates"
            ) from error


def trace_capacity_for(gamma: float, lam: float, cutoff: float, per_step: int) -> int:
    """How many features can be eligible at once, from the decay rate.

    Geometric decay bounds this: a component falls below the cutoff
    after log(cutoff)/log(gamma*lambda) steps, and at most per_step
    features activate each step. Sizing the buffer from that rather
    than guessing is what lets the hot path allocate nothing.
    """
    decay = gamma * lam
    if decay <= 0.0:
        return per_step
    if decay >= 1.0:
        raise ValueError("gamma*lambda >= 1 means traces never decay; no finite capacity exists")

    steps = int(np.ceil(np.log(cutoff) / np.log(decay)))
    return steps * per_step