Tokenized Game and Environment Models › From frames to environment tokens

The starting point is a trajectory: a sequence of interactions with an environment. If we write the observation, action, reward, and termination flag at time $t$, the natural notation for a trajectory is

$$(o_0, a_0, r_1, d_1, o_1, a_1, r_2, d_2, \ldots, o_{T-1}, a_{T-1}, r_T, d_T, o_T),$$

where:

• $o_t \in \mathbb{R}^{H \times W \times C}$: the rendered observation at time $t$ in the image-based setting studied here. It may be an RGB or grayscale frame, or a stack of frames, possibly with a HUD overlay. Other interfaces, such as RAM observations, require different notation.
• $a_t \in \mathcal{A}$: the discrete control chosen at time $t$. In the Arcade Learning Environment, the full discrete Atari action set has 18 controls; a game's reduced action set may contain far fewer. A multi-game dataset may use a union of the included games' actions, with each game using its applicable subset.
• $r_{t+1} \in \mathbb{R}$: the reward received after the transition into time $t+1$. Depending on the game, rewards may be sparse or frequent; pickups, scores, and deaths are possible reward events.
• $d_{t+1} \in \{0,1\}$: the termination flag for the transition into time $t+1$. When $d_{t+1}=1$, the episode ends; a subsequent reset begins a new episode.
• $T$: the length of the episode.

We use $t$ for the observation and action indices and $t+1$ for the reward index on purpose. Reward is associated with the transition after action $a_t$; in a partially observed environment it may depend on hidden state, not solely on the two rendered frames. Part II, Ch 1: Markov and Partially Observable Decision Processes used this indexing convention. A model that misaligns rewards with actions can learn the wrong decision signal even if frame reconstruction looks good.

The tuple above is a single trajectory. The index $t$ labels the observation, $t+1$ labels its reward and termination, and the episode ends at $T$, so a trajectory contains $T+1$ observations and $T$ action-reward bundles. Any model we build will be asked to predict the next element of this sequence given the preceding ones.

It is worth pausing on how much structure the raw trajectory leaves implicit. The observation $o_t$ is high-dimensional and redundant. The action $a_t$ is a small discrete label with no meaning until we connect it to an effect. The reward $r_{t+1}$ and termination $d_{t+1}$ carry explicit task feedback, yet a naive "predict the next frame" objective does not train separate predictions for them. They may also be rare compared with ordinary visual changes. The rest of this section is about the first compression step, and the following section is about the second.

A preprocessed $84 \times 84$ RGB frame contains $84 \cdot 84 \cdot 3 = 21{,}168$ channel values; this is not the native resolution of every game interface. A sixty-second clip at 30 frames per second has $30 \cdot 60 = 1{,}800$ frames. Feeding all pixel values directly into a sequence model is costly, while adjacent frames share much of their content. The model may devote capacity to predicting unchanged walls instead of small events that alter the game's state.

A tokenizer can map each observation (or short clip) into a shorter sequence of discrete symbols. If the token count is reduced by $100\times$ relative to a specified raw-pixel tokenization, the dynamics model sees roughly one-hundredth as many symbols per frame. Its predictions become class labels rather than continuous intensities. A softmax over those labels represents a categorical distribution over code outcomes; whether that distribution is calibrated must be tested separately.

The token count per frame $L$ and vocabulary size $K$ are two design knobs. $L$ is the number of codes emitted per observation; $K$ is the number of entries in the shared vocabulary. At a fixed token-count context window, increasing $L$ shortens the frame history that fits. Increasing $K$ at fixed $L$ does not change that token count, but changes the information capacity per token and the model's embedding and output-classification costs. These are separate tradeoffs, not interchangeable ways of setting the context length.

The combinatorial space of possible code sequences grows with both knobs. An autoregressive factorization is one tractable way to model a distribution over that space; it is not the only structured alternative.

The tokenizer can be one of several kinds:

• Deterministic and fixed: a hand-designed encoder, such as a color histogram or a downsample plus palette quantization.
• Deterministic and learned: a VQ-VAE or similar autoencoder whose encoder is trained to assign a specific code to each input under a nearest-neighbor rule.
• Stochastic and learned: a variational autoencoder with a discrete or relaxed-discrete bottleneck, where the encoder outputs a distribution over codes.
• Temporal: a video tokenizer that maps a short clip $(o_t, \ldots, o_{t+k})$ into a sequence of codes, so the dynamics model predicts a chunk of frames at once.

Which one you pick changes what the dynamics model sees. If the tokenizer is trained separately and frozen, the dynamics model inherits its quirks. If they are trained jointly, the dynamics objective can shape the tokenizer to keep information that matters for prediction, which is sometimes good and sometimes leads to degenerate codes. We will come back to this tradeoff.

Vector quantization is a common discrete-tokenizer design. Walking through its codebook lookup helps explain representation errors that matter later in this chapter.

Suppose an encoder $f_\phi$ maps an observation $o$ to a spatial feature map of shape $L \times D$: $f_\phi(o) \in \mathbb{R}^{L \times D}$. Think of it as a grid of $L$ vectors, each of dimension $D$. We maintain a codebook of $K$ vectors $e_1, \ldots, e_K \in \mathbb{R}^D$. For each spatial position $\ell$, we replace the feature vector with the nearest codebook entry. This is the quantize step, and it is the only place where the continuous encoder output becomes a discrete symbol:

$$\begin{aligned} z_\ell &= \arg\min_{k \in \{1,\ldots,K\}} \|f_\phi(o)_\ell - e_k\|_2^2 && \text{(nearest-neighbor lookup)} \\ \hat{f}_\ell &= e_{z_\ell} && \text{(replace with the codebook entry)} \end{aligned}$$

where:

• $f_\phi(o)_\ell$: the encoder output at spatial position $\ell$, a vector in $\mathbb{R}^D$ that we want to quantize
• $e_k$: the $k$-th entry of the learned codebook, a representative vector in the same $D$-dimensional space
• $\|f_\phi(o)_\ell - e_k\|_2^2$: the squared Euclidean distance between the encoder output at position $\ell$ and codebook entry $k$, which measures how close the two vectors are
• $z_\ell$: the index of the nearest codebook entry, an integer in $\{1,\ldots,K\}$ that becomes the token for position $\ell$
• $\hat{f}_\ell$: the quantized feature at position $\ell$, obtained by looking up the nearest codebook vector $e_{z_\ell}$, so the feature map is replaced by codebook entries before decoding

The discrete codes $z = (z_1, \ldots, z_L)$ are the tokens. The decoder $g_\psi$ reconstructs an observation from the quantized features: $\hat{o} = g_\psi(\hat{f})$. Training minimizes a reconstruction loss plus a commitment term that keeps the encoder close to the codebook and a codebook term that keeps the codebook close to the encoder. The two auxiliary terms, and the straight-through estimator that lets gradients pass through the non-differentiable $\arg\min$, are engineering details; the important point for us is that the tokens are category labels, not numbers, and that the codebook is a jointly learned dictionary. A code is not "a little bit like" another code in the way two nearby real numbers are close; the codes are unordered categories, and the model has to learn any similarity structure between them from data.

The quantization step partitions the encoder's feature space by nearest codebook vector. The figure below shows the assignments; it does not depict codebook collapse, because every illustrated entry is used.