Tokenized Game and Environment Models › Action-conditioned autoregressive dynamics

With tokens in hand, the next step is to model how they change. An autoregressive sequence model conditioned on the action history is one common choice. A frame contains $L$ tokens plus reward and termination variables; a naive categorical table over all token sequences alone would have $K^L$ entries. Autoregression replaces that table with a chain of per-token conditionals. Other structured model families are possible, but this factorization gives us a precise example to analyze.

Here is one factorization, which we will use as a running example. It decomposes the joint prediction of the next frame's tokens, reward, and termination into a product over tokens within the frame, multiplied by a single conditional for the reward and termination given the completed frame:

$$p_\theta(z_{t+1}, r_{t+1}, d_{t+1} \mid h_t, a_t) = \left[\prod_{i=1}^{L} p_\theta(z_{t+1,i} \mid h_t, a_t, z_{t+1,<i})\right] \, p_\theta(r_{t+1}, d_{t+1} \mid h_t, a_t, z_{t+1}),$$

where:

• $p_\theta(z_{t+1}, r_{t+1}, d_{t+1} \mid h_t, a_t)$: the joint distribution the model assigns to the next frame's tokens, reward, and termination flag, conditioned on the history and the current action
• $\prod_{i=1}^{L} p_\theta(z_{t+1,i} \mid h_t, a_t, z_{t+1,<i})$: the product of per-token conditionals that generates the next frame token by token, each token conditioned on the history, the action, and the tokens already emitted for that frame

Continuing the list of terms:

• $h_t$: a bounded history of previously encoded observations and actions. Depending on the model, $h_t$ might be $\{z_{t-k}, a_{t-k}, \ldots, z_t\}$ for some window $k$, or it might be a recurrent summary.
• $a_t$: the current action.
• $z_{t+1,<i}$: the tokens of the next frame that have already been generated in the autoregressive order.
• $r_{t+1}$ and $d_{t+1}$: the reward and termination flag for the transition.

Every term in this factorization deserves a comment.

The product $\prod_i p_\theta(z_{t+1,i} \mid h_t, a_t, z_{t+1,<i})$ means the model generates the next frame token by token, conditioning each prediction on the tokens already emitted for that frame. The order is a design choice: a spatial grid can be flattened in raster order, for example. The chain rule permits any fixed order to represent the same joint distribution if the conditionals are sufficiently expressive. With finite data and model capacity, however, the order can affect both training and generated behavior.

The final factor $p_\theta(r_{t+1}, d_{t+1} \mid h_t, a_t, z_{t+1})$ asks for a reward and termination flag given the generated next frame. Unlike the per-token factors, it does not need a partial next-frame prefix: all frame tokens are available. Implementations may use separate heads with their own classification or regression losses. Modeling these variables matters because a visually plausible frame with the wrong reward or termination can mislead a planner.

Training an autoregressive token model is a supervised learning problem once the tokenizer is fixed. You take a batch of trajectories, convert each observation to a token sequence, and ask the model to predict the next token sequence given the previous tokens and the observed action.

Concretely, for a training pair $(h_t, a_t, z_{t+1})$, the loss is the sum of cross-entropies over the $L$ tokens:

$$\mathcal{L}_{\text{token}}(\theta) = -\sum_{i=1}^{L} \log p_\theta(z_{t+1,i} \mid h_t, a_t, z_{t+1,<i}),$$

where:

• $\mathcal{L}_{\text{token}}(\theta)$: the token-prediction loss as a function of the model parameters $\theta$
• $p_\theta(z_{t+1,i} \mid h_t, a_t, z_{t+1,<i})$: the model's predicted probability of the $i$-th token of the next frame given the history, the action, and the already-emitted earlier tokens
• $L$: the number of tokens per frame

plus analogous loss terms for $r_{t+1}$ and $d_{t+1}$. A continuous reward may use a likelihood model or a regression loss; a small discrete reward set may use cross-entropy. Token cross-entropy is the negative log-likelihood of the observed categorical target under the model. It uses ground-truth categorical targets to train the model's conditional probabilities, but it does not eliminate the train–rollout distribution shift or make greedy, temperature-adjusted, or truncated sampling match the training distribution.

The tokens used as inputs are the actual tokens from the recorded trajectory, not the model's own samples. This is teacher forcing: we condition on ground-truth history and predict one step ahead. It is a common training regime for autoregressive models and it is what Part V, Ch 3: Autoregressive Token World Models described. Two implications matter for evaluation:

1. It trains the model to be accurate on data distribution histories, not on its own histories. During a planning rollout, small errors can compound and move the model's inputs away from the training distribution.
2. It supplies ground-truth $h_t$ during training. A model may exploit cues in that history that its own rollouts later fail to preserve; teacher forcing does not itself require such reliance, and $h_t,a_t$ need not determine a stochastic next frame.

The second point is subtle. Suppose the next frame depends on whether the agent holds the key, and that fact is visible in $h_t$ because the key sprite has disappeared. The model may predict observed next frames accurately by using that cue. Such accuracy alone does not show that its own generated histories will preserve the cue or that it will handle a new pickup sequence. During rollout, an error in the generated history can remove the information on which later predictions depend. That is one way strong one-step accuracy can coexist with weak multi-step simulation; it is not proof that every teacher-forced model failed to learn the rule.

Causal masking enforces this autoregressive structure inside a transformer: token $i$ attends only to tokens $j < i$ and to the context $h_t, a_t$. If training exposed later target tokens to that predictor, the stated next-token likelihood would leak information. The mask prevents that leakage; other generative architectures can enforce a valid factorization in different ways.

Once trained, we generate a future by sampling from the factorized distribution. Which sampling mode we choose changes the distribution of the futures we observe, and it changes how the model behaves inside a planner. There are several sampling modes:

• Greedy: take the argmax at each step. Deterministic but can produce degenerate, repetitive outputs when the model is uncertain.
• Temperature sampling: sample from $\text{softmax}(\text{logits} / \tau)$, where $\tau > 0$ is the temperature parameter: a scalar that scales the logits before the softmax is applied. Small $\tau$ approaches greedy; large $\tau$ flattens the distribution. Temperature changes concentration and diversity; its effect on fidelity depends on the task, metric, and calibration.
• Top-$k$ or nucleus (top-$p$) sampling: when their thresholds exclude tokens, restrict sampling to a high-probability subset. This can remove rare but decision-relevant events such as a pickup or death.

Sampling mode changes the futures a planner considers. Greedy decoding gives one model future. Repeated temperature samples expose several model futures at additional cost, but they are only useful if the model distribution is accurate and calibrated. A top-$k$ or top-$p$ threshold that excludes a low-probability pickup or death removes that event from sampled plans, biasing the evaluation of those plans.

When the environment is stochastic, the same conditioned history and action can have several valid next outcomes. The observed sample still supplies a valid negative-log-likelihood training target; across repeated samples, cross-entropy can learn the conditional distribution. A single point prediction or an overconfident distribution would instead misrepresent outcome variation, such as random enemy spawns or item drops.

Token prediction alone supplies code labels, not explicit reward or termination predictions. If a known reward or stopping rule can be evaluated from the predicted trajectory, including its actions and transitions where relevant, it can provide those signals; otherwise the model needs learned heads or another estimator. A decoded picture may look like the game, but visual fidelity alone does not tell a planner the return or remaining horizon.

A planner cannot reliably score candidate trajectories by return without a reliable reward mapping for its predicted trajectories, whether learned or known. Without a reliable termination signal, a model rollout may continue past an event that ends the real episode. Frame reconstruction alone does not test either signal.

Game trajectories have several temporal conventions that a dynamics model must respect.

• Action repeat. Many agents operate on a frame-level control rate of 60 Hz, but the agent chooses an action every $k$ frames, and the environment executes that action for the next $k$ consecutive frames. The value of $k$ becomes part of the effective time step of the model: one action corresponds to $k$ consecutive raw frames. If you forget this in the data pipeline, the model learns to predict an action's effect over an inconsistent time horizon.
• No-op controls. Some datasets include a "do nothing" action, sometimes as a legal action and sometimes as a technique for skipping events. The model treats it as just another action, so it must see it in the data.
• Frame skip and image stacking. The observation is often a stack of $k$ frames or a max-pooled pair, so the model's input is not a single image. The dynamics model learns a different causal structure than it would from single frames.
• Control latency. In some games and many robots, the action takes effect one or more steps after being issued. A model that assumes zero latency will systematically mis-predict transitions where the true effect lags.

These conventions are common in benchmark pipelines, though their exact settings vary. The preprocessing pipeline is part of the model's action interface. If an action means "press fire for four frames" during training but "press fire for one frame" at test time, the learned conditional transition no longer matches the queried process. The discrepancy can appear as systematic prediction error even when the model's reported uncertainty is low.

There are two scales at which you can train a tokenized game model. A single-game model is trained on trajectories from one environment. A multi-game model is trained on a union of trajectories from many environments, sometimes with an explicit game identifier and sometimes with environment-specific action adapters. The two regimes face different problems.

For a single game, $a_t$ has a fixed meaning within the chosen action interface, and the tokenizer and dynamics can be fitted to that game's data. Its tokens may specialize to the game's art style. The codebook size, training cost, and accuracy still depend on the architecture, data, and task; specialization alone guarantees none of them.

For multi-game, several complications arise:

• Observation representations must be compatible with the shared dynamics model. Differing resolutions, aspect ratios, or palettes may be handled by a shared tokenizer with preprocessing, or by game-specific encoders or adapters that produce compatible model inputs. Downsampling and style normalization are options, not requirements.
• Game identity or context is required. Because the dynamics depend on the game rules, the model must condition on something that identifies the game. This can be a one-hot game embedding, a learned "environment ID" token prepended to the sequence, or natural language. Text-prompt conditioning is taken up in Chapter 53 on Cosmos and omnimodal world foundation models, in Part IX: Foundation and World-Action Models.
• Action sets differ. A shared integer ID for "button 3" need not mean the same underlying effect in every game. "Jump" in one platformer might be "dash" in another and "no-op" in a third. Action masking can constrain the legal choices; per-game adapters or game-conditioned shared dynamics can model differing effects. The interface should state which semantics are shared and which depend on context.

The scaling opportunity is real: a shared tokenizer sees more visual variety, and a shared dynamics model sees more rule variety. Whether that leads to transfer depends on how well the model can condition on game identity and reuse behavior across games. It is a scaling hypothesis, not a guarantee. A model that shares a tokenizer but not a dynamics model can reuse visual features without necessarily transferring rules. A model that shares both needs enough game context to avoid mixing incompatible transition laws.