Genie and Foundation World Models

A foundation world model is a video model that has been forced to explain each frame as the result of a choice.

Video is the largest source of experience about how the world behaves, and almost none of it comes with actions attached. A clip of a platformer game shows the character jumping; it does not record the button. Genie (Bruce et al., 2024, arXiv:2402.15391) showed that an interactive, controllable environment can be learned from such footage alone, by inferring a small discrete action for every transition and then training a dynamics model conditioned on it. This note rebuilds that recipe one component at a time. The mechanics of vector quantization and the small-scale version of a latent action model are in the latent action; this post is about what changes at scale and what the components look like in code.

The simplest version: passive video prediction#

Start with a model that predicts the next frame from the previous ones.

def passive_rollout(predict_next, frames, steps):
    """A video model with no actions: the future is whatever the data usually did."""
    for _ in range(steps):
        frames.append(predict_next(frames))
    return frames

This can be trained on any video, and it can produce plausible continuations. It is not a world model in the useful sense, because nothing in its interface lets an agent or a player choose. The model averages over the choices the people in the data made. To plan, or to play, the future has to be conditioned on an action, and the data has none.

Three problems follow, and Genie has one component for each:

All three share one building block.

The building block: a spatiotemporal transformer#

Full attention over every token of every frame costs $O((T \cdot S)^2)$ for $T$ frames of $S$ tokens. Genie's ST-transformer factorizes it: a spatial layer attends over the $S$ tokens within one frame, then a causal temporal layer attends over the same token position across $T$ frames. The cost becomes $O(T S^2 + S T^2)$, linear in the number of frames for the spatial part.

import torch, torch.nn as nn, torch.nn.functional as F
 
class STBlock(nn.Module):
    def __init__(self, d, heads):
        super().__init__()
        self.n1, self.n2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.spatial = nn.MultiheadAttention(d, heads, batch_first=True)
        self.temporal = nn.MultiheadAttention(d, heads, batch_first=True)
        self.ff = nn.Sequential(nn.LayerNorm(d), nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))
 
    def forward(self, x):                                   # x: (B, T, S, d)
        B, T, S, d = x.shape
        h = self.n1(x).reshape(B * T, S, d)                 # each frame on its own
        x = x + self.spatial(h, h, h, need_weights=False)[0].reshape(B, T, S, d)
        h = self.n2(x).transpose(1, 2).reshape(B * S, T, d) # each position across time
        future = torch.triu(torch.ones(T, T, dtype=torch.bool, device=x.device), 1)
        h = self.temporal(h, h, h, attn_mask=future, need_weights=False)[0]
        x = x + h.reshape(B, S, T, d).transpose(1, 2)
        return x + self.ff(x)

The causal mask matters twice: it lets the dynamics model be trained on all frames of a clip in parallel, and it is what lets the latent action model's output at step $t+1$ see the frame $x_{t+1}$ while the output at step $t$ cannot.

Step 1: a video tokenizer that sees time#

The tokenizer is a VQ-VAE whose encoder and decoder are stacks of STBlocks. The paper calls it ST-ViViT and uses 1,024 codes with a patch size of 4. Because its encoder attends causally over time, the token for frame $t$ can depend on frames before it, which makes tokens temporally consistent.

The paper's tokenizer ablation reports FVD of 114.5 for a per-frame ViT tokenizer, 272.7 for C-ViViT, and 81.4 for ST-ViViT (lower is better). Seeing time helps, and the factorized form did best.

Bottleneck after Step 1: the video is now a sequence of token grids, but the transitions between them are still unexplained.

Step 2: a latent action model#

The latent action model (LAM) takes the frames $x_{1:t}$ and the next frame $x_{t+1}$, and must produce a code $a_t$ such that a decoder can reconstruct $x_{t+1}$ from $x_{1:t}$ and $a_t$ alone. The code is quantized to a vocabulary of only $|A| = 8$. That tiny bottleneck is the whole trick: eight codes cannot encode the next frame, so they are pushed to encode the change that the past frames could not predict, which in a game is mostly the player's input.

class LatentActionModel(nn.Module):
    """Sketch. Encoder sees x_1..x_{t+1}; decoder must rebuild x_{t+1} from x_1..x_t and a_t."""
    def __init__(self, embed, encoder, decoder, d=512, d_a=32, n_actions=8):
        super().__init__()
        self.embed, self.encoder, self.decoder = embed, encoder, decoder   # patchify, STBlocks, STBlocks
        self.to_action = nn.Linear(d, d_a)
        self.codebook = nn.Embedding(n_actions, d_a)
 
    def infer(self, frames):                                # frames: (B, T+1, C, H, W) pixels
        h = self.encoder(self.embed(frames))                # (B, T+1, S, d), causal in time
        z = self.to_action(h[:, 1:].mean(dim=2))            # one vector per transition: (B, T, d_a)
        idx = torch.cdist(z, self.codebook.weight).argmin(-1)
        q = self.codebook(idx)
        return z + (q - z).detach(), idx, z, q              # straight-through estimator
 
    def loss(self, frames, beta=0.25):
        a, idx, z, q = self.infer(frames)
        pred = self.decoder(frames[:, :-1], a)              # predict x_2..x_{T+1}
        recon = F.mse_loss(pred, frames[:, 1:])
        vq = F.mse_loss(q, z.detach()) + beta * F.mse_loss(z, q.detach())
        return recon + vq, idx

Two details from the paper are worth keeping. The LAM reads pixels, not tokenizer output: feeding it tokens lowered controllability in their ablation, presumably because quantized tokens throw away small motions that reveal the action. And the LAM is used only for training. At inference, the encoder and decoder are discarded; only the 8-entry codebook survives, so that a player's button press (an integer from 0 to 7) can be looked up as an action embedding.

Step 3: a masked dynamics transformer#

The dynamics model is a stack of STBlocks over token grids. It receives the tokens $z_{1:t}$ and the latent actions $a_{1:t-1}$ (added to the frame embeddings, not concatenated), and predicts the tokens of the next frame.

Training is masked token prediction in the style of MaskGIT: replace a random fraction of tokens with a mask token and predict the originals. Genie samples the mask rate uniformly in $[0.5, 1]$, so the model learns everything from "fill in a few gaps" to "generate the frame from nothing but the past and the action".

def dynamics_loss(dyn, tokens, action_idx, action_codebook, mask_id):
    """tokens: (B, T, S) ints from the frozen tokenizer; action_idx: (B, T-1) ints from the LAM."""
    B, T, S = tokens.shape
    rate = torch.empty(B, 1, 1, device=tokens.device).uniform_(0.5, 1.0)
    mask = torch.rand(B, T, S, device=tokens.device) < rate
    mask[:, 0] = False                                   # the first frame is always context
    inp = tokens.masked_fill(mask, mask_id)
    act = action_codebook(action_idx)                    # (B, T-1, d_a), from the LAM's codebook
    logits = dyn(inp, act)                               # (B, T, S, vocab); a_{t-1} added at frame t
    return F.cross_entropy(logits[mask], tokens[mask])

Step 4: decoding a frame in a few parallel passes#

Autoregressive decoding would need one forward pass per token: about 900 passes for a 160×90 frame at patch size 4. MaskGIT decoding fills the frame in a fixed number of passes. Start fully masked; at each pass, sample every masked token, keep the most confident ones, and re-mask the rest according to a cosine schedule.

import numpy as np
MASK = -1
 
def maskgit_decode(predict_logits, n_tokens, steps=8, temperature=1.0, rng=np.random.default_rng(0)):
    """predict_logits(tokens) -> (n_tokens, vocab) logits for a partially filled frame."""
    tokens = np.full(n_tokens, MASK)
    for step in range(steps):
        logits = predict_logits(tokens) / temperature
        probs = np.exp(logits - logits.max(-1, keepdims=True)); probs /= probs.sum(-1, keepdims=True)
        sampled = np.array([rng.choice(len(p), p=p) for p in probs])
        conf = probs[np.arange(n_tokens), sampled]
        conf[tokens != MASK] = np.inf                    # committed tokens stay committed
        still_masked = int(np.floor(n_tokens * np.cos(np.pi / 2 * (step + 1) / steps)))
        new = np.where(tokens == MASK, sampled, tokens)
        new[np.argsort(conf)[:still_masked]] = MASK      # re-mask the least confident
        tokens = new
    return tokens

For 64 tokens and 8 steps, the schedule leaves 62, 59, 53, 45, 35, 24, 12 and finally 0 tokens masked: very few commitments early, when the model knows least, and many at the end. Genie uses 25 steps per frame at temperature 2.

Put together, interactive generation is a loop: the player supplies an integer, the codebook turns it into an embedding, and the dynamics model decodes the next frame from the last 16 frames of tokens.

def play(tokenizer, dyn, action_codebook, prompt_frame, get_action, steps=100, context=16):
    tokens, actions = [tokenizer.encode(prompt_frame)], []
    for _ in range(steps):
        actions.append(get_action())                          # an integer in 0..7
        z, a = tokens[-context:], action_codebook(actions[-(context - 1):])
        nxt = maskgit_decode(lambda partial: dyn.next_frame_logits(z, a, partial),
                             n_tokens=z[0].size, steps=25, temperature=2.0)
        tokens.append(nxt)
        yield tokenizer.decode(nxt)

How controllability is measured#

A model can produce beautiful video while ignoring the action. Genie's $\Delta_t$PSNR checks this directly: generate frame $t$ once with the latent actions inferred from the real clip and once with random latent actions, and compare both with the real frame.

A large gap means the actions carry information about the future. It is one instance of the action-sensitivity layer in how to evaluate a world model.

What changes at scale#

The paper trained on about 30,000 hours of 2D platformer gameplay (6.8M clips of 16 seconds at 10 FPS, 160×90), with a final model of about 11B parameters, and ran a scaling study from 40M to 2.7B. It also trained a 2.5B model on robot videos without using their action labels, reaching an FVD of 82.7 on the test split. The same recipe has been carried to other domains. GAIA-1 (2023) applied action- and text-conditioned video world models to driving, and Genie 2 (2024) extended interactive generation to 3D environments. In surgery, where action annotation is expensive and footage is plentiful, the Surgical Vision World Model (Koju et al., 2025, arXiv:2503.02904) adapts Genie-style latent actions to unlabeled surgical video from the SurgToolLoc-2022 dataset.

What these models still get wrong#

Object permanence. An object that leaves the context window, or is occluded for long enough, can reappear changed or not at all. The dynamics model attends to 16 frames; anything older exists only if it is still visible.

Physics drift. Each frame is sampled, and errors are fed back as context. Gravity weakens, a character clips through a platform, and the next frame treats the error as the new truth. Plausible frames are not the same as consistent physics.

Long horizons. Short memory and compounding error together mean quality degrades over long rollouts. The paper lists this limitation directly, along with hallucinated, unrealistic futures.

Action semantics. The 8 codes mean whatever the training data made them mean. A code that is "jump" in one game may be "move right and up" in another, and nothing guarantees that a code means the same thing in every context. Codes are also capped at eight discrete options, which is far from the continuous control a robot needs.

Speed. The paper reports roughly 1 FPS generation, too slow for real-time interaction without further work.

What to try next#

1. Train a LAM on a small game with known actions, then measure code-to-action purity with $|A| \in \{4, 8, 16, 64\}$, and watch purity fall as the bottleneck loosens.
2. Compute $\Delta_t$PSNR for your own model at horizons 1, 4 and 16 frames.
3. Run an occlusion probe: generate with an object hidden for $k$ frames, sweep $k$ past the context length, and track whether it returns.
4. Sweep MaskGIT steps (4, 8, 25) and temperature; record FVD and $\Delta_t$PSNR together, since sharper samples and more controllable samples are not always the same setting.
5. Swap the LAM's pixel input for tokenizer input and confirm the controllability drop the paper reports.
6. Test whether a latent action code transfers across scenes: apply the same code in ten different contexts and check by eye whether the motion is the same.