code/solver.py

"""Adapter exposing the trained planner as a ``stable_worldmodel`` solver.

Mirrors :class:`lejepa_control.solver.ControllerSolver` so the planner drops
into ``WorldModelPolicy`` wherever the baseline controller or ``CEMSolver``
goes. Same env, same wrappers, same preprocessing, same receding-horizon
execution — the only thing that differs is how the plan is produced, which is
what makes the baseline comparison a fair swap.
"""

import gymnasium as gym
import torch

from lejepa_control_2.planner import RecursivePlanner


class PlannerSolver:
    """Runs the three-loop recursion instead of CEM's sampling loop.

    Args:
        model: The frozen ``LeWM``.
        planner: A trained :class:`RecursivePlanner`.
        device: Device to plan on.
        cycles / inner: Override ``T`` / ``n`` at eval time. This is the
            anytime-inference sweep (ablation 1): the recursion is weight-tied,
            so the same checkpoint can be run at any depth. ``None`` keeps the
            trained values.
    """

    def __init__(self, model, planner, device='cuda', cycles=None, inner=None):
        self.model = model
        self.planner = planner.to(device).eval()
        self.device = device
        self.cycles = cycles
        self.inner = inner
        self._n_envs = 1
        self._horizon = planner.horizon
        self._action_dim = planner.action_dim
        self._action_block = planner.frameskip

    def configure(self, *, action_space: gym.Space, n_envs: int, config) -> None:
        self._n_envs = n_envs
        self._horizon = config.horizon
        self._action_block = config.action_block
        self._action_dim = int(action_space.shape[-1])

        assert self._action_block == self.planner.frameskip, (
            f'action_block {self._action_block} != planner frameskip '
            f'{self.planner.frameskip}'
        )

    @property
    def action_dim(self) -> int:
        return self._action_dim * self._action_block

    @property
    def n_envs(self) -> int:
        return self._n_envs

    @property
    def horizon(self) -> int:
        return self._horizon

    def _encode(self, pixels):
        with torch.no_grad():
            return self.model.encode({'pixels': pixels.to(self.device)})['emb']

    @torch.no_grad()
    def solve(self, info_dict: dict, init_action=None) -> dict:
        """Plan for every env in ``info_dict``; returns ``{'actions': ...}``."""
        pixels = info_dict['pixels']
        if pixels.ndim == 4:  # (B, C, H, W) -> single context frame
            pixels = pixels.unsqueeze(1)
        B, T = pixels.shape[:2]

        ctx = self._encode(pixels)

        goal = info_dict['goal']
        if goal.ndim == 4:
            goal = goal.unsqueeze(1)
        goal_emb = self._encode(goal)[:, -1]

        num_context = self.planner.num_context
        if T < num_context:  # early in an episode: repeat the oldest frame
            pad = ctx[:, :1].expand(B, num_context - T, -1)
            ctx = torch.cat([pad, ctx], dim=1)
        elif T > num_context:
            ctx = ctx[:, -num_context:]

        block_dim = self._action_block * self._action_dim
        past = info_dict.get('action_history')
        if past is None:
            past = ctx.new_zeros(B, num_context - 1, block_dim)
        else:
            past = past.to(self.device).float()
            if past.size(1) < num_context - 1:
                pad = past.new_zeros(
                    B, num_context - 1 - past.size(1), block_dim
                )
                past = torch.cat([pad, past], dim=1)
            else:
                past = past[:, -(num_context - 1) :]
            past = torch.nan_to_num(past, 0.0)

        out = self.planner(
            self.model,
            ctx,
            past,
            goal_emb,
            horizon=self._horizon,
            cycles=self.cycles,
            inner=self.inner,
        )

        actions = out['blocks']  # (B, H, block*d_a)
        terminal = out['distances'][:, -1]

        result = {
            'actions': actions.detach().float().cpu(),
            'costs': terminal.detach().float().cpu(),
            'terminal_distance': terminal.mean().item(),
        }
        if out['cycle_distances'] is not None:
            # the stage-B diagnostic, measured at deployment rather than on
            # the training distribution
            result['per_cycle'] = (
                out['cycle_distances'].mean(dim=(0, 1)).detach().cpu().tolist()
            )
        return result

    __call__ = solve


def load_planner(path, device='cuda', cycles=None, inner=None, horizon=None):
    """Rebuild a planner from a training checkpoint."""
    ckpt = torch.load(path, map_location=device, weights_only=False)
    saved = ckpt['args']
    a_mean = torch.tensor(ckpt['action_mean'])
    a_std = torch.tensor(ckpt['action_std'])

    planner = RecursivePlanner(
        latent_dim=ckpt['latent_dim'],
        width=saved['width'],
        hidden=saved['hidden'],
        inner=saved['inner'],
        cycles=saved['cycles'],
        horizon=ckpt.get('horizon', saved.get('horizon', 5)),
        use_feedback=saved['use_feedback'],
        warm_start=saved['warm_start'],
        lambda_z=saved['lambda_z'],
        learn_lambda_z=saved['learn_lambda_z'],
        detach_schedule=saved['detach_schedule'],
        action_center=(-a_mean / a_std),
        action_scale=(1.0 / a_std),
    )
    planner.load_state_dict(ckpt['state_dict'])
    planner.to(device).eval()
    if cycles is not None:
        planner.cycles = cycles
    if inner is not None:
        planner.inner = inner
    if horizon is not None:
        planner.horizon = horizon
    return planner, ckpt