Deep Q-Network (DQN)

"""DQN - the whole method transcribed, network included.

No autodiff and no tensors: a two-layer MLP with hand-written backprop, so the
semi-gradient is visible as code. Watch for where the gradient is NOT taken -
through the target - because that omission is the entire reason none of the
usual convergence results apply here.
"""

import math
import random


def make_network(n_inputs, n_hidden, n_actions):
    """Weights as lists of lists. Small random init, never zeros: identical
    hidden units receive identical gradients and stay identical forever."""
    scale = 1.0 / math.sqrt(n_inputs)
    return {
        "w1": [[random.uniform(-scale, scale) for _ in range(n_inputs)] for _ in range(n_hidden)],
        "b1": [0.0] * n_hidden,
        "w2": [[random.uniform(-scale, scale) for _ in range(n_hidden)] for _ in range(n_actions)],
        "b2": [0.0] * n_actions,
    }


def copy_network(net):
    """The target network is a frozen snapshot of the online one. Freezing it
    is what stops the target moving every time the estimate does."""
    return {
        "w1": [row[:] for row in net["w1"]],
        "b1": net["b1"][:],
        "w2": [row[:] for row in net["w2"]],
        "b2": net["b2"][:],
    }


def forward(net, state):
    """Returns (q_values, hidden). The hidden activations are kept because the
    backward pass needs them and recomputing them would be the same work."""
    hidden = []
    for row, bias in zip(net["w1"], net["b1"]):
        total = bias
        for weight, feature in zip(row, state):
            total += weight * feature
        hidden.append(total if total > 0.0 else 0.0)        # ReLU

    q_values = []
    for row, bias in zip(net["w2"], net["b2"]):
        total = bias
        for weight, activation in zip(row, hidden):
            total += weight * activation
        q_values.append(total)

    return q_values, hidden


def backward(net, state, hidden, action, td_error, lr):
    """One SGD step on the squared TD error, for the taken action only.

    The derivative of (y - Q(s,a))^2 with respect to theta is
    -2 (y - Q(s,a)) dQ(s,a)/dtheta. But y depends on theta too, through the
    target network, and that dependence is simply not differentiated. The
    discarded term is what makes this a SEMI-gradient method: the direction
    being descended is not the gradient of this loss or of any other function.
    """
    # Only the taken action's output unit appears in the loss at all. The
    # other |A| - 1 outputs get no gradient, which is why one transition
    # teaches the network about exactly one action.
    grad_output = -td_error

    grad_hidden = [0.0] * len(hidden)
    for index, weight in enumerate(net["w2"][action]):
        grad_hidden[index] = grad_output * weight
        net["w2"][action][index] -= lr * grad_output * hidden[index]
    net["b2"][action] -= lr * grad_output

    for unit in range(len(hidden)):
        if hidden[unit] <= 0.0:
            continue                                        # ReLU is closed here
        delta = grad_hidden[unit]
        for index, feature in enumerate(state):
            net["w1"][unit][index] -= lr * delta * feature
        net["b1"][unit] -= lr * delta


def train(env, n_inputs, n_actions, steps=200_000, n_hidden=64, lr=1e-3,
          gamma=0.99, batch_size=32, buffer_size=50_000, target_period=1_000,
          warmup=1_000, epsilon_start=1.0, epsilon_end=0.05, epsilon_steps=50_000):
    online = make_network(n_inputs, n_hidden, n_actions)
    target = copy_network(online)

    buffer = []                  # (state, action, reward, next_state, done)
    state = env.reset()

    for step in range(steps):
        fraction = min(1.0, step / epsilon_steps)
        epsilon = epsilon_start + fraction * (epsilon_end - epsilon_start)

        if random.random() < epsilon:
            action = random.randrange(n_actions)
        else:
            q_values, _ = forward(online, state)
            action = max(range(n_actions), key=lambda a: q_values[a])

        next_state, reward, done = env.step(action)

        buffer.append((state, action, reward, next_state, done))
        if len(buffer) > buffer_size:
            buffer.pop(0)                                   # a ring buffer, by hand

        state = env.reset() if done else next_state

        # Training on a nearly empty buffer is training on one correlated
        # episode, which is precisely what replay exists to prevent.
        if len(buffer) < warmup:
            continue

        for _ in range(batch_size):
            # Uniform sampling is the decorrelation mechanism. Consecutive
            # transitions are almost the same sample; a network trained on
            # that stream fits the last few seconds and forgets the rest.
            s, a, r, s_next, terminal = random.choice(buffer)

            # The bootstrap comes from the FROZEN copy. Taking it from the
            # online network closes a feedback loop between the estimate and
            # its own target, which is the configuration that diverges.
            if terminal:
                bootstrap = 0.0                             # no future past a terminal
            else:
                next_q, _ = forward(target, s_next)
                bootstrap = max(next_q)

            y = r + gamma * bootstrap
            q_values, hidden = forward(online, s)
            td_error = y - q_values[a]

            backward(online, s, hidden, a, td_error, lr)

        if step % target_period == 0:
            # Refresh the frozen copy. Too often and the feedback loop comes
            # back; too rarely and value information stops propagating at all.
            target = copy_network(online)

    return online

"""DQN - typed, Double-Q targets, Huber loss, truncation handled correctly."""

from __future__ import annotations

from dataclasses import dataclass
from typing import NamedTuple, Protocol

import numpy as np
import torch
from torch import Tensor, nn


class Transition(NamedTuple):
    state: np.ndarray
    action: int
    reward: float
    next_state: np.ndarray
    # Terminal means the MDP ended and there is no future. Truncated means the
    # episode hit a wall-clock or step limit and the future is still there.
    # Collapsing the two teaches the agent the world ends at the time limit,
    # and the resulting policy is visibly short-sighted.
    terminal: bool


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


@dataclass(frozen=True)
class DqnConfig:
    n_inputs: int
    n_actions: int
    hidden: tuple[int, ...] = (128, 128)
    learning_rate: float = 1e-4
    gamma: float = 0.99
    batch_size: int = 32
    buffer_size: int = 100_000
    warmup_steps: int = 5_000
    target_update_period: int = 2_000
    train_frequency: int = 4
    grad_clip_norm: float = 10.0
    huber_delta: float = 1.0

    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.warmup_steps < self.batch_size:
            raise ValueError("warmup must exceed the batch size, or the first batch repeats samples")
        if self.n_actions < 2:
            raise ValueError("a single-action problem is not a control problem")


class ReplayBuffer:
    """A fixed-capacity ring over pre-allocated arrays.

    A list of Python tuples reallocates, fragments, and costs a gather per
    sample. One contiguous array per field costs one fancy-index per batch.
    """

    def __init__(self, capacity: int, n_inputs: int) -> None:
        if capacity < 1:
            raise ValueError("capacity must be positive")
        self._states = np.zeros((capacity, n_inputs), dtype=np.float32)
        self._next_states = np.zeros((capacity, n_inputs), dtype=np.float32)
        self._actions = np.zeros(capacity, dtype=np.int64)
        self._rewards = np.zeros(capacity, dtype=np.float32)
        self._terminals = np.zeros(capacity, dtype=np.bool_)
        self._capacity = capacity
        self._cursor = 0
        self._size = 0

    def __len__(self) -> int:
        return self._size

    def add(self, transition: Transition) -> None:
        index = self._cursor
        self._states[index] = transition.state
        self._next_states[index] = transition.next_state
        self._actions[index] = transition.action
        self._rewards[index] = transition.reward
        self._terminals[index] = transition.terminal
        self._cursor = (self._cursor + 1) % self._capacity
        self._size = min(self._size + 1, self._capacity)

    def sample(self, batch_size: int, rng: np.random.Generator) -> tuple[Tensor, ...]:
        if self._size < batch_size:
            raise ValueError(f"buffer holds {self._size}, cannot sample {batch_size}")
        index = rng.integers(0, self._size, size=batch_size)
        return (
            torch.from_numpy(self._states[index]),
            torch.from_numpy(self._actions[index]),
            torch.from_numpy(self._rewards[index]),
            torch.from_numpy(self._next_states[index]),
            torch.from_numpy(self._terminals[index]),
        )


def build_network(config: DqnConfig) -> nn.Module:
    layers: list[nn.Module] = []
    width = config.n_inputs
    for size in config.hidden:
        layers += [nn.Linear(width, size), nn.ReLU()]
        width = size
    layers.append(nn.Linear(width, config.n_actions))
    return nn.Sequential(*layers)


def double_dqn_target(
    online: nn.Module,
    target: nn.Module,
    rewards: Tensor,
    next_states: Tensor,
    terminals: Tensor,
    gamma: float,
) -> Tensor:
    """Online network SELECTS the action, target network VALUES it.

    A plain max over the target network's own outputs picks whichever action
    the noise happened to favour, and that bias compounds under function
    approximation into value estimates that climb without bound. Splitting
    selection from evaluation makes the two sources of noise independent, so
    the max can no longer systematically select for its own error.
    """
    with torch.no_grad():
        best_actions = online(next_states).argmax(dim=1, keepdim=True)
        next_values = target(next_states).gather(1, best_actions).squeeze(1)
        # Terminal transitions have no future. Truncated ones DO, which is why
        # the flag stored is terminal and not done.
        next_values = next_values.masked_fill(terminals, 0.0)
        return rewards + gamma * next_values


def train_step(
    online: nn.Module,
    target: nn.Module,
    optimizer: torch.optim.Optimizer,
    batch: tuple[Tensor, ...],
    config: DqnConfig,
) -> float:
    states, actions, rewards, next_states, terminals = batch

    y = double_dqn_target(online, target, rewards, next_states, terminals, config.gamma)
    predicted = online(states).gather(1, actions.unsqueeze(1)).squeeze(1)

    # Huber rather than MSE: one transition with an outsized TD error would
    # otherwise produce a gradient large enough to destroy the network in a
    # single step, and there is no recovering from that.
    loss = nn.functional.smooth_l1_loss(predicted, y, beta=config.huber_delta)

    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    # A second bound on step size, for the case where a whole batch is bad.
    nn.utils.clip_grad_norm_(online.parameters(), config.grad_clip_norm)
    optimizer.step()

    return float(loss.item())

"""DQN - uint8 frame storage, index-based stacking, vectorized envs.

The two costs that actually bind are replay memory and simulator throughput.
Neither is the network, which is why almost nothing here is about the network.
"""

from __future__ import annotations

import numpy as np
import torch
from torch import Tensor

STACK = 4


class FrameRingBuffer:
    """A million transitions of stacked 84x84 frames, in about 7 GB.

    Two decisions do all the work. Frames are stored as uint8 and converted at
    sample time, which is a 4x saving over float32 for free - the network wants
    floats, the buffer does not. And the stack is assembled by INDEX rather
    than stored: consecutive transitions share three of their four frames, so
    materializing the stack per transition would store every frame four times
    and turn 7 GB into 28. That single distinction decides whether the buffer
    fits in RAM, and it is the usual reason one does not.
    """

    def __init__(self, capacity: int, height: int, width: int) -> None:
        # One flat allocation per field, single dtype, written in place.
        self._frames = np.zeros((capacity, height, width), dtype=np.uint8)
        self._actions = np.zeros(capacity, dtype=np.int64)
        self._rewards = np.zeros(capacity, dtype=np.float32)
        self._terminals = np.zeros(capacity, dtype=np.bool_)
        self._capacity = capacity
        self._cursor = 0
        self._size = 0

    def add(self, frame: np.ndarray, action: int, reward: float, terminal: bool) -> None:
        index = self._cursor
        self._frames[index] = frame          # in place; no per-transition allocation
        self._actions[index] = action
        self._rewards[index] = reward
        self._terminals[index] = terminal
        self._cursor = (self._cursor + 1) % self._capacity
        self._size = min(self._size + 1, self._capacity)

    def _stack_indices(self, index: np.ndarray) -> np.ndarray:
        """Gather STACK consecutive frame indices per sample, wrapping the ring.

        One broadcast subtraction builds the whole (batch, STACK) index matrix,
        so no Python loop over samples appears anywhere in the hot path.
        """
        offsets = np.arange(STACK - 1, -1, -1, dtype=np.int64)
        return (index[:, None] - offsets[None, :]) % self._capacity

    def sample(self, batch_size: int, rng: np.random.Generator, device: torch.device):
        # Avoid the cursor: a window straddling it mixes the newest transitions
        # with the oldest ones and produces a state that never existed.
        valid_low = STACK
        index = rng.integers(valid_low, self._size - 1, size=batch_size)

        states = self._frames[self._stack_indices(index)]
        next_states = self._frames[self._stack_indices(index + 1)]

        return (
            # One transfer per field, one dtype conversion on the device rather
            # than |batch| conversions on the host.
            torch.from_numpy(states).to(device, non_blocking=True).float().div_(255.0),
            torch.from_numpy(self._actions[index]).to(device, non_blocking=True),
            torch.from_numpy(self._rewards[index]).to(device, non_blocking=True),
            torch.from_numpy(next_states).to(device, non_blocking=True).float().div_(255.0),
            torch.from_numpy(self._terminals[index]).to(device, non_blocking=True),
        )


def n_step_returns(rewards: np.ndarray, gamma: float, n: int) -> np.ndarray:
    """Accumulate n rewards before bootstrapping, as one reversed pass.

    Trades the bias of a bad value estimate for the variance of a longer
    sampled return, and it is the cheapest real improvement available to a
    DQN - a few lines for a change in learning speed that usually exceeds
    anything gained by tuning the network.
    """
    out = np.zeros_like(rewards)
    running = 0.0
    for position in range(len(rewards) - 1, -1, -1):
        running = rewards[position] + gamma * running
        out[position] = running
    return out


@torch.no_grad()
def act_vectorized(
    online: torch.nn.Module,
    states: Tensor,
    epsilon: float,
    n_actions: int,
    generator: torch.Generator,
) -> Tensor:
    """One forward pass covering every environment in the vector.

    Simulator throughput, not gradient computation, is what bounds wall-clock
    on most problems. Stepping N environments in lockstep amortizes both the
    interpreter overhead and the per-call launch cost of the forward pass over
    the whole batch, which is a far larger lever than a faster accelerator.
    """
    greedy = online(states).argmax(dim=1)
    explore = torch.rand(states.shape[0], generator=generator, device=states.device) < epsilon
    random_actions = torch.randint(
        n_actions, (states.shape[0],), generator=generator, device=states.device
    )
    return torch.where(explore, random_actions, greedy)