Genie and Interactive Generated Environments › Latent Action Discovery

Standard action-conditioned world models, the kind we studied in Recurrent State-Space Models and MuZero and Value-Equivalent Models, assume a dataset of tuples $(o_t, a_t, o_{t+1})$ where the action $a_t$ is known, with

• $o_t$: the observation at time $t$
• $a_t$: the action taken at time $t$
• $o_{t+1}$: the resulting observation at time $t+1$

In an Atari replay buffer, $a_t$ is a recorded button press; in a robot demonstration it might be a recorded control command. Ordinary uploaded video, including gameplay clips, often lacks that paired control stream. Original Genie trained on such unlabeled gameplay video. One option is to restrict training to action-labeled recordings; another is to infer a compact action-like variable from the observed transitions.

The usual supervised recipe gives the model a state, a recorded action, and the resulting state, then minimizes prediction error. Without an action record, an unconditional video model cannot offer that same user-controlled branch. Hand-labeling some clips is possible but adds annotation work and may cover only selected behaviors. Genie asks what action-like structure can be inferred from the video itself.

Genie's answer to this problem is to introduce a latent action. Between two consecutive frames of a video, there is a change. Whatever caused that change is, in a loose sense, an "action." Even in a nature documentary, the next frame is different from the current one because something happened: the camera panned, a bird flapped its wings, a leaf fell. The latent action is a discrete code that summarizes that change. Crucially, it is not required to correspond to any human-meaningful command. It just has to be a useful, compact summary of what moved. The word "action" here is a modeling convenience, not a claim about agency. The code does not need to know who acted or why. It only needs to help predict what happens next.

The simplest abstraction is an inverse model that receives two frames and returns a code from a small discrete codebook. In the original Genie, the latent-action model's own reconstruction decoder tests whether its code preserves information needed to reconstruct the next frame; the dynamics model consumes that code without its loss updating the latent-action encoder through the action input. The two are co-trained after tokenizer pretraining, with separate objectives. For our pair-based toy abstraction, write the code distribution as

where:

• $u_t$: the latent action that summarizes the change between frame $t$ and frame $t+1$
• $o_t$: the observation (frame) at time $t$
• $o_{t+1}$: the observation (frame) at time $t+1$
• $q_\phi(\cdot \mid \cdot)$: a learned distribution over codes $u_t \in \{1, \dots, K\}$
• $K$: the cardinality of the latent action codebook; Genie uses $K = 8$ in the original report, but the exact number is a hyperparameter

The model is called "inverse" because it runs the prediction question backward: given the earlier and later frames, what action-like code describes their change? Recorded actions remove the need to discover such labels for training, though predicting an action from images can still be ambiguous. Genie has no recorded actions, so its latent-action model supplies the interface. Its reconstruction objective pressures the code to preserve next-frame information, and a small codebook limits the vocabulary a user has to explore.

The forward dynamics model predicts next-frame visual tokens given past tokens and a latent action. The original Genie first trains its video tokenizer, then co-trains the latent-action model and dynamics model; it stops the dynamics loss gradient at the latent-action input. The latent-action model itself has an encoder, quantized codebook, and reconstruction decoder. At inference the user selects a code, and the latent-action encoder and its training decoder are not needed. In this chapter we use $o_t$ for toy observations and $z_t$ for full-model visual tokens:

where:

• $z_t$: the visual tokens for frame $t$ (for example, patches or discrete codes produced by a tokenizer)
• $z_{\leq t}$: the sequence of visual tokens from frame $1$ through frame $t$, the conditioning context
• $u_{\leq t}$: the sequence of latent actions from transition $1$ through transition $t$
• $p_\theta(\cdot \mid \cdot)$: a learned distribution over the next frame's visual tokens

The latent-action model encodes change and reconstructs the future from that code and the past; this reconstruction objective makes the code informative. The dynamics model then learns to predict video tokens using the resulting codes. The original method does not simply backpropagate the dynamics loss through the latent-action input, so "co-trained" should not be mistaken for a single end-to-end inverse/forward loss. This separation also matters when we later compare the published model with our hand-engineered toy.

In the original method, the latent-action reconstruction objective learns an informative quantized code, while a separate cross-entropy objective trains the masked-token dynamics model on video tokens and stop-gradient latent actions. The tokenizer is pretrained before this stage. This is a staged training procedure, not one end-to-end optimization of the simplified objective below.

Codebook collapse remains a risk in quantized models: if all transitions use one code, the dynamics model has no varied action signal. Genie's reported latent-action training uses a VQ-VAE-style objective and a small codebook, but the paper does not establish that a separate code-usage-balancing term guarantees every code will be used. Code utilization and action responsiveness therefore need empirical checks.

As a simplified next-frame abstraction, a conditional token-prediction objective is the negative log-probability of observed next-frame tokens given prior tokens and actions:

where:

• $\mathcal{L}_{\text{forward}}(\theta)$: the training loss for the forward model's parameters $\theta$, written as the negative log-likelihood summed across all frame transitions $t$
• $\log p_\theta(z_{t+1} \mid z_{\leq t},\, u_{\leq t})$: the log-probability that the model assigns to the actual next-frame tokens $z_{t+1}$ given the earlier tokens and latent actions
• the leading minus sign: converts the log-likelihood into a loss, so a lower value means the model predicts the observed future better

The original latent-action model has a different training signal: its decoder reconstructs the next frame from the past and quantized inferred action, with the VQ-VAE objective training the encoder and codebook. Maximizing $q_\phi$ on codes chosen by that same $q_\phi$, without reconstruction or another grounding term, would not make the codes useful. The forward expression above is likewise a teaching abstraction; Genie's dynamics model predicts masked frame tokens with cross-entropy, then samples each frame in multiple MaskGIT refinement steps. Our toy is more distant still: its $K$-means codes come from hand-extracted displacements rather than a learned reconstruction model.

Why learn the action code instead of using optical flow or a hand-designed state? Flow describes motion but is not necessarily a compact control interface, while a handcrafted variable for one game may not transfer to another. The latent-action bottleneck restricts the code to a small vocabulary of changes. It makes that vocabulary available to the dynamics model, but does not force the dynamics model to use it: an action-insensitive predictor could still rely on visual history. We test responsiveness separately for precisely that reason.

The first thing to internalize about latent actions is that they have a permutation ambiguity. A code is a compressed summary of change, not a semantic label. Suppose you have trained both models jointly, and everything works. Now swap the labels of code 1 and code 2 in the codebook. The forward model never saw the labels, only the embedding vectors, so as long as you also swap the corresponding embedding rows, the loss is unchanged. The trained system is just as good as before. But the meaning you attach to "code 1" and "code 2" has changed. If you had been using code 1 to mean "left arrow," you now have to use code 2. The labels are an arbitrary naming convention, not a property of the learned model.

The permutation ambiguity is not a rare corner case. It is a structural property of a discrete latent codebook without label supervision. The assignment of integer names to learned codes is arbitrary: simultaneously permuting the encoder's assignments and the dynamics model's code embeddings leaves a given model's predictions unchanged. Retraining with another seed may change the code names, but can also change predictions for unrelated optimization reasons. Evaluation should therefore ask whether codes can be aligned with useful effects, not whether code 0 has a prescribed direction. As our toy shows, even the existence of a clean one-to-one alignment can fail when different button presses produce the same observed transition.

The second thing to internalize is that latent actions are not automatically identifiable physical controls. In a grid world with four buttons, one might hope for four corresponding codes, but blocked moves create a fifth observable outcome (no displacement) and obscure the button that caused it. In a cooking video, codes might summarize camera movement, hand movement, or both. The training objective seeks useful compression of change, not a guarantee of human button semantics.

This point is easy to underestimate because the grid-world example is so clean. In the grid world, the change between frames is exactly the agent's movement, and movement has four discrete directions. In real video, the change between frames can come from camera motion, object motion, lighting changes, occlusions, and cuts between scenes. The inverse model may allocate codes to any of these. A code might mean "the camera panned left" in one clip and "the player moved left" in another, if those two changes happen to be similar in token space. The model has no built-in notion of cause, intention, or control. It has a notion of predictive usefulness. That is enough to make a controllable latent space in favorable cases, but it is not enough to guarantee a human-interpretable action space.

The third thing to internalize is that a latent action does not become a keyboard command or a portable interface until you explicitly align it with one. If you want a human to press the left arrow and see the character move left, you need one of the following:

• A small set of labeled interaction data that pairs latent codes with real commands.
• A user interface that lets the user pick a code and remembers the mapping.
• An auxiliary loss that pushes codes toward known controls.

Without such alignment, a person or outside policy can still select the opaque code IDs and learn their effects. What they cannot assume is that one ID already means "left arrow" or that its meaning transfers to another trained model. That is the difference between a model that has learned a control abstraction and one that has learned your control abstraction. The alignment step can be small, but it is separate from latent action discovery: discovery produces a vocabulary of change, while alignment maps your commands into that vocabulary.

Let us make this concrete with a tiny grid world. We will generate trajectories from a known environment, throw away the action labels, and see what latent actions the model discovers.

We build an 8 by 8 grid with a handful of walls, four actions (up, down, left, right), and deterministic transitions. If the agent attempts to move into a wall or off the grid, the move is a no-op and the agent stays in place. This kind of blocked move will matter later, because it creates pairs of actions with the same observable effect.

We now generate a dataset of random trajectories. Each trajectory starts from a random free cell, and at each step we sample one of the four actions uniformly at random. We collect the current frame, the true action (which we will hide during latent-action inference), and the next frame.

First split by whole trajectory: 160 trajectories train the clusterer and predictors, and 40 remain for one-step testing. Adjacent frames from one trajectory therefore cannot appear on both sides of that split. From each frame pair we extract agent displacement. K-means fits four clusters on training displacements only; it assigns codes to the held-out transitions afterward. This is a hand-engineered, hindsight-based stand-in for $q_\phi(u_t \mid o_t,o_{t+1})$: assigning a test code requires the test next frame, so the held-out latent-code score below is a reconstruction diagnostic, not a deployable next-frame forecast.

The use of $K$-means here deserves a word of explanation. The published latent-action model learns an encoder, codebook, and reconstruction decoder; our toy replaces that problem with clustering. Displacement is one of five possible vectors: up, down, left, right, or zero for a blocked move. We ask for four clusters, so zero displacement must share a code with some movement case. The clustering step never sees button labels, but our hand-designed displacement extractor already builds in knowledge of agent position and motion. Its centers are movement or mixed movement/no-op prototypes, not proof that a generic inverse model would discover the same actions. K-means does not refine its codes using a learned reconstruction or next-frame prediction objective.

The full-data confusion matrix below is descriptive: its held-out labels do not fit K-means or calibrate the lookup. Code labels can be permuted without changing the clustering. Blocked button presses all have zero displacement, so one code can mix those moves with a direction.

Notice that even when the pattern is clean, the naming of the codes is arbitrary. Nothing in the clustering objective forces code 0 to mean "up" rather than "down." Our action_perm lookup uses only training action labels to assign one code to each button by column majority. It happens to be one-to-one in this seeded run, but that rule does not guarantee a permutation in other datasets. This supervised calibration supplies extra information for prospective tests; it is not part of label-free code discovery.

The confusion matrix is also useful because it exposes the no-op wrinkle directly. If many zero-delta transitions are assigned to a single latent code, that code may look like a direction in the confusion matrix while actually absorbing all the blocked moves. In a larger model, this could mean that one latent action becomes a "nothing happened" code, which is useful but not a direction. That is not a failure. It is the inverse model doing its job: summarizing change, including the absence of change. The evaluation question is whether the codes are useful for prediction and control, not whether they line up one-to-one with human button names.