🤖 AI

Distilling Game Code World Model Generation into Lightweight Large Language Models

This paper proposes a scalable post-training pipeline combining Supervised Fine-Tuning and Reinforcement Learning with Verifiable Rewards to distill Game Code World Model generation capabilities from frontier models into a lightweight 3B-parameter LLM, enabling it to accurately generate executable game environments from natural language rules.

Original authors: Tyrone Serapio, Arjun Prakash, Haoyang Xu, Kevin Wang, Amy Greenwald

Published 2026-05-26

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Tyrone Serapio, Arjun Prakash, Haoyang Xu, Kevin Wang, Amy Greenwald

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you want to build a robot that can play any board game you describe to it, from Chess to Poker. To do this, the robot needs a "rulebook" that tells it exactly how the game works: how pieces move, what counts as a win, and what happens when a player makes a mistake.

In the past, building these rulebooks for new games was like hiring a team of expert engineers to write code from scratch every time. It was slow, expensive, and required very smart (and very costly) computers to figure out the details.

This paper proposes a new way to do it: teaching a small, cheap computer to write the rulebooks itself.

Here is the breakdown of their approach using simple analogies:

1. The Problem: The "Genius Tutor" vs. The "Apprentice"

Currently, if you want a computer to generate a game rulebook (which the authors call a Game Code World Model), you have to ask a "Genius Tutor" (a massive, expensive AI like GPT-4 or Gemini).

The Old Way: You ask the Genius Tutor, "Write the rules for this new game." If it makes a mistake, you check the code, point out the error, and ask the Genius to try again. You might have to do this loop dozens of times. It's like hiring a world-class chef to cook a simple sandwich, but you have to keep tasting it and sending it back to the kitchen until it's perfect. It works, but it's slow and costs a fortune.

2. The Solution: Distilling the Knowledge

The authors wanted to see if they could train a small, lightweight AI (an "Apprentice") to do this job without needing the Genius Tutor to check its work every time. They call this process "Distilling."

Think of it like a master chef (the large AI) teaching a junior chef (the small AI) how to cook. Instead of the junior chef asking the master for help every time they chop an onion, the master spends time training the junior chef until the junior chef can cook the dish perfectly on their own.

3. The Training Process: Two Steps

The authors used a two-step training pipeline to turn the small AI (Qwen2.5-3B) into a game-rule expert:

Step 1: Supervised Fine-Tuning (SFT) - "The Cookbook"
They showed the small AI 30 different games (like Tic-Tac-Toe, Poker, and Blackjack) along with the correct code for how those games work. It's like giving the apprentice a stack of perfect cookbooks and saying, "Memorize these recipes."
- Result: The AI got much better at writing code that didn't have typos or syntax errors (like forgetting a comma or a bracket).
Step 2: Reinforcement Learning with Verifiable Rewards (RLVR) - "The Taste Test"
Just memorizing recipes isn't enough; the food has to taste right. The authors built an automated "Taste Tester" (a verification framework).
- The AI tries to write a game rulebook.
- The "Taste Tester" runs the code. Does the game crash? Do the rules make sense? Does the game end correctly?
- If the code passes the test, the AI gets a "reward" (points). If it fails, it gets a penalty.
- The AI tries again and again, learning from these rewards to write code that actually works, not just code that looks right.

4. The Results: What Worked and What Didn't

The authors tested their trained "Apprentice" on games it had never seen before (Out-of-Distribution games).

The Good News: The small, trained AI became significantly better at writing valid game code. It learned to avoid syntax errors and followed the basic structure of game rules much better than it did before training. It proved that you can teach a small model to do this job without needing the expensive "Genius Tutor" to check every single line of code.
The Bad News: The AI still struggled with the most complex games, especially those with hidden information (like Poker, where you don't know what cards your opponent has).
- The Metaphor: The apprentice can cook a perfect grilled cheese sandwich (simple games) and even a complex lasagna (perfect information games). But when asked to cook a dish that requires guessing what the other person is thinking (imperfect information games), the apprentice gets confused and sometimes just gives up or writes a simplified, incorrect version.
- Interestingly, even the "Genius Tutor" (GPT-4) struggled with these complex hidden-information games, suggesting this is a very hard problem for all current AIs.

5. The Bottom Line

This paper shows that we can take the ability to generate complex game rules from expensive, giant AI models and "distill" it into smaller, cheaper models through training.

Before: You needed a supercomputer and a lot of time to generate a game engine.
Now: You can train a small, efficient model to do it, provided the game isn't too complex.

The authors conclude that while this small model isn't perfect yet (especially for tricky games with hidden secrets), it is a major step toward making it easy and cheap to automatically create digital worlds for AI agents to play in. They did not claim this works for medical diagnosis, financial trading, or other real-world applications, only for generating code that simulates game environments.

Technical Summary: Distilling Game Code World Model Generation into Lightweight Large Language Models

Problem Statement

The construction of World Models—internal representations of an environment that allow AI agents to simulate action consequences—is critical for multi-agent systems. While recent work has introduced Code World Models (CWMs), where Large Language Models (LLMs) translate natural language game descriptions into executable Python code compatible with solvers like Monte Carlo Tree Search (MCTS), current approaches face significant scalability barriers.

Existing methods for generating Game Code World Models (GameCWMs) rely heavily on "frontier" models (e.g., Gemini 2.5 Pro, GPT-4) and extensive test-time compute. These approaches utilize iterative refinement loops where an LLM generates code, verifies it against unit tests or ground-truth trajectories, and regenerates code upon failure. This process is computationally expensive and limits accessibility. The central question of this work is whether the capabilities required to generate valid GameCWMs can be distilled into smaller, lightweight models (specifically 3B parameter models) through post-training, thereby eliminating the dependency on frontier models and inference-time refinement loops.

Methodology

The authors propose a three-component pipeline to distill GameCWM generation capabilities into Qwen2.5-3B-Instruct:

1. Dataset Curation

A dataset of 30 extensive-form games was constructed, spanning:

Perfect Information Games: e.g., Tic-Tac-Toe, Connect4, Generalized Chess.
Imperfect Information Games: e.g., Blackjack, Kuhn Poker, Leduc Poker, Gin Rummy.
Out-of-Distribution (OOD) Games: Novel games not present in standard LLM training data, including a new game named Converge.

The dataset consists of prompt-code pairs $(P_G, C_G)$ , where the prompt includes the API specification, natural language game description, explicit action strings, and an example trajectory. The target output is ground-truth Python code implementing the game engine.

2. Hierarchical Verification Framework

To enable training without ground-truth game engines for novel games, the authors developed a four-tier verification framework that evaluates generated code:

Tier 1 (Static Analysis): Verifies syntactic validity, API completeness (all required methods present), type signatures, and successful compilation.
Tier 2 (Dynamics/Fuzzing): Executes 100 random synthetic trajectories to check for determinism, non-terminal logic, terminal logic, crash-free execution, and state immutability.
Tier 3 (Semantic Rule Adherence): Compares the generated engine's behavior against a small set of "scenario traces" (game histories) generated by a frontier LLM from the game rules. This tests boundary conditions and typical gameplay.
Tier 4 (Information Consistency): Specific to imperfect-information games, this tier evaluates the resample_history function. It simulates a random walk to a state, extracts an observation, and verifies that the resampled history produces the same observation, ensuring epistemic boundaries are maintained.

3. Post-Training Pipeline

The training process involves two stages:

Supervised Fine-Tuning (SFT): The base Qwen2.5-3B-Instruct model is fine-tuned on the prompt-code pairs using LoRA. This stage aims to align the model with the required API structure and syntax.
Reinforcement Learning with Verifiable Rewards (RLVR): Following SFT, the model is further refined using Group Relative Policy Optimization (GRPO). The reward signal is a composite score derived from the four-tier verification framework:
$r_i = \sum_{t=1}^{4} w_t S_t$
Where $S_t$ is the pass rate at tier $t$ , and weights ( $w_1=0.15, w_2=0.25, w_3=0.3, w_4=0.3$ ) prioritize semantic logic and information handling over mere syntax.

Key Results

Experiments were conducted on 7 held-out games (including OOD and imperfect-information games), generating 25 samples per game.

Overall Performance: The full SFT+RLVR pipeline produced the strongest model, achieving a mean verification score of 53.2%, significantly outperforming the base model (47.3%) and the SFT-only model (49.0%). However, it remained below the performance of frontier models like GPT-4o (66.7%).
Syntactic and Structural Gains: SFT significantly improved Static (Tier 1) and Fuzz (Tier 2) scores, demonstrating that fine-tuning effectively teaches the model the required API structure and robust transition logic.
Semantic Adherence: RLVR provided the most significant boost to Semantic (Tier 3) scores. The SFT+RLVR model achieved a 14.4% semantic score compared to 8.5% for the base model, indicating that RL helps the model adhere to game rules beyond just syntax.
Imperfect Information Challenges: All models, including GPT-4o, struggled with Information (Tier 4) consistency in imperfect-information games. The best Qwen model (SFT+RLVR) achieved only 5.3%, compared to 14.0% for GPT-4o.
Limitations Observed:
- Overfitting in SFT: On the in-distribution game Gin Rummy (not in the training set), SFT actually decreased performance (from 49.7% to 31.4%) due to the model blindly applying list-comprehension templates learned during training, introducing syntax errors.
- Reasoning Boundaries: The models often inserted comments acknowledging the difficulty of implementing resample_history logic, opting for simplified stubs (e.g., returning the last action) rather than generating complex reasoning paths. This suggests the base model's latent space may lack the necessary "Theory of Mind" reasoning capabilities for these tasks.

Significance and Claims

The paper claims that GameCWM generation can be partially distilled into lightweight LLMs through post-training. The primary significance of this work is the demonstration that:

Resource Efficiency: It is possible to reduce reliance on expensive frontier models and iterative test-time refinement loops by shifting the computational burden to post-training.
Two-Stage Necessity: A combination of SFT (for structure/syntax) and RLVR (for semantic adherence) is more effective than either method alone.
Scalability: The approach offers a scalable path toward automatic environment generation from natural language, making GameCWMs accessible for a broader range of applications.

However, the authors remain modest about the current state of the technology. They explicitly state that complex semantic reasoning and imperfect-information handling remain challenging and exceed the current capacity of both the lightweight models and the training pipeline. The work does not claim to have solved the problem of perfect GameCWM generation but rather provides a viable, more efficient alternative to current frontier-model-dependent approaches.