PyTorch dataset for preprocessed DROID trajectories.

Wrapper around Stable Diffusion 3's VAE encoder (AutoencoderKL). Supports both scaling factor and shift factor for latent normalization.

Self-attention with RoPE, optional causal mask, and QK-norm. RoPE is precomputed at init from rope_grid_shape. At forward time: - If rope_ids

Container for a single DROID trajectory

Pre-norm attention block for spatial-temporal transformers. - SPATIAL: within-frame attention with PIXEL RoPE. Pass rope_ids for SPRINT. - TE

Cross-attention: queries attend to a separate key-value source.

MLP that maps (current state, action chunk) → future joint + gripper positions.

Loads (joint, gripper) trajectories from DROID-format annotation files.

I3D (Inflated Inception V3) feature extractor for FVD computation. Uses the I3D model pretrained on Kinetics-400, which is the standard f

Inception V3 feature extractor for FID computation. Returns 2048-dimensional features from the last pooling layer.

All accumulated data from the M-chunk imagination loop for one batch.

Wrapper around Stable Diffusion's VAE encoder. Produces latents for images (B, C, H, W) or videos (B, T, C, H, W).