Section 38.4: Transformer world models (IRIS)

"The tokenizer is not preprocessing; it is the alphabet the world model is allowed to think in."

A Latent Sequence That Behaves Like Language

Problem First

Recurrent world models summarize history in a fixed-size hidden state. That is efficient, but it can bottleneck long-range structure. Transformer world models were introduced to test whether image-token sequences and causal attention can capture temporally extended dependencies more faithfully, especially when the environment behaves like a structured visual language.

Core Model

IRIS discretizes visual observations with a tokenizer, then models action-conditioned token sequences autoregressively: $$c_t = \mathrm{Tokenizer}(o_t), \qquad p(c_{t+1} \mid c_{\le t}, a_{\le t}) = \prod_i p(c_{t+1,i} \mid c_{\le t}, a_{\le t}, c_{t+1,

This changes the inductive bias. An RSSM assumes a compact hidden state should summarize the past; IRIS assumes the model can recover what matters by attending over a token history. The benefit is flexible long-range dependency modeling. The cost is quadratic sequence processing and a strong dependence on the tokenizer's ability to preserve the variables the controller needs.

For control, the model must still satisfy the same decision criterion as any world model: generated futures must be action-conditional, temporally stable, and sufficiently aligned with reward-relevant state that a policy trained in imagination transfers back to real trajectories.

Minimal Probe

The probe below mirrors the core IRIS idea with toy tokens. It rolls a short token sequence forward under actions and checks whether the generated symbol stream still preserves the task-relevant state transition pattern.

# Roll a tokenized world state forward under action-conditioned updates.
# The token history acts like a tiny visual language for the world model.
token_state = [3, 7, 2]
actions = [1, 0, 2]
generated = []
for action in actions:
    next_token = (token_state[-1] + action + token_state[0]) % 10
    generated.append(next_token)
    token_state = token_state[1:] + [next_token]
print({"generated_tokens": generated, "final_context": token_state})

Expected behavior: The generated token pattern should depend on both the rolling context and the chosen actions. If changing the actions barely changes the sampled future tokens, the model has become a passive video predictor instead of an action-conditioned world model.

Practical Recipe

1. Audit the tokenizer first; if it destroys small but action-relevant details, no amount of attention depth will fix the problem.
2. Measure action sensitivity by changing only the action prefix and checking whether sampled futures diverge in the correct way.
3. Keep track of context length because attention quality can improve even while compute and memory costs become unacceptable for control loops.
4. Compare the transformer against an RSSM on matched rollout horizon and wall-clock budget, not only sample efficiency.

IRIS is important less because transformers always beat recurrent models and more because it reframed world modeling as discrete sequence prediction. That move lets researchers borrow a mature set of tools from language modeling: tokenization, causal masking, long-context studies, and sampling diagnostics.

The hidden engineering challenge is semantic granularity. If one token change corresponds to a tiny texture difference, the model may look visually sharp but remain weak for control. If tokens are too coarse, the model may ignore subtle collision, grasp, or lane-position cues. In other words, tokenization is where control relevance is won or lost.