code/controller.py

"""Amortized iterative MPC controller over a frozen LeWM world model.

The controller holds one hidden token per action block and repeatedly refines
those tokens. Each refinement decodes the tokens to bounded action blocks,
rolls them through the *frozen* predictor, reads the consequences, and emits a
correction. Both transformers are shared across refinements, so K adds compute
and depth of reasoning but no parameters.
"""

import torch
from torch import nn

from lejepa_control.rollout import (
    goal_distance,
    rollout_contexts,
    rollout_plan,
)


class Block(nn.Module):
    """Pre-norm transformer block."""

    def __init__(self, dim, heads, mlp_ratio=4, dropout=0.1):
        super().__init__()
        self.norm1 = nn.LayerNorm(dim)
        self.attn = nn.MultiheadAttention(
            dim, heads, dropout=dropout, batch_first=True
        )
        self.norm2 = nn.LayerNorm(dim)
        self.mlp = nn.Sequential(
            nn.Linear(dim, mlp_ratio * dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(mlp_ratio * dim, dim),
            nn.Dropout(dropout),
        )

    def forward(self, x):
        h = self.norm1(x)
        x = x + self.attn(h, h, h, need_weights=False)[0]
        return x + self.mlp(self.norm2(x))


class Encoder(nn.Module):
    """Stack of pre-norm blocks with a final norm."""

    def __init__(self, dim, depth, heads, dropout=0.1):
        super().__init__()
        self.blocks = nn.ModuleList(
            Block(dim, heads, dropout=dropout) for _ in range(depth)
        )
        self.norm = nn.LayerNorm(dim)

    def forward(self, x):
        for block in self.blocks:
            x = block(x)
        return self.norm(x)


class IterativeController(nn.Module):
    """Predicts and iteratively refines a continuous action plan.

    Args:
        latent_dim: World-model latent width (192 for LeWM PushT).
        action_dim: Native env action dim (2 for PushT).
        frameskip: Env actions per world-model transition (5).
        horizon: Plan length in world-model transitions (H).
        num_context: Context frames the predictor consumes (N=3).
        refinements: Refinement iterations K.
        action_center / action_scale: Per-dim tanh bounds, expressed in the
            *normalized* action space the world model was trained on. Defaults
            correspond to raw PushT actions in [-1, 1].
    """

    def __init__(
        self,
        latent_dim=192,
        action_dim=2,
        frameskip=5,
        horizon=5,
        num_context=3,
        refinements=3,
        width=256,
        depth=4,
        heads=8,
        dropout=0.1,
        action_center=0.0,
        action_scale=1.0,
        no_latent_proj=False,
        fused=False,
    ):
        super().__init__()
        if no_latent_proj:
            assert width == latent_dim, (
                '--no-latent-proj requires width == latent_dim, got '
                f'width={width} latent_dim={latent_dim}'
            )
        self.horizon = horizon
        self.refinements = refinements
        self.num_context = num_context
        self.frameskip = frameskip
        self.action_dim = action_dim
        self.block_dim = frameskip * action_dim
        self.latent_dim = latent_dim
        self.no_latent_proj = no_latent_proj
        self.fused = fused

        center = torch.as_tensor(action_center).float().expand(action_dim)
        scale = torch.as_tensor(action_scale).float().expand(action_dim)
        self.register_buffer('action_center', center.clone())
        self.register_buffer('action_scale', scale.clone())

        # --- conditioning: context frames + goal ---------------------------
        # identity when tokens already live in the world model's own
        # coordinates (width == latent_dim); a learned projection otherwise
        self.latent_proj = (
            nn.Identity() if no_latent_proj else nn.Linear(latent_dim, width)
        )
        self.context_pos = nn.Parameter(torch.randn(1, num_context, width) * 0.02)
        self.goal_token = nn.Parameter(torch.randn(1, 1, width) * 0.02)

        # --- plan tokens ---------------------------------------------------
        self.plan_query = nn.Parameter(torch.randn(1, horizon, width) * 0.02)
        self.plan_pos = nn.Parameter(torch.randn(1, horizon, width) * 0.02)

        if fused:
            # one operator Phi: per-slot feature [y_j, x_hat, x_hat - x_G,
            # ||x_hat - x_G||^2] -> one projection -> one transformer, used
            # both to seed Y^0 (zeros trick) and to emit refinement deltas
            self.slot_proj = nn.Linear(width + 2 * latent_dim + 1, width)
            self.net = Encoder(width, depth * 2, heads, dropout)
        else:
            # --- consequence encoder F_theta --------------------------------
            # per horizon step: [y_j, x_hat, x_hat - x_G, ||x_hat - x_G||^2]
            self.consequence_proj = nn.Linear(width + 2 * latent_dim + 1, width)
            self.consequence_net = Encoder(width, depth, heads, dropout)

            # --- refinement transformer G_theta ------------------------------
            self.refine_proj = nn.Linear(2 * width, width)
            self.refine_net = Encoder(width, depth, heads, dropout)

        # small (not zero) init: refinements start near-identity, but the
        # consequence transformer still receives gradient on the first step
        self.delta_head = nn.Linear(width, width)
        nn.init.normal_(self.delta_head.weight, std=0.01)
        nn.init.zeros_(self.delta_head.bias)

        # one learned step size per refinement, squashed into (0, 1)
        self.step_logit = nn.Parameter(torch.zeros(refinements))

        # --- action head ---------------------------------------------------
        self.action_head = nn.Sequential(
            nn.LayerNorm(width),
            nn.Linear(width, width),
            nn.GELU(),
            nn.Linear(width, self.block_dim),
        )

    def step_sizes_repr(self):
        """Learned refinement step sizes, for logging."""
        with torch.no_grad():
            return [round(v, 3) for v in torch.sigmoid(self.step_logit).tolist()]

    def condition(self, ctx_emb, goal_emb):
        """Build conditioning tokens from context latents and the goal."""
        ctx = self.latent_proj(ctx_emb) + self.context_pos
        goal = self.latent_proj(goal_emb).unsqueeze(1) + self.goal_token
        return torch.cat([ctx, goal], dim=1)  # (B, N+1, W)

    def to_actions(self, plan_tokens):
        """Decode plan tokens to bounded action blocks ``(B, H, 5*d_a)``."""
        raw = self.action_head(plan_tokens)
        raw = raw.unflatten(-1, (self.frameskip, self.action_dim))
        bounded = self.action_center + self.action_scale * torch.tanh(raw)
        return bounded.flatten(-2)

    def _run_refine(self, plan_tokens, cond, consequence):
        """Split ``G_theta`` body: plan tokens + consequences -> hidden."""
        x = self.refine_proj(torch.cat([plan_tokens, consequence], dim=-1))
        x = x + self.plan_pos
        x = self.refine_net(torch.cat([x, cond], dim=1))
        return x[:, : self.horizon]

    def _run_fused(self, slot_features, cond):
        """Fused ``Phi`` body: raw per-slot features -> hidden, one operator."""
        x = self.slot_proj(slot_features)
        x = x + self.plan_pos
        x = self.net(torch.cat([x, cond], dim=1))
        return x[:, : self.horizon]

    def initial_plan(self, cond):
        """``Y^(0)`` from learned queries, with no consequences known yet."""
        queries = self.plan_query.expand(cond.size(0), -1, -1)
        if self.fused:
            zeros = queries.new_zeros(
                queries.size(0), self.horizon, 2 * self.latent_dim + 1
            )
            slot_features = torch.cat([queries, zeros], dim=-1)
            # absolute, not a delta: same "zeros trick" semantics as the
            # split path's initial_plan, routed through the same operator
            return self._run_fused(slot_features, cond)
        return self._run_refine(queries, cond, torch.zeros_like(queries))

    def refine(self, plan_tokens, cond, pred, goal_emb):
        """One refinement: read consequences, emit a correction to the plan."""
        delta_goal = pred - goal_emb.unsqueeze(1)
        dist = delta_goal.pow(2).mean(dim=-1, keepdim=True)
        features = torch.cat([plan_tokens, pred, delta_goal, dist], dim=-1)

        if self.fused:
            return self.delta_head(self._run_fused(features, cond))

        cons = self.consequence_proj(features) + self.plan_pos
        cons = self.consequence_net(torch.cat([cons, cond], dim=1))
        cons = cons[:, : self.horizon]

        return self.delta_head(self._run_refine(plan_tokens, cond, cons))

    def forward(self, model, ctx_emb, past_actions, goal_emb):
        """Run the full refinement loop against the frozen world model.

        Args:
            model: Frozen ``LeWM``.
            ctx_emb: ``(B, N, D)`` context latents.
            past_actions: ``(B, N-1, 5*d_a)`` normalized executed blocks.
            goal_emb: ``(B, D)`` goal latent.

        Returns:
            Dict with per-iteration ``plans``, ``rollouts``, ``distances`` and
            ``contexts``; each list holds ``K+1`` entries (initial plan plus K
            refinements).
        """
        cond = self.condition(ctx_emb, goal_emb)
        tokens = self.initial_plan(cond)

        plans, rollouts, distances, contexts = [], [], [], []
        for k in range(self.refinements + 1):
            actions = self.to_actions(tokens)
            pred, frames = rollout_plan(
                model, ctx_emb, past_actions, actions, return_frames=True
            )

            plans.append(actions)
            rollouts.append(pred)
            distances.append(goal_distance(pred, goal_emb))
            contexts.append(rollout_contexts(frames, self.num_context))

            if k == self.refinements:
                break

            delta = self.refine(tokens, cond, pred, goal_emb)
            # running more refinements at eval than were trained reuses the
            # last learned step size rather than indexing past it
            idx = min(k, self.step_logit.numel() - 1)
            tokens = tokens + torch.sigmoid(self.step_logit[idx]) * delta

        return {
            'plans': plans,
            'rollouts': rollouts,
            'distances': distances,
            'contexts': contexts,
            'step_sizes': torch.sigmoid(self.step_logit),
        }