Model-Based RL & MPC

"""Model-based RL with random-shooting MPC, transcribed.

Two pieces. A dynamics model fitted by ordinary regression - the next state IS
the label, which is why every transition here carries a whole state vector of
information instead of a single scalar reward. And a planner that searches over
action sequences by rolling them out through that model.

Watch the last line of the control loop: only the FIRST action of the optimized
sequence is executed, and the rest is thrown away and recomputed. That looks
wasteful and is the entire mechanism - replanning from a real observation every
step means model error never gets a chance to accumulate.
"""

import random


def make_model(n_state, n_action):
    """Predicts the state DELTA: s' - s = A s + B a + c.

    Predicting the delta rather than the absolute next state matters more than
    it looks. Consecutive states are highly correlated, so a model predicting
    absolutes scores beautifully on held-out error while having learned almost
    nothing about how the system actually changes.
    """
    return {
        "a": [[0.0] * n_state for _ in range(n_state)],
        "b": [[0.0] * n_action for _ in range(n_state)],
        "c": [0.0] * n_state,
    }


def predict(model, state, action):
    next_state = []
    for row in range(len(model["c"])):
        delta = model["c"][row]
        for index, value in enumerate(state):
            delta += model["a"][row][index] * value
        for index, value in enumerate(action):
            delta += model["b"][row][index] * value
        next_state.append(state[row] + delta)
    return next_state


def fit_step(model, state, action, next_state, lr):
    """One SGD step on squared error over the delta.

    Ordinary supervised learning: the label was handed over by the environment
    for free, alongside the reward that a model-free method would have kept
    instead. That asymmetry is the whole sample-efficiency argument.
    """
    for row in range(len(model["c"])):
        target = next_state[row] - state[row]

        prediction = model["c"][row]
        for index, value in enumerate(state):
            prediction += model["a"][row][index] * value
        for index, value in enumerate(action):
            prediction += model["b"][row][index] * value

        error = prediction - target
        for index, value in enumerate(state):
            model["a"][row][index] -= lr * error * value
        for index, value in enumerate(action):
            model["b"][row][index] -= lr * error * value
        model["c"][row] -= lr * error


def rollout_return(model, reward_fn, state, sequence, gamma):
    """Score one candidate action sequence by imagining it.

    Every term past the first has the model predicting from its OWN previous
    prediction. Errors therefore compound: a per-step error that looks
    negligible is multiplied down the rollout, which is why the horizon is the
    most important knob here and why it should be set from measured multi-step
    error rather than guessed.
    """
    imagined = list(state)
    total = 0.0
    discount = 1.0

    for action in sequence:
        imagined = predict(model, imagined, action)
        total += discount * reward_fn(imagined, action)
        discount *= gamma

    return total


def plan_random_shooting(model, reward_fn, state, n_action, horizon, candidates,
                         action_low, action_high, gamma):
    """The simplest planner there is: sample sequences, keep the best.

    No gradients and no structure, which is exactly why it works when the model
    is non-differentiable. It is also the baseline the cross-entropy method
    improves on, by refitting a distribution to the best candidates instead of
    resampling uniformly every iteration.
    """
    best_sequence = None
    best_score = -float("inf")

    for _ in range(candidates):
        sequence = [
            [random.uniform(action_low, action_high) for _ in range(n_action)]
            for _ in range(horizon)
        ]
        score = rollout_return(model, reward_fn, state, sequence, gamma)
        if score > best_score:
            best_sequence, best_score = sequence, score

    return best_sequence


def run(env, reward_fn, n_state, n_action, steps=10_000, horizon=15,
        candidates=200, lr=1e-3, gamma=0.99, action_low=-1.0, action_high=1.0,
        warmup=200):
    model = make_model(n_state, n_action)
    state = env.reset()

    for step in range(steps):
        if step < warmup:
            # Planning against a model that knows nothing is searching for the
            # maxima of noise, and the planner is very good at finding them.
            action = [random.uniform(action_low, action_high) for _ in range(n_action)]
        else:
            sequence = plan_random_shooting(
                model, reward_fn, state, n_action, horizon, candidates,
                action_low, action_high, gamma,
            )
            action = sequence[0]        # only the first. Receding horizon.

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

        # Fit as we go. The planner immediately moves the system into states the
        # model has not seen, so a model fitted once is fitted to a policy that
        # no longer exists.
        fit_step(model, state, action, next_state, lr)

        state = env.reset() if done else next_state

    return model

"""Model-based RL - probabilistic ensemble, CEM planner, uncertainty penalty."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Callable

import torch
from torch import Tensor, nn

LOG_VAR_BOUNDS = (-10.0, 0.5)


@dataclass(frozen=True)
class PlannerConfig:
    n_state: int
    n_action: int
    action_low: float = -1.0
    action_high: float = 1.0
    # The decisive knob. Set it from MEASURED multi-step rollout error, not from
    # intuition: compounding is what makes a horizon that worked at ten steps
    # useless at thirty.
    horizon: int = 15
    candidates: int = 500
    elite_fraction: float = 0.1
    iterations: int = 5
    gamma: float = 0.99
    # How much predicted return is discounted where the ensemble disagrees.
    # This is the explicit defence against the planner exploiting model error.
    uncertainty_penalty: float = 1.0

    def __post_init__(self) -> None:
        if self.horizon < 1:
            raise ValueError("horizon must be at least one step")
        if not 0.0 < self.elite_fraction <= 1.0:
            raise ValueError(f"elite fraction must be in (0, 1], got {self.elite_fraction}")
        if self.action_low >= self.action_high:
            raise ValueError("action bounds must be ordered")

    @property
    def elite_count(self) -> int:
        return max(2, int(self.candidates * self.elite_fraction))


class ProbabilisticDynamics(nn.Module):
    """Predicts a DISTRIBUTION over the state delta, not a point.

    The variance head is not decoration. It separates the two kinds of
    uncertainty a planner has to treat differently:

      - the predicted variance is ALEATORIC - noise in the environment, which
        more data will not reduce and which should NOT make a planner cautious
      - disagreement ACROSS ensemble members is EPISTEMIC - ignorance about the
        dynamics, which more data does reduce and which absolutely should

    Conflating them yields a planner that is timid in a noisy but well-mapped
    region and confident in a region it has never visited.
    """

    def __init__(self, config: PlannerConfig, hidden: int = 200) -> None:
        super().__init__()
        width = config.n_state + config.n_action
        self.trunk = nn.Sequential(
            nn.Linear(width, hidden), nn.SiLU(),
            nn.Linear(hidden, hidden), nn.SiLU(),
        )
        self.mean_head = nn.Linear(hidden, config.n_state)
        self.log_var_head = nn.Linear(hidden, config.n_state)

    def forward(self, states: Tensor, actions: Tensor) -> tuple[Tensor, Tensor]:
        features = self.trunk(torch.cat([states, actions], dim=-1))
        # Bounded: an unconstrained log-variance runs away to explain a hard
        # sample as pure noise, and the Gaussian likelihood rewards it for that.
        log_var = self.log_var_head(features).clamp(*LOG_VAR_BOUNDS)
        return self.mean_head(features), log_var


def gaussian_nll(mean: Tensor, log_var: Tensor, target_delta: Tensor) -> Tensor:
    """Negative log-likelihood of the observed delta.

    The target is the DELTA, never the absolute next state: consecutive states
    are correlated enough that predicting levels gives excellent held-out error
    and almost no information about how the system changes.
    """
    inverse_var = torch.exp(-log_var)
    return (((mean - target_delta) ** 2) * inverse_var + log_var).mean()


class DynamicsEnsemble(nn.Module):
    """Independently initialized members, trained on bootstrapped batches.

    Independence is the point: correlated members agree everywhere, including
    where they are all wrong, and the disagreement signal disappears exactly
    when it is needed.
    """

    def __init__(self, config: PlannerConfig, size: int = 5) -> None:
        super().__init__()
        self.members = nn.ModuleList([ProbabilisticDynamics(config) for _ in range(size)])

    def predict(self, states: Tensor, actions: Tensor) -> tuple[Tensor, Tensor]:
        means = torch.stack([member(states, actions)[0] for member in self.members])
        # Mean prediction, and the epistemic term as spread ACROSS members.
        return means.mean(dim=0), means.std(dim=0).sum(dim=-1)


def cross_entropy_plan(
    ensemble: DynamicsEnsemble,
    reward_fn: Callable[[Tensor, Tensor], Tensor],
    state: Tensor,
    config: PlannerConfig,
    warm_start: Tensor | None = None,
) -> Tensor:
    """Refit a distribution over action SEQUENCES to its own elites.

    Random shooting resamples uniformly every time and learns nothing from what
    it already tried. CEM keeps a Gaussian over sequences and refits it to the
    best candidates each iteration, which concentrates the search where the
    returns are and converges far faster for the same evaluation budget.
    """
    mean = (
        warm_start
        if warm_start is not None
        else torch.zeros(config.horizon, config.n_action)
    )
    std = torch.full_like(mean, (config.action_high - config.action_low) / 2.0)

    for _ in range(config.iterations):
        noise = torch.randn(config.candidates, config.horizon, config.n_action)
        sequences = (mean + std * noise).clamp(config.action_low, config.action_high)

        returns = evaluate_sequences(ensemble, reward_fn, state, sequences, config)
        elites = sequences[returns.topk(config.elite_count).indices]

        mean = elites.mean(dim=0)
        std = elites.std(dim=0) + 1e-6

    return mean


@torch.no_grad()
def evaluate_sequences(
    ensemble: DynamicsEnsemble,
    reward_fn: Callable[[Tensor, Tensor], Tensor],
    state: Tensor,
    sequences: Tensor,
    config: PlannerConfig,
) -> Tensor:
    """Roll every candidate forward through the model and score it.

    Each step predicts from the model's OWN previous prediction, so error
    compounds down the horizon - which is why the epistemic penalty accumulates
    alongside the reward rather than being applied once at the end.
    """
    imagined = state.expand(sequences.shape[0], -1).clone()
    total = torch.zeros(sequences.shape[0])
    discount = 1.0

    for step in range(config.horizon):
        actions = sequences[:, step, :]
        delta, disagreement = ensemble.predict(imagined, actions)
        imagined = imagined + delta

        reward = reward_fn(imagined, actions)
        # Penalize where the ensemble disagrees, not where it predicts noise:
        # only ignorance is reducible, and only ignorance should deter a planner.
        total += discount * (reward - config.uncertainty_penalty * disagreement)
        discount *= config.gamma

    return total


def control_loop(env, ensemble, reward_fn, config: PlannerConfig, steps: int) -> None:
    state = torch.as_tensor(env.reset(), dtype=torch.float32)
    previous_plan: Tensor | None = None

    for _ in range(steps):
        plan = cross_entropy_plan(ensemble, reward_fn, state, config, previous_plan)

        # Receding horizon: execute the first action only, then replan from a
        # real observation. The plan is a sequence and all but one step of it is
        # discarded, which is what stops model error from accumulating.
        next_state, _, done = env.step(plan[0].tolist())

        # Warm start: shift the surviving plan forward one step. The world moved
        # by one step and the old plan is still nearly right for the rest.
        previous_plan = torch.cat([plan[1:], plan[-1:].clone()])

        state = torch.as_tensor(env.reset() if done else next_state, dtype=torch.float32)
        if done:
            previous_plan = None

"""MPC - the whole planner as one batched tensor program.

Planning cost is candidates x horizon x ensemble per decision, inside a control
cycle. It is also embarrassingly parallel in candidates AND in ensemble
members, which makes this one of the few places in this category where the
work maps onto batched hardware essentially perfectly.

The horizon is the one axis that genuinely cannot be parallelized: step k+1
depends on step k. Everything else collapses.
"""

from __future__ import annotations

from typing import Callable

import torch
from torch import Tensor


class BatchedEnsemble(torch.nn.Module):
    """All ensemble members as one weight tensor with a leading member axis.

    Separate modules mean E small GEMMs per layer, each too small to saturate
    anything, plus E kernel launches. A leading member dimension turns every
    layer into one baddbmm over (E, batch, width) - identical arithmetic, one
    launch - and the members stay independent, which is what the disagreement
    signal requires.
    """

    def __init__(self, size: int, n_input: int, n_state: int, hidden: int = 200) -> None:
        super().__init__()
        self.w1 = torch.nn.Parameter(torch.empty(size, n_input, hidden))
        self.b1 = torch.nn.Parameter(torch.zeros(size, 1, hidden))
        self.w2 = torch.nn.Parameter(torch.empty(size, hidden, n_state))
        self.b2 = torch.nn.Parameter(torch.zeros(size, 1, n_state))
        for tensor in (self.w1, self.w2):
            torch.nn.init.orthogonal_(tensor, gain=1.0)

    def forward(self, joint: Tensor) -> Tensor:
        """joint is (E, batch, n_input); returns (E, batch, n_state) deltas."""
        hidden = torch.baddbmm(self.b1, joint, self.w1)
        hidden = torch.nn.functional.silu(hidden)
        return torch.baddbmm(self.b2, hidden, self.w2)


class PlanBuffers:
    """Every tensor the planner needs, allocated once at construction.

    A planning loop that allocates runs inside a control cycle and does so
    thousands of times per second. Sizing the buffers once turns the planner
    into a fixed-cost operation, which is what a real-time deadline requires.
    """

    def __init__(self, ensemble_size: int, candidates: int, horizon: int,
                 n_state: int, n_action: int, device: torch.device) -> None:
        # Single dtype throughout, contiguous, reused every iteration.
        self.imagined = torch.zeros((ensemble_size, candidates, n_state),
                                    dtype=torch.float32, device=device)
        self.joint = torch.zeros((ensemble_size, candidates, n_state + n_action),
                                 dtype=torch.float32, device=device)
        self.returns = torch.zeros(candidates, dtype=torch.float32, device=device)
        self.noise = torch.zeros((candidates, horizon, n_action),
                                 dtype=torch.float32, device=device)
        self.n_state = n_state


@torch.no_grad()
def evaluate_batched(
    ensemble: BatchedEnsemble,
    reward_fn: Callable[[Tensor, Tensor], Tensor],
    state: Tensor,
    sequences: Tensor,
    buffers: PlanBuffers,
    gamma: float,
    penalty: float,
) -> Tensor:
    """Every candidate and every ensemble member, advanced together.

    One model call per HORIZON STEP covers the entire search - the naive form
    makes candidates x ensemble separate calls per step. The horizon loop is
    all that remains, and it has to: step k+1 predicts from step k's output.
    """
    ensemble_size, candidates, n_state = buffers.imagined.shape

    buffers.imagined.copy_(state.view(1, 1, n_state).expand(ensemble_size, candidates, n_state))
    buffers.returns.zero_()
    discount = 1.0

    for step in range(sequences.shape[1]):
        actions = sequences[:, step, :].unsqueeze(0).expand(ensemble_size, -1, -1)

        # Written into the pre-allocated joint buffer rather than concatenated
        # into a fresh tensor on every one of thousands of steps.
        buffers.joint[..., :n_state].copy_(buffers.imagined)
        buffers.joint[..., n_state:].copy_(actions)

        deltas = ensemble(buffers.joint)          # (E, candidates, n_state)
        buffers.imagined.add_(deltas)             # in place; deltas, not levels

        # Epistemic term: spread ACROSS members, computed on the same tensor
        # that was just produced rather than from a second pass.
        disagreement = buffers.imagined.std(dim=0).sum(dim=-1)
        mean_state = buffers.imagined.mean(dim=0)

        reward = reward_fn(mean_state, sequences[:, step, :])
        buffers.returns.add_(discount * (reward - penalty * disagreement))
        discount *= gamma

    return buffers.returns


@torch.no_grad()
def cem_plan(
    ensemble: BatchedEnsemble,
    reward_fn: Callable[[Tensor, Tensor], Tensor],
    state: Tensor,
    buffers: PlanBuffers,
    previous_plan: Tensor | None,
    horizon: int,
    n_action: int,
    iterations: int,
    elite_count: int,
    bounds: tuple[float, float],
    gamma: float,
    penalty: float,
) -> Tensor:
    """CEM, warm-started from the surviving tail of the last plan.

    Warm-starting is the cheapest real speedup in MPC and it is not an
    approximation: the world advanced one step, so last cycle's plan shifted
    forward by one is already close to optimal, and the search starts from
    there instead of from scratch. It typically removes an entire CEM iteration.
    """
    if previous_plan is None:
        mean = torch.zeros(horizon, n_action, device=state.device)
    else:
        mean = torch.cat([previous_plan[1:], previous_plan[-1:]])

    std = torch.full_like(mean, (bounds[1] - bounds[0]) / 2.0)

    for _ in range(iterations):
        # normal_ fills the pre-allocated noise buffer in place: no allocation
        # anywhere in the planning loop.
        buffers.noise.normal_()
        sequences = torch.addcmul(mean, buffers.noise, std).clamp_(*bounds)

        returns = evaluate_batched(
            ensemble, reward_fn, state, sequences, buffers, gamma, penalty
        )
        elites = sequences[returns.topk(elite_count).indices]

        mean = elites.mean(dim=0)
        std = elites.std(dim=0).clamp_min_(1e-6)

    return mean