| """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)
|
| )
|
|
|
| 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)
|
|
|
