File size: 5,683 Bytes
d0c9d97 | 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 158 159 160 161 162 163 164 | """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
|