File size: 5,666 Bytes
872cf4d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 | """Adapter exposing the trained controller as a ``stable_worldmodel`` solver.
Implements the ``Solver`` protocol so the controller drops into
``WorldModelPolicy`` wherever ``CEMSolver`` goes, which makes the CEM
comparison an apples-to-apples swap: same env, same wrappers, same
preprocessing, same receding-horizon execution.
"""
import gymnasium as gym
import torch
from lejepa_control.controller import IterativeController
class ControllerSolver:
"""Runs K amortized refinements instead of CEM's sampling loop.
Args:
model: The frozen ``LeWM``.
controller: A trained :class:`IterativeController`.
device: Device to plan on.
refinements: Override the controller's K at eval time (for the
success-vs-refinement-count ablation). ``None`` keeps the trained
value.
"""
def __init__(self, model, controller, device='cuda', refinements=None):
self.model = model
self.controller = controller.to(device).eval()
self.device = device
self._refinements = refinements
self._n_envs = 1
self._horizon = controller.horizon
self._action_dim = controller.action_dim
self._action_block = controller.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._horizon == self.controller.horizon, (
f'plan horizon {self._horizon} != controller horizon '
f'{self.controller.horizon}'
)
assert self._action_block == self.controller.frameskip, (
f'action_block {self._action_block} != controller frameskip '
f'{self.controller.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):
"""Encode ``(B, T, C, H, W)`` frames to ``(B, T, D)`` latents."""
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': ...}``.
Expects ``pixels`` ``(B, T, C, H, W)``, ``goal`` ``(B, ...)`` and,
when the policy carries history, ``action_history``
``(B, T-1, block*d_a)``.
"""
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.controller.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)
k = self._refinements
original = self.controller.refinements
if k is not None:
self.controller.refinements = k
try:
out = self.controller(self.model, ctx, past, goal_emb)
finally:
self.controller.refinements = original
actions = out['plans'][-1] # (B, H, block*d_a)
terminal = out['distances'][-1][:, -1]
return {
'actions': actions.detach().float().cpu(),
'costs': terminal.detach().float().cpu(),
'terminal_distance': terminal.mean().item(),
}
__call__ = solve
def load_controller(path, latent_dim=192, device='cuda', refinements=None):
"""Rebuild a controller 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'])
controller = IterativeController(
latent_dim=latent_dim,
horizon=saved['horizon'],
refinements=saved['refinements'],
width=saved['width'],
depth=saved['depth'],
heads=saved['heads'],
dropout=saved['dropout'],
action_center=(-a_mean / a_std),
action_scale=(1.0 / a_std),
no_latent_proj=saved.get('no_latent_proj', False),
fused=saved.get('fused', False),
)
controller.load_state_dict(ckpt['state_dict'])
controller.to(device).eval()
if refinements is not None:
controller.refinements = refinements
return controller, ckpt
|