Parse excluded trajectory ids from tokens like '60', '60-79'.

Return True if annotation passes both optional filters. Args: anno: Trajectory annotation dict. filter_episode_id: K

Save a figure comparing GT reward_progress and WM-inferred reward.

Defer SIGINT until after the protected block completes. Prevents Ctrl+C from killing the policy-server connection mid-inference.

Read source video, apply frameskip, resize, reorder cameras. Returns (T, H, 3*W, 3) uint8 RGB array.

Read a specific range of frames from a stacked video using random access. This is much faster than reading the entire video when only a small chu

Render the reward curve as an (frame_height, frame_width, 3) uint8 image. Plots the full curve in blue with a red marker at current_idx.

Compute N-dimensional RoPE frequency table (cached). Args: spatial_shape: Tuple of axis sizes, e.g. (T,) for 1D or (H, W) for 2D.

Half-split rotation for RoPE: [-x2, x1].

Save a GT (top row) / predicted (bottom row) side-by-side video. If reward_values is provided, appends a reward curve row at the bottom of ev

Save a synthetic trajectory: video to videos/<idx>.mp4 and annotation to annotations/<idx>.json. Returns the video path.

Apply SD3 latent scaling: (latents - shift) * scaling

Overlay accumulated reward and (optionally) per-step reward/critic on each frame. Args: frames: (T, H, W, 3) uint8 numpy array.

Send one action to the Franka robot via deoxys. Converts 7-D joint-velocity + 1-D gripper action to a joint-position command and forwards it

Reverse SD3 latent scaling: latents / scaling + shift

Sample a chunk from a trajectory Returns: Dict containing: - obs: Dict with image latents and states

EMA wrapper using Diffusers' EMAModel. Tracks ONLY parameters that require gradients. Includes warmup where decay ramps from

Args: root: Root directory of preprocessed DROID dataset split: 'train' or 'val' horizon: Number of frame

Truncate time dimension to bootstrap steps if specified.

Build memory_tokens from the buffer at exact t_memory spacing.

Compute FID from two dataloaders. Args: real_dataloader: DataLoader yielding real images fake_dataloader: DataLoader yie

Compute SSIM & PSNR in smaller chunks to avoid OOM. Args: real, fake: (N, C, H, W) or (B, T, C, H, W) batch_size: how ma

Decodes: (B, N, D) → (B, 3, H, W) (B, T, N, D) → (B, T, 3, H, W)

Compute FID, FVD, LPIPS for each pred set vs real videos. Args: real_videos: Dict mapping img_key -> (N, T, 3, H, W) tensor or List[

Args: x: Images of shape (B, 3, H, W) in range [0, 1] Returns: features: (B, 2048) feature vectors

Args: x: Videos of shape (B, 3, T, H, W) in range [0, 1] Returns: features: (B, 400) pre-softmax logits

joint: (B, action_dim) or (B, 1, action_dim) — current joint positions joint_vel: (action_num, action_dim) or (B, acti

Encodes: (B, C, H, W) → (B, 4, H/8, W/8) (B, T, C, H, W) → (B, T, 4, H/8, W/8)

Encodes: (B, C, H, W) → (B, N, D) where N = (H/8/spatial_size) * (W/8/spatial_size), D = latent_channels * spatial_size^2

z: [B, N, D] where N = 16*16 = 256 returns: [B, 3, 256, 256]

x: (B, latent_dim) returns state: (B, out_dim)

Args: x: (B, N, D) input tokens. rope_ids: Optional indices into the RoPE table for sparse/reordered

Args: x: Input tensor (B, T*N, D) T: Number of timesteps rope_ids: Integer indices into the RoPE table fo