File size: 2,701 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 | """Differentiable rollout of an action plan through the frozen LeWM predictor.
Mirrors ``LeWM.rollout`` semantics exactly, but operates on cached latents
instead of pixels and keeps the graph intact so gradients reach the plan.
The action/frame alignment is the subtle part and follows ``lewm.py:83-87``:
the action block at index ``k`` is the block *leaving* context frame ``k``.
With ``N`` context frames there are ``N-1`` past blocks between them, and the
current frame pairs with the first block of the plan.
"""
import torch
def rollout_plan(model, ctx_emb, past_actions, plan_actions, return_frames=False):
"""Roll a plan through the frozen predictor, autoregressively in latent space.
Args:
model: The frozen ``LeWM``.
ctx_emb: ``(B, N, D)`` encoded context frames.
past_actions: ``(B, N-1, 5*d_a)`` normalized blocks between them.
plan_actions: ``(B, H, 5*d_a)`` normalized candidate blocks.
return_frames: Also return the full ``(B, N+H, D)`` frame sequence,
which the support loss slices into per-step context windows.
Returns:
``(B, H, D)`` predicted latents, and optionally the full sequence.
"""
B, N, _ = ctx_emb.shape
H = plan_actions.size(1)
history_size = model.predictor.num_frames
act_emb = model.action_encoder(
torch.cat([past_actions, plan_actions], dim=1)
) # (B, N-1+H, A); index k = the block leaving frame k
frames = list(ctx_emb.unbind(dim=1))
for t in range(H):
lo = max(0, N + t - history_size)
emb_win = torch.stack(frames[lo:], dim=1)
act_win = act_emb[:, lo : N + t]
frames.append(model.predict(emb_win, act_win)[:, -1])
pred = torch.stack(frames[N:], dim=1)
if return_frames:
return pred, torch.stack(frames, dim=1)
return pred
def rollout_contexts(frames, num_context):
"""Context window seen by each plan step, as a single batched tensor.
Args:
frames: ``(B, N+H, D)`` full rollout sequence from ``rollout_plan``.
num_context: Window size ``N``.
Returns:
``(B, H, N, D)`` where entry ``j`` is the window preceding block ``j``.
"""
H = frames.size(1) - num_context
return torch.stack(
[frames[:, j : j + num_context] for j in range(H)], dim=1
)
def goal_distance(pred, goal):
"""Per-step latent goal distance ``d_j = ||x_j - x_G||^2 / D``.
Args:
pred: ``(B, H, D)`` predicted latents.
goal: ``(B, D)`` goal latent.
Returns:
``(B, H)`` distances.
"""
return (pred - goal.unsqueeze(1)).pow(2).mean(dim=-1)
|