Autoregressive Token World Models › Next-Token Dynamics

We now formalize what the model learns. Let a rollout be serialized as a single token stream interleaving observations and actions:

$$x = \big(\underbrace{o_0}_{\text{obs tokens}},\ \underbrace{a_0}_{\text{action tokens}},\ \underbrace{o_1}_{\text{obs tokens}},\ a_1,\ o_2,\ \dots\big),$$

where:

• $x$: the token stream for one rollout
• $o_t$: the observation tokens emitted at step $t$
• $a_t$: the action tokens at step $t$, ordered before observation block $o_{t+1}$

Each $o_t$ and $a_t$ may itself expand into several tokens. A causal transformer with parameters $\theta$ defines an autoregressive factorization over any such stream:

$$p_\theta(x_1, \dots, x_n) = \prod_{t=1}^{n} p_\theta\big(x_t \mid x_{<t}\big),$$

where:

• $x_1, \dots, x_n$: the tokens of the stream $x$, listed in sequence order
• $x_{<t}$: all tokens preceding position $t$, that is, $(x_1, \dots, x_{t-1})$
• $p_\theta(x_t \mid x_{<t})$: the categorical distribution over the vocabulary that the transformer, with parameters $\theta$, assigns to the next token
• $n$: the total number of tokens in the stream

The chain rule permits any ordering of a joint distribution. A causal next-token implementation additionally needs shifted targets: the input at the prediction position must not already contain the target. The usual mask allows self-attention and only excludes later input positions. Finite context restricts the represented conditionals, whereas the unrestricted chain rule itself makes no independence assumption. Expected negative log-likelihood is minimized by the true conditional distribution when it is representable and optimized in expectation; the optimum can be sharp or diffuse.

For prediction under externally specified actions, define the causally conditioned rollout law below. This is not generally the observational conditional obtained by selecting trajectories with an entire future action sequence: under an adaptive logging policy, future actions can reveal earlier observations. The causal rollout construction instead supplies actions without that retrospective selection. Calling its factors true intervention dynamics additionally requires adequate conditioning information, action support, and no relevant hidden confounding.

$$p_\theta\big(o_{1:H} \Vert a_{0:H-1}; o_0\big) := \prod_{t=1}^{H} \prod_{j=1}^{m} p_\theta\Big(o_t^{(j)} \mid o_{<t},\ a_{<t},\ o_t^{(<j)}\Big),$$

where:

• $o_{1:H}$: the predicted observation blocks at steps 1 through $H$, with $o_t=(o_t^{(1)},\ldots,o_t^{(m)})$
• $\Vert$: causal conditioning on supplied actions, not selection on future logged actions
• $H$: the rollout horizon
• $o_0$: the initial observation tokens, given as context rather than predicted
• $a_{0:H-1}$: the externally specified actions for steps 0 through $H-1$
• $m$: the number of tokens per observation block
• $o_t^{(j)}$: the $j$-th token of the observation block at step $t$
• $o_{<t}$: all observation tokens from steps before $t$
• $a_{<t}$: actions before $t$, including $a_{t-1}$; future actions are excluded by the causal rollout construction, not by an identity for arbitrary observational conditioning
• $o_t^{(<j)}$: the first $j-1$ tokens of observation block $t$ itself

Read this carefully, because it encodes several design decisions:

• Actions are conditions in this objective. We mask action targets and learn observation-token factors. Joint models can also learn behavior-policy probabilities without mathematically invalidating the same observation conditionals. Either objective needs support and causal assumptions before its predictions are trusted under intervention. In our toy, controls are randomized independently of initial states and the transition mechanism is known; a planner in another dataset cannot assume that masking actions removes confounding or provides evidence for unseen actions.
• Within-frame dependence is a decoding choice. In raster next-token decoding, a token can read earlier tokens in its frame and accessible earlier frames. Independent simultaneous outputs cannot represent residual correlations unless those are resolved by conditioning or other latent structure. Iterative masked decoding permits a different within-frame information graph. Neither access pattern guarantees coherent generated frames.
• Available context determines usable evidence. For strict recomputation, state includes the last $W$ token IDs and positional metadata. Rolling caches can carry additional contextual information. A short window suffices only when the tokenized observation-action record preserves a sufficient predictive state. In partial observability, a Bayesian posterior would summarize history under a specified model; learned transformer activations need not equal that posterior or be calibrated. Windowed prediction can approximate belief-based prediction using accessible evidence without explicitly representing the full belief.

Training is by maximum likelihood under teacher forcing:

Under teacher forcing, the model is trained to maximize the likelihood of the true world tokens given their true contexts:

$$\mathcal{L}(\theta) = -\, \mathbb{E}_{x \sim \text{data}} \left[ \sum_{t : x_t \text{ is a world token}} \log p_\theta\big(x_t \mid x_{<t}\big) \right],$$

where:

• $\mathcal{L}(\theta)$: the expected negative log-likelihood of the world tokens under parameters $\theta$, to be minimized
• $x$: a token stream drawn from the training data, so the expectation is taken under the data distribution rather than the model's own
• $x_t$: the target token at position $t$, with $x_{<t}$ the true context fed to the model as input
• the sum ranges only over positions whose target is a world (observation) token; action tokens are masked out of the loss

Teacher forcing gives one loss term per supervised world token, uses recorded rather than self-generated histories, and permits layerwise-parallel evaluation across sequence positions. It avoids error feedback in the training inputs, but does not guarantee numerical stability or sufficient representations. Dense attention still has quadratic pairwise work. When learned conditionals reproduce the data process under matched actions, generated and recorded histories can have the same distribution; mismatch is a risk of imperfect models and changed deployment conditions, not an unavoidable consequence of teacher forcing.

It is worth pausing on a distinction this book keeps returning to, because token world models blur it more than most. The model above predicts observations: given a history and a candidate action sequence, it rolls out a distribution over future observation tokens. That is an open-loop forecasting capability, and everything in this chapter is about its quality.

Using predictions is a separate task from producing them. Model-predictive control searches candidate actions, executes an initial action, then replans from new evidence. Fresh observations can correct visible drift, but robustness still depends on model accuracy, sensing, constraints, and controller design. Differentiable planning through discrete choices needs an expectation, relaxation, or gradient estimator; an agent can also learn from imagined experience or distill a planner's actions. Search need not find a global optimum and may exploit overestimated rewards. Conversely, a model can retain enough task-relevant structure despite poor texture reconstruction. Predictive fidelity and decision usefulness therefore require separate tests, without assuming a general correlation.