Upload code/solver.py with huggingface_hub

code/solver.py

@@ -0,0 +1,163 @@
+"""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