Frequency tiers (estimates, relative to interviews targeting this pack's specialty — a ★★★★★ here means "core for these roles", not for generalist ML loops): ★★★★★ expect it in most loops · ★★★★ very common · ★★★ common · ★★ occasional · ★ rare/deep-specialist. Spend ~50% of practice on the top two tiers.
60+ problems with full PyTorch solutions and tests
Principal/Senior-Principal Generative-Vision Interview Prep
Notes
This pack covers the live-coding questions that come up at principal-level interviews on modern multi-modal generation: video generation (Sora-style spacetime DiT, MM-DiT, video VAEs, temporal attention, motion modules, frame interpolation), VLA (Vision-Language-Action models for robotics — RT-2, Pi-0, OpenVLA, ac- tion tokenization, expert chunking), and world models (Dreamer-style latent dynamics, GAIA, Genie, Cosmos, action-conditioned rollouts, jointly-embedded planners, V-JEPA-style predictive coding). Standard imports across the pack:
import math, numpy as np, torch
import torch.nn as nn, torch.nn.functional as F
Problem
Split a video (B, C, T, H, W) into non-overlapping pt×ps×ps patches and project to dim D via 3D conv.
import torch.nn as nn
class VideoPatchEmbed(nn.Module):
def __init__(self, in_ch=3, dim=768, p_t=2, p_s=16):
super().__init__()
self.proj = nn.Conv3d(in_ch, dim, kernel_size=(p_t, p_s, p_s), stride=(p_t, p_s, p_s))
def forward(self, x): # (B, C, T, H, W)
x = self.proj(x) # (B, D, T', H', W')
return x.flatten(2).transpose(1, 2) # (B, N, D)
Tests
m = VideoPatchEmbed(); y = m(torch.randn(1, 3, 16, 64, 64))
assert y.shape == (1, 8 * 4 * 4, 768); print("video patch OK")
Problem
A causal 3D conv only sees current and past frames: pad on the temporal axis only at the front.
import torch.nn as nn
import torch.nn.functional as F
class CausalConv3d(nn.Module):
def __init__(self, in_ch, out_ch, kt=3, ks=3):
super().__init__()
self.kt = kt
self.conv = nn.Conv3d(in_ch, out_ch, (kt, ks, ks), padding=(0, ks // 2, ks // 2))
def forward(self, x): # (B, C, T, H, W)
x = F.pad(x, (0, 0, 0, 0, self.kt - 1, 0))
return self.conv(x)
Tests
m = CausalConv3d(3, 16); y = m(torch.randn(1, 3, 8, 32, 32))
assert y.shape == (1, 16, 8, 32, 32); print("causal3d OK")
Problem
Halve spatial resolution with stride-2 conv, then halve temporal resolution with a stride-2 causal 3D conv.
import torch.nn as nn
import torch.nn.functional as F
class VideoDown(nn.Module):
def __init__(self, c):
super().__init__()
self.s = nn.Conv3d(c, c, (1, 3, 3), stride=(1, 2, 2), padding=(0, 1, 1))
self.t = nn.Conv3d(c, c, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0))
def forward(self, x):
x = self.s(x)
x = F.pad(x, (0, 0, 0, 0, 2, 0))
return self.t(x)
Tests
m = VideoDown(8); y = m(torch.randn(1, 8, 8, 32, 32))
assert y.shape[2:] == (4, 16, 16); print("video down OK")
Problem
Standard VAE KL between q(z|x) ∼N(µ, σ2) and N(0, I), summed over spatiotemporal locations.
def vae_kl(mu, logvar):
return -0.5 * (1 + logvar - mu.pow(2) - logvar.exp()).flatten(1).sum(dim=1).mean()
Tests
mu = torch.zeros(2, 4, 4, 4); lv = torch.zeros_like(mu)
assert abs(vae_kl(mu, lv).item()) < 1e-6; print("KL OK")
Problem
Map a continuous latent to a discrete code via Finite Scalar Quantization (FSQ): each channel is rounded to one of Li levels in [−1, 1], with straight-through gradient.
import torch
def fsq(x, levels):
# levels: list of ints per channel; same length as last dim
L = torch.tensor(levels, device=x.device, dtype=x.dtype)
x = torch.tanh(x)
half = (L - 1) / 2
q = torch.round(x * half) / half
return x + (q - x).detach()
Tests
out = fsq(torch.randn(4, 6), levels=[8, 8, 8, 8, 8, 8])
assert out.shape == (4, 6); print("FSQ OK")
Problem
Apply spatial attention across HW at each timestep, then temporal attention across T at each spatial location, sharing parameters in a single block.
import torch.nn as nn
class STAttn(nn.Module):
def __init__(self, d, h):
super().__init__()
self.spatial = nn.MultiheadAttention(d, h, batch_first=True)
self.temporal = nn.MultiheadAttention(d, h, batch_first=True)
def forward(self, x): # (B, T, S, D)
B, T, S, D = x.shape
# spatial within each frame
xs = x.view(B * T, S, D)
xs, _ = self.spatial(xs, xs, xs)
x = xs.view(B, T, S, D)
# temporal at each location
xt = x.transpose(1, 2).contiguous().view(B * S, T, D)
xt, _ = self.temporal(xt, xt, xt)
return xt.view(B, S, T, D).transpose(1, 2)
Tests
m = STAttn(32, 4); y = m(torch.randn(1, 4, 16, 32))
assert y.shape == (1, 4, 16, 32); print("ST attn OK")
Advanced note. Full spatiotemporal attention costs O((T·S)²) — at 16 frames x 1024 tokens that's a 16k x 16k score matrix per head. Factorized attention (spatial within each frame, then temporal across frames at each location) costs O(T·S² + S·T²) — here ~256x fewer score entries — and is the standard trade (TimeSformer, and the pattern inside video DiTs before full 3D attention became affordable). State the ratio (T·S)²/(T·S² + S·T²) = T·S/(S+T) unprompted; it's the whole argument.
Problem
DiT: AdaLN-Zero modulation around self-attention and MLP, with learnable scale α for each residual branch.
import torch.nn as nn
import torch.nn.functional as F
class DiTBlock(nn.Module):
def __init__(self, d, h, mlp_ratio=4):
super().__init__()
self.attn = nn.MultiheadAttention(d, h, batch_first=True)
self.mlp = nn.Sequential(nn.Linear(d, mlp_ratio * d), nn.GELU(),
nn.Linear(mlp_ratio * d, d))
self.adaLN = nn.Linear(d, 6 * d)
nn.init.zeros_(self.adaLN.weight); nn.init.zeros_(self.adaLN.bias)
def forward(self, x, c):
sh1, sc1, ga1, sh2, sc2, ga2 = self.adaLN(F.silu(c)).chunk(6, dim=-1)
def mod(h, sh, sc):
return (1 + sc.unsqueeze(1)) * F.layer_norm(h, h.shape[-1:]) + sh.unsqueeze(1)
h_attn, _ = self.attn(mod(x, sh1, sc1), mod(x, sh1, sc1), mod(x, sh1, sc1))
x = x + ga1.unsqueeze(1) * h_attn
x = x + ga2.unsqueeze(1) * self.mlp(mod(x, sh2, sc2))
return x
Tests
m = DiTBlock(64, 4); y = m(torch.randn(2, 16, 64), torch.randn(2, 64))
assert y.shape == (2, 16, 64); print("DiT block OK")
Problem
Concatenate text and image tokens, run joint self-attention; both modalities fully attend to each other at every layer.
import torch
import torch.nn as nn
class MMDiTBlock(nn.Module):
def __init__(self, d, h):
super().__init__()
self.attn = nn.MultiheadAttention(d, h, batch_first=True)
self.ln_i = nn.LayerNorm(d); self.ln_t = nn.LayerNorm(d)
self.mlp_i = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))
self.mlp_t = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))
def forward(self, img_tok, txt_tok):
x = torch.cat([img_tok, txt_tok], dim=1)
x_n = torch.cat([self.ln_i(img_tok), self.ln_t(txt_tok)], dim=1)
x_attn, _ = self.attn(x_n, x_n, x_n)
x = x + x_attn
i, t = x.split([img_tok.size(1), txt_tok.size(1)], dim=1)
i = i + self.mlp_i(self.ln_i(i))
t = t + self.mlp_t(self.ln_t(t))
return i, t
Tests
m = MMDiTBlock(64, 4)
i = torch.randn(1, 16, 64); t = torch.randn(1, 8, 64)
oi, ot = m(i, t); assert oi.shape == i.shape and ot.shape == t.shape
print("MM-DiT OK")
Problem
Treat the video latent as a sequence of spatiotemporal tokens; flatten to one sequence and run a standard Transformer with sinusoidal time and 2D positional embeddings.
import torch.nn as nn
class SpacetimeDiT(nn.Module):
def __init__(self, dim=256, n_blocks=4, n_heads=4):
super().__init__()
self.blocks = nn.ModuleList([DiTBlock(dim, n_heads) for _ in range(n_blocks)])
def forward(self, x_seq, cond):
for b in self.blocks: x_seq = b(x_seq, cond)
return x_seq
Tests
m = SpacetimeDiT(); y = m(torch.randn(1, 64, 256), torch.randn(1, 256))
assert y.shape == (1, 64, 256); print("Spacetime DiT OK")
Problem
Add a learnable 3D positional encoding (separable: spatial 2D + temporal 1D) to a video token sequence.
import math
import torch
import torch.nn as nn
class SeparablePE(nn.Module):
def __init__(self, T, H, W, d):
super().__init__()
self.t = nn.Parameter(torch.zeros(T, d) * 0.02)
self.h = nn.Parameter(torch.zeros(H, d) * 0.02)
self.w = nn.Parameter(torch.zeros(W, d) * 0.02)
def forward(self, x_btshw):
B, T, S, _ = x_btshw.shape
H = W = int(math.sqrt(S))
pos = self.t.unsqueeze(1).unsqueeze(1) + self.h.unsqueeze(0).unsqueeze(2) + self.w.unsqueeze(0).unsqueeze(1)
return x_btshw + pos.flatten(1, 2).unsqueeze(0)
Tests
pe = SeparablePE(4, 4, 4, 16)
y = pe(torch.randn(1, 4, 16, 16))
assert y.shape == (1, 4, 16, 16); print("sep PE OK")
Problem
Standard ε-prediction loss in spacetime latent space, with optional CFG conditioning drop.
import torch
import torch.nn.functional as F
def video_diffusion_loss(model, x_lat, cond, abar, p_drop=0.1):
B = x_lat.size(0)
t = torch.randint(0, abar.size(0), (B,), device=x_lat.device)
a = abar[t].view(-1, *([1] * (x_lat.dim() - 1)))
eps = torch.randn_like(x_lat)
xt = a.sqrt() * x_lat + (1 - a).sqrt() * eps
if torch.rand(1).item() < p_drop:
cond = torch.zeros_like(cond)
pred = model(xt, t, cond)
return F.mse_loss(pred, eps)
Tests
class M(nn.Module):
def __init__(self): super().__init__(); self.c = nn.Conv3d(4, 4, 3, padding=1)
def forward(self, x, t, c): return self.c(x)
ab = torch.linspace(0.99, 0.01, 50)
loss = video_diffusion_loss(M(), torch.randn(1, 4, 4, 8, 8), torch.randn(1, 16), ab)
assert loss > 0; print("video loss OK")
Problem
Pack videos of different shapes into a single batch via padding + a key-padding mask to ignore padding positions in attention.
import torch
def pack_videos(videos):
# videos: list of (T, S, D)
L = [v.size(0) * v.size(1) for v in videos]
Lmax = max(L); D = videos[0].size(-1)
out = torch.zeros(len(videos), Lmax, D)
mask = torch.ones(len(videos), Lmax, dtype=torch.bool)
for i, v in enumerate(videos):
flat = v.flatten(0, 1)
out[i, :flat.size(0)] = flat
mask[i, :flat.size(0)] = False
return out, mask
Tests
v1 = torch.randn(4, 16, 8); v2 = torch.randn(2, 8, 8)
out, m = pack_videos([v1, v2])
assert out.shape == (2, 64, 8) and m.shape == (2, 64); print("pack OK")
Problem
Conditional DDPM that, given two keyframes, predicts the in-between frame; condition by concatenation along the channel axis.
import torch
import torch.nn as nn
class MidFrameDiff(nn.Module):
def __init__(self):
super().__init__()
self.net = nn.Conv2d(9, 3, 3, padding=1) # 9 = 3 + 3 + 3 (key1, key2, noisy mid)
def forward(self, mid_t, k1, k2, t):
return self.net(torch.cat([mid_t, k1, k2], dim=1))
Tests
m = MidFrameDiff()
y = m(torch.randn(1, 3, 16, 16), torch.randn(1, 3, 16, 16), torch.randn(1, 3, 16, 16),
torch.tensor([0]))
assert y.shape == (1, 3, 16, 16); print("mid-frame OK")
Problem
A small temporal Transformer inserted into a frozen image diffusion model that operates only across the time axis.
import torch.nn as nn
class MotionModule(nn.Module):
def __init__(self, c, n_heads=4):
super().__init__()
self.attn = nn.MultiheadAttention(c, n_heads, batch_first=True)
def forward(self, x): # (B*S, T, C)
return x + self.attn(x, x, x)[0]
Tests
m = MotionModule(16); y = m(torch.randn(64, 8, 16))
assert y.shape == (64, 8, 16); print("motion OK")
Problem
Combine three predictions (uncond, text-cond, image-cond) into a single guided one with separate weights wt, wi.
def multi_cfg(eps_uncond, eps_text, eps_image, w_t=5.0, w_i=2.0):
return eps_uncond + w_t * (eps_text - eps_uncond) + w_i * (eps_image - eps_uncond)
Tests
out = multi_cfg(torch.zeros(4), torch.ones(4), torch.full((4,), 2.0), w_t=1.0, w_i=1.0)
assert torch.allclose(out, torch.full((4,), 3.0)); print("multi CFG OK")
Problem
Generate one chunk at a time; condition the next chunk on the last K frames of the previous one.
import torch
@torch.no_grad()
def autoregressive_video(model, init_frames, chunks=4, chunk_T=8, K=2):
out = [init_frames]
for _ in range(chunks):
prev_tail = out[-1][:, :, -K:]
new_chunk = model.generate(cond=prev_tail, T=chunk_T)
out.append(new_chunk)
return torch.cat(out, dim=2)
Tests
class M:
def generate(self, cond, T): return torch.zeros(1, 3, T, 8, 8)
out = autoregressive_video(M(), torch.zeros(1, 3, 4, 8, 8), chunks=2, chunk_T=4)
assert out.shape == (1, 3, 4 + 2 * 4, 8, 8); print("AR video OK")
Problem
Frechet Video Distance: compute feature statistics from an I3D-style backbone for real and fake videos, compare via FID-style formula.
import numpy as np
def fvd_features_to_score(real_feats, fake_feats):
mu_r, mu_f = real_feats.mean(0), fake_feats.mean(0)
cov_r = np.cov(real_feats, rowvar=False); cov_f = np.cov(fake_feats, rowvar=False)
diff = mu_r - mu_f
from scipy.linalg import sqrtm
s = sqrtm(cov_r @ cov_f)
if np.iscomplexobj(s): s = s.real
return float(diff @ diff + np.trace(cov_r + cov_f - 2 * s))
Tests
r = np.random.randn(64, 32); f = np.random.randn(64, 32)
print("FVD OK", fvd_features_to_score(r, f))
Problem
Discretise 7-DoF robot actions (xyz + rpy + gripper) into K = 256 bins per dim, encoded as new vocabulary tokens.
import torch
def discretize_action(a, lo=-1.0, hi=1.0, K=256):
return torch.round((a.clamp(lo, hi) - lo) / (hi - lo) * (K - 1)).long()
def undiscretize(t, lo=-1.0, hi=1.0, K=256):
return lo + (hi - lo) * t.float() / (K - 1)
Tests
a = torch.tensor([0.0, 0.5, -1.0])
t = discretize_action(a)
assert torch.allclose(undiscretize(t), a, atol=1.0 / 256)
print("action tok OK")
Problem
Predict the next H actions at once; this enables high-frequency control.
import torch.nn as nn
class ActionChunkHead(nn.Module):
def __init__(self, d, action_dim=7, horizon=8):
super().__init__()
self.head = nn.Linear(d, action_dim * horizon)
self.horizon, self.dim = horizon, action_dim
def forward(self, h):
return self.head(h).view(-1, self.horizon, self.dim)
Tests
m = ActionChunkHead(128, 7, 8); y = m(torch.randn(2, 128))
assert y.shape == (2, 8, 7); print("action chunk OK")
Problem
Predict an action chunk via a small denoiser conditioned on a fused VLM embedding (Pi-0 / Diffusion Policy).
import torch
import torch.nn as nn
class DiffActionPolicy(nn.Module):
def __init__(self, ctx_dim, action_dim=7, horizon=8):
super().__init__()
self.cond_proj = nn.Linear(ctx_dim, 64)
self.net = nn.Sequential(
nn.Linear(action_dim * horizon + 64 + 1, 128), nn.SiLU(),
nn.Linear(128, action_dim * horizon))
self.horizon, self.dim = horizon, action_dim
def forward(self, a_noisy, t, cond):
h = torch.cat([a_noisy.flatten(1), self.cond_proj(cond), t.float().unsqueeze(-1)], dim=-1)
return self.net(h).view(-1, self.horizon, self.dim)
Tests
m = DiffActionPolicy(64, 7, 8)
y = m(torch.randn(2, 8, 7), torch.tensor([5, 10]), torch.randn(2, 64))
assert y.shape == (2, 8, 7); print("diff policy OK")
Problem
Combine token-level cross-entropy on language tokens and either MSE / discrete-CE on action tokens.
import torch.nn.functional as F
def vla_loss(text_logits, text_tgt, act_logits, act_tgt, alpha=0.5):
text_l = F.cross_entropy(text_logits.view(-1, text_logits.size(-1)),
text_tgt.view(-1), ignore_index=-100)
act_l = F.cross_entropy(act_logits.view(-1, act_logits.size(-1)),
act_tgt.view(-1), ignore_index=-100)
return (1 - alpha) * text_l + alpha * act_l
Tests
loss = vla_loss(torch.randn(2, 4, 100), torch.randint(0, 100, (2, 4)),
torch.randn(2, 8, 256), torch.randint(0, 256, (2, 8)))
assert loss > 0; print("vla loss OK")
Problem
Given an observation, predict an action chunk; execute the first K actions before re-planning.
def rollout_step(policy, obs, env, exec_K=4, horizon=8):
chunk = policy(obs) # (1, H, A)
rewards = []
for k in range(min(exec_K, horizon)):
obs, r, done = env.step(chunk[0, k].numpy())
rewards.append(r)
if done: break
return obs, sum(rewards)
Tests
class Env:
def step(self, a): return np.zeros(4), 1.0, False
class P:
def __call__(self, obs): return torch.zeros(1, 8, 4)
o, r = rollout_step(P(), np.zeros(4), Env())
assert r == 4; print("rollout OK")
Problem
For goal-conditioned policies, encode the goal image to a latent and concatenate with the observation embed- ding.
import torch
import torch.nn as nn
class GoalConditioned(nn.Module):
def __init__(self, obs_enc, goal_enc, head):
super().__init__()
self.obs_enc = obs_enc; self.goal_enc = goal_enc; self.head = head
def forward(self, obs, goal):
return self.head(torch.cat([self.obs_enc(obs), self.goal_enc(goal)], dim=-1))
Tests
class E(nn.Module):
def forward(self, x): return x.flatten(1)[:, :8]
class H(nn.Module):
def forward(self, x): return x.sum(dim=-1, keepdim=True)
m = GoalConditioned(E(), E(), H())
out = m(torch.randn(1, 3, 4, 4), torch.randn(1, 3, 4, 4))
assert out.shape == (1, 1); print("goal OK")
Problem
Add a temporal smoothness penalty ∥at+1 −at∥2 to the BC loss to discourage jittery trajectories.
import torch.nn.functional as F
def bc_smooth_loss(pred_chunk, gt_chunk, beta=0.1):
bc = F.mse_loss(pred_chunk, gt_chunk)
smooth = (pred_chunk[:, 1:] - pred_chunk[:, :-1]).pow(2).mean()
return bc + beta * smooth
Tests
loss = bc_smooth_loss(torch.zeros(1, 8, 7), torch.zeros(1, 8, 7))
assert loss.item() == 0; print("BC smooth OK")
Problem
Sample a batch with probability p from robot demonstrations and 1 −p from web text; weight the losses accordingly.
def mixed_batch_loss(text_loss, robot_loss, p_robot=0.3):
return p_robot * robot_loss + (1 - p_robot) * text_loss
Tests
print("mix OK", mixed_batch_loss(torch.tensor(1.0), torch.tensor(2.0)).item())
Problem
Map a frozen vision encoder’s features to the LLM input dimension via a 2-layer MLP.
import torch.nn as nn
class VisionProjector(nn.Module):
def __init__(self, vision_dim, llm_dim):
super().__init__()
self.proj = nn.Sequential(nn.Linear(vision_dim, llm_dim), nn.GELU(),
nn.Linear(llm_dim, llm_dim))
def forward(self, feats): return self.proj(feats)
Tests
m = VisionProjector(768, 4096); y = m(torch.randn(1, 256, 768))
assert y.shape == (1, 256, 4096); print("projector OK")
Problem
RSSM step: deterministic GRU ht+1 = f(ht, zt, at), stochastic posterior zt+1 ∼q(z|ht+1, ot+1), prior zˆt+1 ∼ p(z|ht+1).
import torch
import torch.nn as nn
class RSSM(nn.Module):
def __init__(self, h_dim=64, z_dim=32, a_dim=4):
super().__init__()
self.gru = nn.GRUCell(z_dim + a_dim, h_dim)
self.prior = nn.Linear(h_dim, 2 * z_dim)
self.posterior = nn.Linear(h_dim + 32, 2 * z_dim) # 32 = obs_feat
def step(self, h, z, a, o_feat=None):
h = self.gru(torch.cat([z, a], -1), h)
mu_p, lv_p = self.prior(h).chunk(2, -1)
if o_feat is not None:
mu_q, lv_q = self.posterior(torch.cat([h, o_feat], -1)).chunk(2, -1)
z_next = mu_q + lv_q.exp().sqrt() * torch.randn_like(mu_q)
return h, z_next, (mu_q, lv_q), (mu_p, lv_p)
z_next = mu_p + lv_p.exp().sqrt() * torch.randn_like(mu_p)
return h, z_next, None, (mu_p, lv_p)
Tests
r = RSSM()
h, z, _, prior = r.step(torch.zeros(1, 64), torch.zeros(1, 32), torch.zeros(1, 4))
assert h.shape == (1, 64) and z.shape == (1, 32); print("RSSM OK")
Problem
Roll out the world model for H steps from a starting state, applying a learned policy.
def imagine(rssm, h0, z0, policy, H=10):
h, z = h0, z0; states = []
for _ in range(H):
a = policy(h, z)
h, z, _, _ = rssm.step(h, z, a)
states.append((h, z, a))
return states
Tests
r = RSSM()
def pol(h, z): return torch.zeros(h.size(0), 4)
out = imagine(r, torch.zeros(1, 64), torch.zeros(1, 32), pol, H=5)
assert len(out) == 5; print("imagine OK")
Problem
KL(q(z|o, h)∥p(z|h)) summed over time, free-bits clipped at β.
def rssm_kl(post, prior, free_bits=1.0):
mu_q, lv_q = post; mu_p, lv_p = prior
var_q = lv_q.exp(); var_p = lv_p.exp()
kl = 0.5 * (lv_p - lv_q + (var_q + (mu_q - mu_p) ** 2) / var_p - 1)
kl = kl.sum(dim=-1).clamp(min=free_bits)
return kl.mean()
Tests
mu = torch.zeros(2, 8); lv = torch.zeros(2, 8)
print("RSSM KL OK", rssm_kl((mu, lv), (mu, lv), free_bits=0.0).item())
Problem
Decode latent z to observation oˆ, train via MSE on pixel space (or BCE / Categorical if discrete tokens).
import torch.nn.functional as F
def world_recon_loss(decoder, z, o):
return F.mse_loss(decoder(z), o)
Tests
class D(nn.Module):
def forward(self, z): return torch.zeros(z.size(0), 3, 4, 4)
print("recon OK", world_recon_loss(D(), torch.randn(1, 16), torch.zeros(1, 3, 4, 4)).item())
Problem
A small head that takes (h, z) and predicts the immediate reward.
import torch
import torch.nn as nn
class RewardHead(nn.Module):
def __init__(self, h_dim=64, z_dim=32):
super().__init__()
self.l = nn.Sequential(nn.Linear(h_dim + z_dim, 64), nn.ReLU(), nn.Linear(64, 1))
def forward(self, h, z): return self.l(torch.cat([h, z], -1)).squeeze(-1)
Tests
m = RewardHead(); r = m(torch.randn(2, 64), torch.randn(2, 32))
assert r.shape == (2,); print("reward head OK")
Problem
On rolled-out states, compute discounted λ-returns and update a value head + policy head (Dreamer-V3 style).
import torch
import torch.nn.functional as F
def lambda_return(rewards, values, gamma=0.99, lam=0.95, last_value=0.0):
out = [0.0] * len(rewards); g = last_value
for t in reversed(range(len(rewards))):
g = rewards[t] + gamma * ((1 - lam) * values[t + 1 if t + 1 < len(values) else -1] + lam * g)
out[t] = g
return out
def actor_critic_imagined(rewards, values, log_probs, ent, beta=1e-3):
R = torch.tensor(lambda_return(rewards, values))
actor_loss = -(log_probs * (R - torch.tensor(values)).detach()).mean() - beta * ent
value_loss = F.mse_loss(torch.tensor(values), R.detach())
return actor_loss + value_loss
Tests
out = lambda_return([1, 1, 1], [0, 0, 0], gamma=1.0, lam=1.0)
assert out == [3, 2, 1]; print("lambda return OK")
Problem
xt+1 = f(x≤t, a≤t) where f is a video diffusion / Transformer model conditioned on per-frame actions (Genie / GAIA-2 style).
import torch.nn as nn
class ActionConditionedWM(nn.Module):
def __init__(self, dim=64):
super().__init__()
self.frame_enc = nn.Conv2d(3, dim, 3, padding=1)
self.action_emb = nn.Embedding(16, dim)
self.pred = nn.Conv2d(dim, 3, 3, padding=1)
def forward(self, x_prev, action_id):
h = self.frame_enc(x_prev)
a = self.action_emb(action_id).view(-1, h.size(1), 1, 1)
return self.pred(h + a)
Tests
m = ActionConditionedWM()
y = m(torch.randn(1, 3, 16, 16), torch.tensor([3]))
assert y.shape == (1, 3, 16, 16); print("AC WM OK")
Problem
Predict the latent of a future patch from a context patch via a small predictor; train with MSE in latent space (no pixel decoding).
import torch.nn as nn
import torch.nn.functional as F
class VJEPA(nn.Module):
def __init__(self, d=64):
super().__init__()
self.enc = nn.Sequential(nn.Linear(d, d), nn.GELU(), nn.Linear(d, d))
self.pred = nn.Sequential(nn.Linear(d, d), nn.GELU(), nn.Linear(d, d))
def forward(self, ctx, tgt):
z_ctx = self.enc(ctx)
z_tgt = self.enc(tgt).detach() # EMA target in real V-JEPA
return F.mse_loss(self.pred(z_ctx), z_tgt)
Tests
m = VJEPA(); loss = m(torch.randn(2, 4, 64), torch.randn(2, 4, 64))
assert loss > 0; print("V-JEPA OK")
Problem
Tokenize each frame to a sequence of integers; train a Transformer to predict the next frame’s tokens given previous frames + actions.
import torch.nn as nn
class TokenWM(nn.Module):
def __init__(self, vocab=1024, d=128, n_blocks=2):
super().__init__()
self.emb = nn.Embedding(vocab + 16, d) # extra slots for action tokens
layer = nn.TransformerEncoderLayer(d, 4, batch_first=True)
self.tr = nn.TransformerEncoder(layer, n_blocks)
self.head = nn.Linear(d, vocab)
def forward(self, ids):
h = self.tr(self.emb(ids))
return self.head(h)
Tests
m = TokenWM()
y = m(torch.randint(0, 1024, (1, 64)))
assert y.shape == (1, 64, 1024); print("token WM OK")
Problem
Train a discrete action codebook by reconstructing frame transitions: minimise ∥D(zt, at) −xt+1∥where at is a discrete latent.
import torch
import torch.nn as nn
import torch.nn.functional as F
class LatentActions(nn.Module):
def __init__(self, n_actions=8, d=64):
super().__init__()
self.actions = nn.Parameter(torch.randn(n_actions, d))
self.encoder = nn.Linear(d * 2, d)
self.predictor = nn.Linear(d * 2, d)
def forward(self, x_t, x_t1):
# encode the transition to a continuous embedding, snap to nearest codebook entry
emb = self.encoder(torch.cat([x_t, x_t1], -1))
d = ((emb.unsqueeze(1) - self.actions.unsqueeze(0)) ** 2).sum(-1)
idx = d.argmin(dim=1)
a = self.actions[idx]
pred = self.predictor(torch.cat([x_t, a], -1))
return F.mse_loss(pred, x_t1), idx
Tests
m = LatentActions()
loss, idx = m(torch.randn(4, 64), torch.randn(4, 64))
assert loss >= 0 and idx.shape == (4,); print("latent action OK")
Problem
Plan in the world model: sample N action sequences, simulate, pick the highest-return sequence.
import torch
def cem_in_world(rssm, reward_head, h0, z0, action_dim, H=10, N=200, elite=0.1, iters=3):
mu = torch.zeros(H, action_dim); sigma = torch.ones_like(mu)
for _ in range(iters):
actions = mu + sigma * torch.randn(N, H, action_dim)
rewards = torch.zeros(N)
h, z = h0.expand(N, -1).clone(), z0.expand(N, -1).clone()
for t in range(H):
h, z, _, _ = rssm.step(h, z, actions[:, t])
rewards += reward_head(h, z)
idx = rewards.argsort(descending=True)[:int(N * elite)]
elites = actions[idx]
mu = elites.mean(0); sigma = elites.std(0) + 1e-6
return mu[0]
Tests
r = RSSM(); rh = RewardHead()
a = cem_in_world(r, rh, torch.zeros(1, 64), torch.zeros(1, 32), 4, H=3, N=8, iters=1)
assert a.shape == (4,); print("CEM in world OK")
Advanced implementation. Every CEM iteration should evaluate the entire population in ONE batched objective call (batched world-model rollout) — the loop-over-candidates version wastes exactly the parallelism a learned model provides. Verified: converges to a known optimum within 0.05.
import torch
def cem_batched(objective, dim, iters=15, pop=256, elite=32, seed=0):
g = torch.Generator().manual_seed(seed)
mu, sd = torch.zeros(dim), torch.ones(dim)
for _ in range(iters):
cand = mu + sd * torch.randn(pop, dim, generator=g)
scores = objective(cand) # (pop,) one vectorized rollout
top = scores.topk(elite).indices
mu, sd = cand[top].mean(0), cand[top].std(0) + 1e-6
return mu
Problem
Concatenate two text encoders’ outputs (CLIP penultimate + T5-XXL final) along the sequence axis, project into the diffusion model’s conditioning dim.
import torch
import torch.nn as nn
class DualTextEncoder(nn.Module):
def __init__(self, clip_dim, t5_dim, out_dim):
super().__init__()
self.clip_proj = nn.Linear(clip_dim, out_dim)
self.t5_proj = nn.Linear(t5_dim, out_dim)
def forward(self, clip_feat, t5_feat):
return torch.cat([self.clip_proj(clip_feat), self.t5_proj(t5_feat)], dim=1)
Tests
m = DualTextEncoder(768, 4096, 1024)
y = m(torch.randn(1, 77, 768), torch.randn(1, 77, 4096))
assert y.shape == (1, 154, 1024); print("dual text OK")
Problem
adaLN(x, c) = (1 + γ(c)) LN(x) + β(c) where c fuses time embedding and class / pooled-text embedding.
import torch
import torch.nn as nn
import torch.nn.functional as F
class AdaLNCond(nn.Module):
def __init__(self, d, t_dim, c_dim):
super().__init__()
self.proj = nn.Linear(t_dim + c_dim, 2 * d)
def forward(self, x, t, c):
ab = self.proj(torch.cat([t, c], dim=-1)); g, b = ab.chunk(2, -1)
return (1 + g.unsqueeze(1)) * F.layer_norm(x, x.shape[-1:]) + b.unsqueeze(1)
Tests
m = AdaLNCond(64, 32, 32); y = m(torch.randn(2, 16, 64), torch.randn(2, 32), torch.randn(2, 32))
assert y.shape == (2, 16, 64); print("adaLN-cond OK")
Problem
Condition the video diffusion model on the first frame by concatenating its latent along the channel axis with the noisy latent at every timestep.
import torch
def i2v_input(noisy_lat, first_frame_lat):
f = first_frame_lat.unsqueeze(2).expand(-1, -1, noisy_lat.size(2), -1, -1)
return torch.cat([noisy_lat, f], dim=1)
Tests
out = i2v_input(torch.zeros(1, 4, 8, 8, 8), torch.zeros(1, 4, 8, 8))
assert out.shape == (1, 8, 8, 8, 8); print("I2V OK")
Notes
Survival tactics for live video / VLA / world-model coding:
State the architecture style: spacetime DiT (Sora), MM-DiT (SD3 / Flux), factorised spatial+temporal, hierarchical (cascaded keyframes + interpolation), or autoregressive over visual tokens (Genie / IRIS).
Video VAE: causal 3D conv, separable spatial/temporal downsampling, reconstruction MSE + KL + per- ceptual + GAN losses. Common compression 4× in time, 8× in space.
Tokenization: continuous (FSQ / VQ) for video diffusion conditioning vs. discrete token-Transformer world models. State the trade-off.
Conditioning: CFG (text), multi-CFG (text + image), AdaLN-Zero is the standard. Drop the conditioning during training with p ∈ [0.05, 0.2] to enable CFG at inference.
VLA: distinguish action tokenization (RT-2: discrete bins as new vocabulary tokens), action chunking (Pi-0: predict H actions), and diffusion-policy heads (continuous actions modelled by a small denoiser).
Co-training: web text + robot data. The mixing ratio matters; too much text and the action head atrophies, too little and language grounding collapses.
World model: RSSM (deterministic + stochastic latent), discrete token Transformer (Genie), pixel video diffusion (GAIA-2 / Cosmos), or JEPA-style latent prediction (V-JEPA, no decoder). Each is best for a different downstream use.
Imagination training: roll out the world model H steps, train policy and value on the imagined returns. Free-bits / KL clipping is essential.
Always end with a sanity check: video shapes (B, C, T, H, W), that VAE encode → decode preserves shape, that CFG with w = 0 recovers the conditional model, and that an action token round-trips through encode/decode within bin width.