Load trajectory latent features into memory.

Convert numpy arrays to torch tensors and move to device.

Encode memory observations and assemble into flattened memory tokens. Uses FlowWM's separate memory projectors to encode, then adds

Save checkpoint to a single checkpoint.pt file (atomic overwrite).

Call AFTER optimizer.step().

Encode observations into embeddings. Token structure: [patches_cam1, patches_cam2, ..., state] Args:

Load checkpoint.

Resize uint8 (H, W, 3) → (180, 320) then zero-pad to (224, 224).

Compute FID between real and generated images. Args: real_images: (N, 3, H, W) tensor in [0, 1] fake_images: (M, 3, H, W

Decode WM output back to feature space. Token structure: [patches_cam1, patches_cam2, ..., state, action, time] Args:

Generate full video predictions by autoregressively chunking over the sequence. The window advances by `horizon` frames each step. If `bootst

Update nested dict `cfg` using dot-notation keys like 'training.lr'

Compute mean and covariance of features. Args: features: (N, D) array of features Returns: mu: (D,) mean vector

Extract features from images in batches

Read a specific range of video frames. If preprocessed .npy frames exist, uses memory-mapped loading (very fast). Otherwise falls bac

Apply precomputed cos/sin RoPE to a tensor. cos, sin have shape (..., D) — already expanded to full head dim.

Write a (T, H, W, 3) uint8 array to an MP4 file.

Build a pyramid scheduling matrix for flow-matching diffusion forcing. Returns a (height, horizon) tensor of t-values in [0, 1] where t=0 is

Copy EMA weights into the given model.

Construct WEAVER and move to device. Args: cfg: Full config namespace (from load_eval_config). device:

FVD using I3D features. Expects (N, C, T, H, W) float [0,1].

Encode memory observations using separate memory projectors. Token structure per memory frame: [patches_cam1, patches_ca

Evaluate model with FID/FVD on master process only. Only rank 0 generates videos and computes metrics. Other ranks wait at a barrier.

Find the latest checkpoint in checkpoint_dir and load it into wm.

Load WEAVER config: merge package defaults with checkpoint config, apply overrides.

Return (state_mean, state_std, action_mean, action_std) as float32 numpy arrays.

Convert PI action chunks (numpy) to WEAVER-ready (fut_a, fut_s) torch tensors.

Diffusion forcing inference with pyramid scheduling. Generates future frames using a pyramid schedule where earlier frames denoise fa

Original lockstep inference: all future frames denoised together with the same timestep schedule. Used when diff_forcing=False.

Pack obs/action/noise-level tensors into transformer tokens.

Reshape precomputed features to match target spatial_size. Precomputed features have shape (B, T, N_src, D_src) where: - D_s

Extract sliding-window clips from a batch of videos (StyleGAN-V style). Args: videos: (N, 3, T, H, W) tensor window_size: nu

Convert to tensor. Handles List[Tensor] (variable-length) by padding to max T.

Calculate Fréchet distance between two Gaussian distributions. FID = ||mu1 - mu2||^2 + Tr(sigma1 + sigma2 - 2*sqrt(sigma1*sigma2))

(T, 7) joint angles → (T, 6) [x, y, z, rx, ry, rz] end-effector pose.

Computes LPIPS safely in chunks to avoid OOM. Works for both images and frames of videos. Args: real, fake: (N, C, H, W) or

Encode a single real observation frame into WM latent features.

Eval-only forward that reuses prefix temporal K/V across denoising steps.

KV-cache eval variant of generate_latent_rollouts. The prefix is memory + history frames. It is clean/fixed for all denoising steps i

Preprocess video frames: normalize and rearrange to (T, C, H, W)

Load a pre-trained Ctrl-World Dynamics model from a .pth checkpoint.

Return a PIL ImageFont at the requested size using Pillow's built-in default.

Stack videos to a fixed length. When pad_short_clips is False, clips shorter than max_t are dropped. When pad_short_clips is True, short clip

Convert B*N PI action chunks to normalized WM video-rate actions and states. Velocity path (dynamics_model=None): Each video-rate step su

Bilinear resize (T, H, W, C) uint8 → (T, h, w, C) uint8.

Resize a single (H, W, C) image to hw=(h, w); return a black frame if img is None.

Sample fixed-length clips for FVD from per-video tensors. Args: vid_list: list of (T, C, H, W) float videos. clip_len: number of

Extract training-style sliding FVD clips from per-video tensors. Mirrors the original WEAVER sliding-window metric implementation: clips are

Free memory by unloading trajectory data