Imagine one chunk and score with reward model + critic. Rolls out WEAVER for future_actions.shape[1] steps starting from context, then evalua

Aggregate individual losses into a single total_loss for backward. Returns: total_loss: scalar tensor (live graph for backward)

Truncate to min length for LPIPS (requires same shape).

Build B PI request dicts using real history frames (no WM decode needed).

Build B*N PI request dicts from decoded imagined frames.

Commit one video-rate frame: update the real-frame buffers and sync GT/imagined panels. Called every RGB_SKIP control steps to keep the video-rat

Phase 3 — Compute per-segment selection scores and return the best sample index. Returns: best_ns: (B,) int tensor — best sample index

Drop image patch tokens based on pre-computed ids_keep. Args: x: (B, T * N_total, D) flattened input T: Numb

Compute 4×4 end-effector transform for 7-DOF Panda joints.

Core sparse gather: keep only selected image tokens, preserve other tokens. Args: x: (B, N_total, D) where first n_img_t

Generate random indices for SPRINT token dropping. Drops from image patch tokens only. Same ids_keep across all timesteps fo

Get cos/sin tensors for rope_mix. Args: rope_ids: (B, N) or (N,) indices into the precomputed table. Returns:

Compute rope_ids for sparse image tokens by gathering from the precomputed dense rope_ids. The dense rope_ids already encode

Load video, split top/bottom and by view. Returns (real_views, pred_views) per key.

Load trajectory annotation file

Download and cache the I3D TorchScript model.

Pad sparse image tokens back to full size with [MASK] tokens. Args: x_sparse: (B, T * N_sparse, D) T: Number

Phase 1 — Encode the real context and query PI for the first action chunk. Encodes the n_history real frames as WM context, queries PI for N acti

Build context dict, history action window, and memory tokens from real-frame buffers.

Extract video-rate states and actions from a raw DROID annotation dict.

Preprocess full trajectory states and actions from annotation. Called once during __init__ to cache preprocessed data. Note: The ann

Rollout via KV-cache (no reward scoring). Builds the combined (context + future) input tensors and calls generate_latent_rollouts_cached, whi

Phase 2 — Run M imagination chunks through WEAVER, interleaved with PI queries. Each chunk: a. Rolls out WEAVER for exec_len video-rate ste

Query PI, imagine all N samples, score, and return the best sample's execution state. Returns: pred_action_chunk: (act_horizon, 8) contro

Return clip start indices for a single video.

Decode all B*N imagined trajectories and save per-segment thumbnail comparison videos.

Save side-by-side GT (top) / imagined WM (bottom) video for one trial.

Phase 4 — Decode, annotate, and save one winning trajectory. Returns the saved video path, or None if the segment was filtered out.

Build memory tokens from t_memory-spaced real observation history.

Strip DDP/compile prefixes (module., _orig_mod.) from checkpoint keys.

Discover all (task_dir, ep_id) pairs across all input roots.

FID using pytorch-fid, streaming frames as PNGs to temp dirs.

Compute FVD between real and generated videos using the StyleGAN-V approach. Instead of extracting one feature vector per video, this uses a

LPIPS (VGG) computed streaming, one video at a time.

PSNR from TorchMetrics for images normalized to [0, 1].

SSIM from TorchMetrics for images normalized to [0, 1].

Temporal consistency loss: MSE((f_h - f_{h-1}) - (v_h - v_{h-1})). For v-pred: v_h = x1_h - x0_h (true velocity), f_h = pred_h (predicted vel

Build a DataLoader for synthetic trajectory generation. Creates a PrecomputedDroid dataset with video frames enabled (needed for the first PI

Decodes: (B, C, H/8, W/8) → (B, 3, H, W) (B, T, C, H/8, W/8) → (B, T, 3, H, W)

Auto-discover camera names from the first available trajectory directory.

Encode (T, H, W, 3) uint8 frames through SD3. Returns (T, tokens, channels).

Eval-only future-token pass using cached prefix temporal K/V.

Eval-only full-sequence pass that returns prefix temporal K/V. This is intentionally separate from ``forward`` so training and torch.

Generate rollouts for a variable number of future frames (no padding). Args: context: Encoded context dict with keys in self._img

Autoregressively predict future frames and return per-frame reward scores. The trajectory is divided into chunks of `bootstrap` frames. Each chun

Get sorted list of trajectory directories present in both dirs with all cameras.

Predict future absolute joint + gripper positions via the dynamics model. Pads the chunk to n_steps if shorter; returns only up to len(chunk) fra

Load a single camera across all trajectories. Returns lists of per-video tensors, each (T_i, C, H, W) float [0,1].

Load a full trajectory: latent obs, GT video frames, states, actions, and memory. obs and actions span [0, end_frame] so the generation loop can

Convert numpy arrays to torch tensors and move to device.