World Models From Scratch 2: Model Training and Dreaming [P]

# World Models From Scratch 2: Model Training and Dreaming **TL;DR:** This second part of the series builds a world model for Super Mario Land by training a transformer to predict the next game frame (as a grid of tokens) from a sequence of past frames and player actions, then demonstrates "dreaming" where the model predicts future frames based solely on its own predictions. ## Recap of Part 1 The first video established the foundation by recording 150,000 frames of *Super Mario Land* gameplay. A tokenizer was trained to convert these raw frames (144x160 pixel matrices with four color values per pixel) into a compact 18x20 grid representation. This grid size was chosen based on the tile layout in Game Boy memory. Each position in this grid is then mapped to a vector of five elements. The inverse process—reconstructing an approximate image from this internal representation—was also learned. ## The World Model Architecture The core idea is to predict the next frame given a sequence of past frames and the actions taken. 1. **Input:** A context window of previous frames (e.g., the last 8 frames) and the most recent player action (a button press or no action). 2. **Processing:** Each frame passes through the tokenizer to become its 18x20 grid of token vectors. 3. **Prediction:** The world model, a transformer, takes this sequence and predicts the next frame in the same token grid format. 4. **Output:** The predicted token grid is passed through the inverse tokenizer to generate an image. ### Key Difference from GPT-style Models The masking strategy for this world model differs significantly from a standard causal language model: * **GPT-style:** Predicts one token at a time, masking all future tokens (a diagonal mask pattern). * **This World Model:** Predicts an entire frame at once. The model sees all tokens *within* a frame but only sees past frames in the context window (frame-level masking). This approach treats an 18x20 grid (320 token values) as the unit of prediction, rather than individual values. ## Model Components and Training The implementation in a Colab notebook involves several key components: ### Embeddings and Positional Encodings The model uses multiple embedding layers: * **Token Embeddings:** Converts the tokenizer's output vectors for each position in the frame grid. * **Action Embeddings:** Encodes the player's action (e.g., jump, right, left, or no action). Positional information is added in two dimensions: * **Spatial Positions:** Encodes the location of each token *within* a single frame (1 to 320). This is added globally to all frames, as the spatial layout is constant. * **Temporal Positions:** Encodes the order of frames *within the context window* (e.g., frame 1 to frame 8). This is added per-frame so the model understands the sequence of events over time. ### Training Setup The model is a transformer block trained with a standard cross-entropy loss, comparing its predicted token for the next frame against the tokenizer's output of the actual frame from the simulator. * **Optimizer:** Adam with weight decay (beta1=0.95, weight_decay=0.1). * **Context Window:** 8 frames, used to predict the 9th frame. * **Training Progress:** The validation accuracy improves and stabilizes over training steps, as measured via **teacher forcing** (where all input context frames are ground-truth frames from the simulator). ## Dreaming: Autoregressive Prediction The most interesting part is "dreaming," where the model operates in a purely autoregressive mode. * **Process:** Starting from an initial state, the model predicts frame t+1 based on the context and an action. It then uses its own prediction for frame t+1 as part of the context to predict frame t+2, and so on. * **Observation:** The video compares three rows of frames: 1. **Ground Truth:** Frames from the original Game Boy simulator. 2. **Tokenizer Reconstruction:** Frames reconstructed from the ground truth via the tokenizer. 3. **Dreamed Frames:** Frames generated entirely by the world model's own predictions. * **Results:** The dreamed frames are coherent and recognizable as *Super Mario Land* gameplay, following the sequence of actions (right, jump). The quality is naturally limited by the tokenizer's reconstruction fidelity. * **Drift Analysis:** An experiment shows that when the model predicts frames based on its own prior predictions, the pixel-level error (drift) compared to a baseline increases over time. This is expected error accumulation in autoregressive generation. ## Conclusion The project demonstrates that a world model can learn to predict future visual states of a simple game environment solely from recorded frame-action sequences, without access to the game's ROM or simulator logic. The model can then "dream" coherent gameplay sequences by predicting its own future states. Further optimization and more data could improve accuracy, but the core principle—autoregressive prediction of a tokenized visual world—is proven feasible. **Source:** [World Models From Scratch 2: Model Training and Dreaming](https://youtu.be/L4kKIHk8a8s)