Reinforcement learning from scratch

def discounted_return(rewards: list[float], gamma: float) -> float:
    """The return G = r_0 + gamma r_1 + gamma^2 r_2 + ...

    Raises:
        ValueError: If ``gamma`` is outside [0, 1].
    """
    if not 0.0 <= gamma <= 1.0:
        raise ValueError(f"gamma must be in [0, 1], got {gamma}")
    total = 0.0
    for reward in reversed(rewards):
        total = reward + gamma * total
    return total

@dataclass(frozen=True)
class GridWorld:
    """A small grid: reach the goal (+1), avoid the pit (-1), every step costs a bit.

    With probability ``slip`` the agent moves to one side of the intended direction
    instead (half to each side). Bumping into a wall or the edge leaves it in place.

    Attributes:
        rows: Grid height.
        cols: Grid width.
        start: Starting cell.
        goal: Terminal cell with reward +1.
        pit: Terminal cell with reward -1.
        walls: Cells that cannot be entered.
        step_reward: Reward for every non-terminal step.
        slip: Probability of moving sideways.
    """

    rows: int = 3
    cols: int = 4
    start: Cell = (2, 0)
    goal: Cell = (0, 3)
    pit: Cell = (1, 3)
    walls: frozenset[Cell] = field(default_factory=lambda: frozenset({(1, 1)}))
    step_reward: float = -0.04
    slip: float = 0.2

    def cells(self) -> list[Cell]:
        """Every cell the agent can be in, walls excluded."""
        return [
            (r, c)
            for r in range(self.rows)
            for c in range(self.cols)
            if (r, c) not in self.walls
        ]

    def is_terminal(self, cell: Cell) -> bool:
        """True for the goal and the pit."""
        return cell in (self.goal, self.pit)

    def _move(self, cell: Cell, action: int) -> Cell:
        dr, dc = ACTIONS[action]
        target = (cell[0] + dr, cell[1] + dc)
        inside = 0 <= target[0] < self.rows and 0 <= target[1] < self.cols
        return target if inside and target not in self.walls else cell

    def _reward(self, cell: Cell) -> float:
        if cell == self.goal:
            return 1.0
        if cell == self.pit:
            return -1.0
        return self.step_reward

    def transitions(self, cell: Cell, action: int) -> list[tuple[float, Cell, float]]:
        """Every outcome of ``action`` in ``cell`` as (probability, next, reward)."""
        outcomes = [(1.0 - self.slip, action)]
        outcomes += [
            (self.slip / 2, (action + 1) % 4),
            (self.slip / 2, (action + 3) % 4),
        ]
        result = []
        for probability, actual in outcomes:
            if probability > 0:
                after = self._move(cell, actual)
                result.append((probability, after, self._reward(after)))
        return result

    def step(self, cell: Cell, action: int, rng: random.Random) -> tuple[Cell, float]:
        """Samples one transition."""
        draw, total = rng.random(), 0.0
        outcomes = self.transitions(cell, action)
        for probability, after, reward in outcomes:
            total += probability
            if draw < total:
                return after, reward
        return outcomes[-1][1], outcomes[-1][2]

def action_values(
    world: GridWorld, values: dict[Cell, float], cell: Cell, gamma: float
) -> list[float]:
    """One-step lookahead: expected r + gamma V(s') for each action in ``cell``."""
    return [
        sum(
            p * (r + gamma * values[after])
            for p, after, r in world.transitions(cell, a)
        )
        for a in range(len(ACTIONS))
    ]


def bellman_sweep(
    world: GridWorld, values: dict[Cell, float], gamma: float
) -> dict[Cell, float]:
    """Applies the Bellman optimality equation once to every cell (from old values)."""
    updated = dict.fromkeys(world.cells(), 0.0)
    for cell in world.cells():
        if not world.is_terminal(cell):
            updated[cell] = max(action_values(world, values, cell, gamma))
    return updated


def value_iteration(
    world: GridWorld, gamma: float = 0.95, tolerance: float = 1e-8
) -> dict[Cell, float]:
    """Solves the Bellman optimality equation by repeating it until it stops changing.

    V(s) = max_a sum_s' P(s' | s, a) [r + gamma V(s')], with V = 0 at terminal cells.

    Returns:
        The optimal value of every cell.
    """
    values = dict.fromkeys(world.cells(), 0.0)
    while True:
        updated = bellman_sweep(world, values, gamma)
        change = max(abs(updated[c] - values[c]) for c in values)
        values = updated
        if change < tolerance:
            return values


def greedy_from_values(
    world: GridWorld, values: dict[Cell, float], gamma: float = 0.95
) -> dict[Cell, int]:
    """The action with the best one-step lookahead in every non-terminal cell."""
    policy = {}
    for cell in world.cells():
        if not world.is_terminal(cell):
            scores = action_values(world, values, cell, gamma)
            policy[cell] = max(range(len(ACTIONS)), key=scores.__getitem__)
    return policy

def policy_evaluation(
    world: GridWorld,
    policy: dict[Cell, list[float]],
    gamma: float = 0.95,
    tolerance: float = 1e-8,
) -> dict[Cell, float]:
    """V of a fixed policy: repeat V(s) = sum_a pi(a|s) [lookahead] until stable.

    Args:
        world: The grid world.
        policy: Action probabilities for every non-terminal cell.
        gamma: Discount factor.
        tolerance: Stop when no value changes by more than this.
    """
    values = dict.fromkeys(world.cells(), 0.0)
    while True:
        updated = dict.fromkeys(world.cells(), 0.0)
        for cell in policy:
            lookahead = action_values(world, values, cell, gamma)
            updated[cell] = sum(
                p * v for p, v in zip(policy[cell], lookahead, strict=True)
            )
        change = max(abs(updated[c] - values[c]) for c in values)
        values = updated
        if change < tolerance:
            return values


def policy_iteration(
    world: GridWorld, gamma: float = 0.95
) -> list[tuple[dict[Cell, int], dict[Cell, float]]]:
    """Evaluate, act greedily on the values, repeat until the policy stops changing.

    Starts from the uniformly random policy.

    Returns:
        The greedy policy and the values it was derived from, for every round.
    """
    uniform = [1 / len(ACTIONS)] * len(ACTIONS)
    probs = {c: uniform for c in world.cells() if not world.is_terminal(c)}
    history: list[tuple[dict[Cell, int], dict[Cell, float]]] = []
    while True:
        values = policy_evaluation(world, probs, gamma)
        greedy = greedy_from_values(world, values, gamma)
        if history and greedy == history[-1][0]:
            return history
        history.append((greedy, values))
        probs = {
            c: [1.0 if a == greedy[c] else 0.0 for a in range(len(ACTIONS))]
            for c in greedy
        }

def prediction_errors(
    world: GridWorld,
    policy: dict[Cell, int],
    *,
    method: str,
    episodes: int = 300,
    alpha: float = 0.05,
    gamma: float = 0.95,
    seed: int = 0,
) -> list[float]:
    """Learns V of a fixed policy from episodes; returns RMS error after each episode.

    Each episode starts in a random cell, so every state keeps being visited.
    ``"mc"`` (Monte Carlo) waits for the episode to end and moves V(s_t) toward the
    actual return G_t. ``"td"`` (TD(0)) moves V(s_t) toward r + gamma V(s_t+1) right
    after each step, using its own current estimate of the next state (bootstrapping).

    Raises:
        ValueError: If ``method`` is not ``"mc"`` or ``"td"``.
    """
    if method not in ("mc", "td"):
        raise ValueError(f"method must be 'mc' or 'td', got {method!r}")
    probs = {c: [1.0 if a == policy[c] else 0.0 for a in range(4)] for c in policy}
    truth = policy_evaluation(world, probs, gamma)
    rng = random.Random(seed)
    values = dict.fromkeys(world.cells(), 0.0)
    live = [c for c in world.cells() if not world.is_terminal(c)]
    errors = []
    for _ in range(episodes):
        cell, visited = rng.choice(live), []
        for _ in range(200):
            after, reward = world.step(cell, policy[cell], rng)
            if method == "td":
                values[cell] += alpha * (reward + gamma * values[after] - values[cell])
            visited.append((cell, reward))
            cell = after
            if world.is_terminal(cell):
                break
        if method == "mc":
            ret = 0.0
            for state, reward in reversed(visited):
                ret = reward + gamma * ret
                values[state] += alpha * (ret - values[state])
        squared = [(values[c] - truth[c]) ** 2 for c in live]
        errors.append((sum(squared) / len(squared)) ** 0.5)
    return errors

def _epsilon_greedy(values: list[float], epsilon: float, rng: random.Random) -> int:
    if rng.random() < epsilon:
        return rng.randrange(len(values))
    return max(range(len(values)), key=values.__getitem__)


def q_learning(
    world: Environment,
    *,
    episodes: int = 3000,
    alpha: float = 0.1,
    epsilon: float = 0.1,
    gamma: float = 0.95,
    seed: int = 0,
    trace: list[tuple[Cell, int, float, Cell, float, float, float]] | None = None,
) -> tuple[dict[Cell, list[float]], list[float]]:
    """Learns action values from experience, without knowing the transition rules.

    After each step: Q(s, a) += alpha * (r + gamma * max_a' Q(s', a') - Q(s, a)).
    Actions are epsilon-greedy: random with probability ``epsilon``, else the best.

    Args:
        world: Any environment with ``start``, ``cells``, ``is_terminal`` and ``step``.
        episodes: Number of episodes (each capped at 200 steps).
        alpha: Step size.
        epsilon: Exploration probability.
        gamma: Discount factor.
        seed: Seed for exploration and slips.
        trace: If given, every update is appended as (state, action, reward, next
            state, old Q value, TD target, new Q value).

    Returns:
        The Q table and the undiscounted return of every episode.
    """
    rng = random.Random(seed)
    q = {cell: [0.0] * len(ACTIONS) for cell in world.cells()}
    returns = []
    for _ in range(episodes):
        cell, total = world.start, 0.0
        for _ in range(200):
            action = _epsilon_greedy(q[cell], epsilon, rng)
            after, reward = world.step(cell, action, rng)
            target = reward
            if not world.is_terminal(after):
                target += gamma * max(q[after])
            old = q[cell][action]
            q[cell][action] = old + alpha * (target - old)
            if trace is not None:
                trace.append(
                    (cell, action, reward, after, old, target, q[cell][action])
                )
            cell, total = after, total + reward
            if world.is_terminal(cell):
                break
        returns.append(total)
    return q, returns

def sarsa(
    world: Environment,
    *,
    episodes: int = 3000,
    alpha: float = 0.1,
    epsilon: float = 0.1,
    gamma: float = 1.0,
    seed: int = 0,
) -> tuple[dict[Cell, list[float]], list[float]]:
    """On-policy TD control: the target uses the action the agent will actually take.

    Q(s, a) += alpha * (r + gamma * Q(s', a') - Q(s, a)), where a' is chosen
    epsilon-greedily, exactly as the agent will act. Q-learning uses max_a' instead,
    which learns about the greedy policy while behaving epsilon-greedily.

    Returns:
        The Q table and the undiscounted return of every episode.
    """
    rng = random.Random(seed)
    q = {cell: [0.0] * len(ACTIONS) for cell in world.cells()}
    returns = []
    for _ in range(episodes):
        cell, total = world.start, 0.0
        action = _epsilon_greedy(q[cell], epsilon, rng)
        for _ in range(200):
            after, reward = world.step(cell, action, rng)
            next_action = _epsilon_greedy(q[after], epsilon, rng)
            target = reward
            if not world.is_terminal(after):
                target += gamma * q[after][next_action]
            q[cell][action] += alpha * (target - q[cell][action])
            cell, action, total = after, next_action, total + reward
            if world.is_terminal(cell):
                break
        returns.append(total)
    return q, returns


def greedy_path(
    world: Environment, q: dict[Cell, list[float]], limit: int = 60
) -> list[Cell]:
    """Cells visited by following the greedy action from the start (no exploration)."""
    rng = random.Random(0)
    cell, path = world.start, [world.start]
    for _ in range(limit):
        cell, _ = world.step(
            cell, max(range(len(ACTIONS)), key=q[cell].__getitem__), rng
        )
        path.append(cell)
        if world.is_terminal(cell):
            break
    return path

@dataclass(frozen=True)
class CliffWalk:
    """The classic cliff: walk along the bottom edge without falling off.

    A 4 x 12 grid. The start is the bottom-left cell and the goal the bottom-right; the
    cells between them on the bottom row are a cliff. Every step costs -1; stepping
    onto the cliff costs -100 and sends the agent back to the start. Moves are
    deterministic.
    """

    rows: int = 4
    cols: int = 12
    start: Cell = (3, 0)
    goal: Cell = (3, 11)

    def cells(self) -> list[Cell]:
        """Every cell, cliff included (the agent never rests on it)."""
        return [(r, c) for r in range(self.rows) for c in range(self.cols)]

    def is_cliff(self, cell: Cell) -> bool:
        """True for the bottom-row cells between start and goal."""
        return cell[0] == self.rows - 1 and 0 < cell[1] < self.cols - 1

    def is_terminal(self, cell: Cell) -> bool:
        """Only the goal ends an episode."""
        return cell == self.goal

    def step(self, cell: Cell, action: int, rng: random.Random) -> tuple[Cell, float]:
        """Moves one cell (``rng`` is unused: the cliff is deterministic)."""
        del rng
        dr, dc = ACTIONS[action]
        after = (
            min(max(cell[0] + dr, 0), self.rows - 1),
            min(max(cell[1] + dc, 0), self.cols - 1),
        )
        if self.is_cliff(after):
            return self.start, -100.0
        return after, -1.0

def epsilon_greedy_bandit(
    means: list[float], epsilon: float, steps: int = 1000, seed: int = 0
) -> list[bool]:
    """Plays a Gaussian-reward bandit; returns whether each pull chose the best arm.

    The agent keeps a running average reward per arm. With probability ``epsilon`` it
    explores (a random arm); otherwise it exploits (the arm with the best average).
    """
    rng = random.Random(seed)
    counts = [0] * len(means)
    estimates = [0.0] * len(means)
    best = max(range(len(means)), key=means.__getitem__)
    chose_best = []
    for _ in range(steps):
        if rng.random() < epsilon:
            arm = rng.randrange(len(means))
        else:
            arm = max(range(len(means)), key=estimates.__getitem__)
        reward = rng.gauss(means[arm], 1.0)
        counts[arm] += 1
        estimates[arm] += (reward - estimates[arm]) / counts[arm]
        chose_best.append(arm == best)
    return chose_best

def ucb_bandit(
    means: list[float], c: float = 1.0, steps: int = 1000, seed: int = 0
) -> list[bool]:
    """Upper confidence bound: pick the arm with the best optimistic estimate.

    Score = average reward + c * sqrt(ln t / pulls). Rarely pulled arms get a large
    bonus, so exploration is directed at arms that might be better, not random.
    """
    rng = random.Random(seed)
    counts = [0] * len(means)
    estimates = [0.0] * len(means)
    best = max(range(len(means)), key=means.__getitem__)
    chose_best = []
    for t in range(1, steps + 1):
        if t <= len(means):
            arm = t - 1
        else:
            scores = [
                estimates[a] + c * math.sqrt(math.log(t) / counts[a])
                for a in range(len(means))
            ]
            arm = max(range(len(means)), key=scores.__getitem__)
        reward = rng.gauss(means[arm], 1.0)
        counts[arm] += 1
        estimates[arm] += (reward - estimates[arm]) / counts[arm]
        chose_best.append(arm == best)
    return chose_best

def dqn(
    world: GridWorld,
    *,
    episodes: int = 300,
    batch: int = 64,
    gamma: float = 0.95,
    sync_every: int = 100,
    seed: int = 0,
) -> torch.nn.Module:
    """Deep Q-learning on the grid world: a network instead of a table.

    The three ingredients that make Q-learning work with a neural network:

    * a Q-network mapping a state (one-hot cell here; pixels in DQN) to a Q per action;
    * a replay buffer, so each update uses a random batch of old transitions instead of
      the last few, highly correlated ones;
    * a target network, a periodically refreshed copy used for the TD target, so the
      target does not move with every gradient step.

    Exploration decays from fully random to 5 percent over the first half of training.

    Returns:
        The trained Q-network.
    """
    torch.manual_seed(seed)
    rng = random.Random(seed)
    cells = world.cells()
    index = {cell: i for i, cell in enumerate(cells)}
    eye = torch.eye(len(cells))
    net = torch.nn.Sequential(
        torch.nn.Linear(len(cells), 64), torch.nn.ReLU(), torch.nn.Linear(64, 4)
    )
    target_net = copy.deepcopy(net)
    optimizer = torch.optim.Adam(net.parameters(), lr=1e-3)
    buffer: deque[tuple[int, int, float, int, float]] = deque(maxlen=5000)
    updates = 0
    for episode in range(episodes):
        epsilon = max(0.05, 1.0 - episode / (episodes / 2))
        cell = world.start
        for _ in range(100):
            with torch.no_grad():
                q_values = net(eye[index[cell]]).tolist()
            action = _epsilon_greedy(q_values, epsilon, rng)
            after, reward = world.step(cell, action, rng)
            done = float(world.is_terminal(after))
            buffer.append((index[cell], action, reward, index[after], done))
            cell = after
            if len(buffer) >= batch:
                sample = rng.sample(list(buffer), batch)
                s, a, r, s2, d = (torch.tensor(column) for column in zip(*sample))
                with torch.no_grad():
                    best_next = target_net(eye[s2]).max(dim=1).values
                    target = r.float() + gamma * (1 - d.float()) * best_next
                predicted = net(eye[s])[torch.arange(batch), a]
                loss = torch.nn.functional.mse_loss(predicted, target)
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
                updates += 1
                if updates % sync_every == 0:
                    target_net.load_state_dict(net.state_dict())
            if done:
                break
    return net

def policy_gradient_samples(
    means: tuple[float, float], baseline: float, samples: int = 5000, seed: int = 0
) -> torch.Tensor:
    """One-sample REINFORCE estimates of d(expected reward)/d(logit of arm 1).

    The policy is a fair coin over two arms (logits 0, 0). Each estimate is
    (reward - baseline) * d log pi(arm) / d logit_1, which is +-0.5 here.

    Returns:
        Shape (samples,): the individual gradient estimates.
    """
    generator = torch.Generator().manual_seed(seed)
    arms = torch.randint(0, 2, (samples,), generator=generator)
    rewards = torch.tensor(means)[arms] + torch.randn(samples, generator=generator)
    score = torch.where(arms == 1, 0.5, -0.5)
    return (rewards - baseline) * score

def reinforce(
    world: GridWorld,
    *,
    episodes: int = 600,
    lr: float = 0.1,
    baseline: bool = True,
    gamma: float = 0.95,
    seed: int = 0,
) -> list[float]:
    """REINFORCE on the grid world with a table of action logits per cell.

    After each episode, every visited (s_t, a_t) is pushed by (G_t - b(s_t)) times the
    gradient of log pi(a_t | s_t). With ``baseline``, b is a learned value table moved
    toward the observed returns (a minimal actor-critic); without it, b = 0.

    Returns:
        The undiscounted return of every episode.
    """
    torch.manual_seed(seed)
    rng = random.Random(seed)
    index = {cell: i for i, cell in enumerate(world.cells())}
    logits = torch.zeros(len(index), len(ACTIONS), requires_grad=True)
    optimizer = torch.optim.Adam([logits], lr=lr)
    values = torch.zeros(len(index))
    returns = []
    for _ in range(episodes):
        cell, steps = world.start, []
        for _ in range(100):
            probs = torch.softmax(logits[index[cell]].detach(), dim=0)
            action = int(torch.multinomial(probs, 1))
            after, reward = world.step(cell, action, rng)
            steps.append((index[cell], action, reward))
            cell = after
            if world.is_terminal(cell):
                break
        rows = torch.tensor([s for s, _, _ in steps])
        actions = torch.tensor([a for _, a, _ in steps])
        future, ret = [], 0.0
        for _, _, reward in reversed(steps):
            ret = reward + gamma * ret
            future.append(ret)
        future_t = torch.tensor(future[::-1])
        advantage = future_t - values[rows] if baseline else future_t
        logp = torch.log_softmax(logits[rows], dim=1)[torch.arange(len(steps)), actions]
        loss = -(advantage * logp).sum()
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if baseline:
            values.index_add_(0, rows, 0.1 * (future_t - values[rows]))
        returns.append(sum(r for _, _, r in steps))
    return returns