Imitation Learning (Behavioral Cloning, DAgger)

"""Behavioural cloning and DAgger - the training-distribution swap, transcribed.

Note what the loss is in BOTH functions: ordinary cross-entropy over
(state, expert_action) pairs. No reward, no bootstrapping, no environment
model. What DAgger changes is never the loss - it is which states end up
labelled in the dataset the loss is computed over.
"""

import math


def softmax(scores):
    # Numerically stabilized: subtract the max before exponentiating.
    m = max(scores)
    exps = [math.exp(s - m) for s in scores]
    total = sum(exps)
    return [e / total for e in exps]


def predict_action(weights, state, n_actions):
    scores = [sum(weights[a][i] * state[i] for i in range(len(state))) for a in range(n_actions)]
    probs = softmax(scores)
    return max(range(n_actions), key=lambda a: probs[a])


def fit_bc(dataset, n_actions, n_features, epochs=200, lr=0.1):
    """Behavioural cloning: fit a linear-softmax policy to demonstrated
    actions. This is the entire objective - no reward appears anywhere.
    """
    weights = [[0.0] * n_features for _ in range(n_actions)]

    for _ in range(epochs):
        for state, expert_action in dataset:
            scores = [sum(weights[a][i] * state[i] for i in range(n_features)) for a in range(n_actions)]
            probs = softmax(scores)

            # Gradient of cross-entropy w.r.t. the scores is (prob - one_hot).
            for a in range(n_actions):
                target = 1.0 if a == expert_action else 0.0
                grad_score = probs[a] - target
                for i in range(n_features):
                    weights[a][i] -= lr * grad_score * state[i]

    return weights


def dagger(env, expert, n_actions, n_features, initial_dataset, rounds=10, epochs=200, lr=0.1):
    """Dataset aggregation: run the current policy, label the states it visits
    with the expert, and refit on the union. The dataset only grows, and it is
    refit from scratch every round.
    """
    dataset = list(initial_dataset)
    weights = fit_bc(dataset, n_actions, n_features, epochs, lr)

    for _ in range(rounds):
        state = env.reset()
        done = False

        while not done:
            action = predict_action(weights, state, n_actions)
            # Query the expert on a state the LEARNER chose to visit - the
            # entire mechanism that closes the distribution gap.
            expert_action = expert(state)
            dataset.append((state, expert_action))

            state, done = env.step(action)

        weights = fit_bc(dataset, n_actions, n_features, epochs, lr)

    return weights

"""Behavioural cloning / DAgger - typed, vectorized per sample, loss tracked."""

from dataclasses import dataclass, field
from typing import Protocol

import numpy as np
from numpy.typing import NDArray

State = NDArray[np.float64]
Weights = NDArray[np.float64]


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


class Expert(Protocol):
    def act(self, state: State) -> int: ...


@dataclass
class BcConfig:
    n_actions: int
    n_features: int
    epochs: int = 200
    learning_rate: float = 0.1

    def __post_init__(self) -> None:
        if self.n_actions < 1 or self.n_features < 1:
            raise ValueError("action and feature spaces must be non-empty")
        if not self.learning_rate > 0.0:
            raise ValueError(f"learning_rate must be positive, got {self.learning_rate}")


@dataclass
class TrainingResult:
    weights: Weights
    epoch_loss: list[float] = field(default_factory=list)

    @property
    def has_converged(self) -> bool:
        """Falling fitting loss says the optimization worked. It says nothing
        about whether the distribution it was fit on matches deployment -
        that gap is the entire subject of this model."""
        if len(self.epoch_loss) < 10:
            return False
        recent = float(np.mean(self.epoch_loss[-5:]))
        earlier = float(np.mean(self.epoch_loss[-10:-5]))
        return recent < earlier * 1.02


def softmax(scores: NDArray[np.float64]) -> NDArray[np.float64]:
    shifted = scores - scores.max(axis=-1, keepdims=True)
    exps = np.exp(shifted)
    return exps / exps.sum(axis=-1, keepdims=True)


def fit_bc(dataset: list[tuple[State, int]], config: BcConfig) -> TrainingResult:
    """Ordinary cross-entropy over (state, expert_action) pairs - the entire
    objective. No reward, no environment interaction, no bootstrapping."""
    weights = np.zeros((config.n_actions, config.n_features), dtype=np.float64)
    result = TrainingResult(weights=weights)

    for _ in range(config.epochs):
        losses: list[float] = []
        for state, expert_action in dataset:
            probs = softmax(weights @ state)
            losses.append(float(-np.log(max(probs[expert_action], 1e-12))))

            target = np.zeros(config.n_actions)
            target[expert_action] = 1.0
            weights -= config.learning_rate * np.outer(probs - target, state)

        result.epoch_loss.append(float(np.mean(losses)))

    return result


def dagger(
    env: Environment,
    expert: Expert,
    config: BcConfig,
    initial_dataset: list[tuple[State, int]],
    rounds: int = 10,
) -> TrainingResult:
    """Dataset aggregation: label the states the current policy visits with
    the expert, then refit on the union. Refitting from scratch every round is
    what makes the total cost grow with the round count squared."""
    dataset = list(initial_dataset)
    result = fit_bc(dataset, config)

    for _ in range(rounds):
        state = env.reset()
        done = False

        while not done:
            action = int(np.argmax(result.weights @ state))
            # Query the expert on a state the LEARNER chose to visit - the
            # mechanism that closes the distribution gap.
            dataset.append((state, expert.act(state)))
            state, done = env.step(action)

        result = fit_bc(dataset, config)

    return result

"""Behavioural cloning - full-batch vectorized gradient, no per-sample loop."""

import numpy as np
from numpy.typing import NDArray


def softmax(scores: NDArray[np.float64]) -> NDArray[np.float64]:
    shifted = scores - scores.max(axis=-1, keepdims=True)
    exps = np.exp(shifted)
    return exps / exps.sum(axis=-1, keepdims=True)


def train_epoch(
    weights: NDArray[np.float64],
    states: NDArray[np.float64],
    expert_actions: NDArray[np.intp],
    learning_rate: float,
) -> float:
    """One full-batch update over the whole dataset at once.

    The previous stage recomputes scores per sample inside a Python loop; here
    the whole dataset is two matrix multiplies - one for the scores, one for
    the gradient - and the per-sample loop is gone entirely.
    """
    n_samples = states.shape[0]

    probs = softmax(states @ weights.T)                # (n_samples, n_actions)

    targets = np.zeros_like(probs)
    targets[np.arange(n_samples), expert_actions] = 1.0

    residual = probs - targets                         # (n_samples, n_actions)
    gradient = residual.T @ states / n_samples          # (n_actions, n_features)
    weights -= learning_rate * gradient

    sample_losses = -np.log(np.clip(probs[np.arange(n_samples), expert_actions], 1e-12, None))
    return float(sample_losses.mean())


def collect_dagger_round_batched(weights, rollout_states, query_expert_batch):
    """Query the expert on an entire batch of visited states in one call,
    instead of one state at a time - the batched analogue of full-batch
    fitting, and it matters whenever the expert call carries fixed overhead
    such as a model server or a human interface.

    Also reports how often the policy already agrees with the expert, which
    is a cheap convergence signal that needs no additional environment steps.
    """
    learner_actions = np.argmax(rollout_states @ weights.T, axis=1)
    expert_actions = query_expert_batch(rollout_states)
    agreement = float(np.mean(learner_actions == expert_actions))
    return rollout_states, expert_actions, agreement