A Brief Overview of Code World Models

Posted on In Views:

Most modern world models represent environment dynamics using neural networks, including recurrent models, Transformers, and diffusion models. An alternative research direction is to represent world models as executable programs.

Such models, referred to as code world models (CWMs), are constructed through LLM-based program synthesis. In this paradigm, the LLM does not directly act as the world simulator. Instead, it synthesizes executable code that implements the simulator, including the environment's state transitions, observations, rewards, and other relevant dynamics.

We will discuss three representative CWM works:

  1. WorldCoder (2024)
  2. DeepMind CWM (2025)
  3. PoE-World (2025)

We focus primarily on WorldCoder, while the latter two are used to complement the discussion with more advanced formulations and design choices and we don't talk about the details of them.

A Brief Introduction to Code World Models

Problem formulation

We consider both fully observable and partially observable environments.

For an MDP, \[ s_{t+1}, r_{t+1} \sim P(\cdot \mid s_t, a_t), \] where \(s_t\) is the environment state and \(a_t\) is the action.

For a POMDP, the underlying state is not directly observed. Instead, \[ s_{t+1}, r_{t+1} \sim P(\cdot \mid s_t, a_t), \qquad o_t \sim O(\cdot \mid s_t), \] where \(o_t\) is the observation, such as an image.

By interacting with the environment, we collect trajectories of states or observations, actions, and rewards. A code world model uses an LLM to synthesize an executable program \(\hat{E}\) from these interactions and, optionally, an environment description: \[ \text{environment description + interaction trajectories} \rightarrow \text{LLM} \rightarrow \hat{E}. \] The goal is for \(\hat{E}\) to reproduce the behavior of the original environment as accurately as possible.

Central question

The central question of CWM research is:

What are the limits of LLM-based code world model synthesis?

More specifically, what kinds of environments can LLMs synthesize as executable world models, and how accurately can the resulting programs reproduce them?

This concerns both the scope of environments that can be modeled and the fidelity of the synthesized simulator.Given the same actions and initial conditions, \(\hat{E}\) should reproduce the observations, rewards, and dynamics of the real environment as accurately as possible.

1. WorldCoder

1.1 Core idea

The overall architecture of WorldCoder, and more generally a CWM agent, is a closed loop:

WorldCoder workflow, from WorldCoder paper Fig. 1

The agent collects transitions from the real environment into a replay buffer. An LLM uses these data to write and refine an executable world model, while a planner acts inside the synthesized model and executes the selected actions in the real environment. New interactions are then added to the replay buffer to further refine the model.

Using AI to synthesize environment simulators is not new, but WorldCoder is one of the first works to systematically study LLM-generated executable world models, although it considers a relatively restricted setting.

WorldCoder assumes a fully observable and deterministic environment. Each task is additionally specified by a natural-language goal \(c\), such as "pick up the ball". The environment is formulated as a contextual MDP: \[ \mathcal{M}_c=(\mathcal{S},\mathcal{A},T,R_c), \] where \[ T:\mathcal{S}\times\mathcal{A}\rightarrow\mathcal{S} \] is the transition function, and \[ R_c:\mathcal{S}\times\mathcal{A}\times\mathcal{S} \rightarrow \mathbb{R}\times{0,1} \] is the goal-conditioned reward and termination function. WorldCoder uses an LLM to synthesize Python implementations of these functions, denoted by \(\hat{T}\) and \(\hat{R}_c\).

WorldCoder also does not operate on raw pixels. Instead, it assumes fully observable object-centric symbolic states, for example:

[
    Agent(x=1, y=2, direction=(0, -1), carrying=None),
    Door(x=2, y=2, color="yellow", state="locked"),
    Key(x=1, y=3, color="yellow"),
    Goal(x=3, y=3),
]

In a nutshell, WorldCoder presents a general framework for CWM tasks and it has following restrictions:

  1. only considers fully observable environments;
  2. only considers deterministic environments;
  3. models symbolic states instead of directly learning pixels.

It's worth noting that all 3 CWM methods introduced here have restriction (3).

1.2 Method: refining a CWM

Generating a CWM is only the first step. The more important question is:

What should we do when the synthesized simulator is wrong?

Suppose the real environment gives us the following transition:

state:  agent is carrying a yellow key
action: toggle the yellow locked door

real environment:
Door.state = open

but our synthesized CWM predicts:

CWM:
Door.state = locked

Then the CWM is clearly wrong.

A natural solution is to show this failure to the LLM:

current CWM code
      +
failed transition
      +
correct result
      ↓
     LLM
      ↓
refined CWM code

For example, the LLM may discover that the transition code forgot to handle the interaction between a key and a locked door, and add the missing rule.

We then run the CWM on the collected transitions again. If another case fails, we repeat the process.

synthesize CWM
      ↓
test it
      ↓
find a failure
      ↓
refine the code
      ↓
test again
      ↓
     ...

This process is usually called CWM refinement or CWM repair.

An analogy to neural network training

This is conceptually similar to training a neural network world model.

For a neural network, we first define an optimization objective, usually a loss function:

prediction
    ↓
compare with ground truth
    ↓
loss

Then gradient descent tells us how to change the parameters:

loss
 ↓
gradient descent
 ↓
updated parameters

For a CWM, we can make a similar decomposition:

Neural world model                 Code world model

prediction                         program execution
    ↓                                   ↓
loss / objective                   evaluation objective
    ↓                                   ↓
gradient descent                   LLM refinement
    ↓                                   ↓
updated weights                    updated code

The major difference is that the evaluation objective of a CWM does not have to be a differentiable numerical loss.

It can simply be a test such as:

Did the CWM reproduce this real transition?

yes → pass
no  → fail

Therefore, CWM "training" can be understood as two problems:

  1. Evaluation: how do we decide whether a candidate CWM is good?
  2. Refinement: if it is not good, how do we ask the LLM to improve it?

Evaluation: what makes a CWM good?

The most obvious criterion is replay consistency.

Suppose interaction with the real environment gives us a replay buffer:

$$ D =

{(s,a,r,s',c,d)}, $$

where \(c\) is the task goal and \(d\) indicates whether the episode terminates.

Every recorded transition gives us a unit test:

    real environment
(s, a) → (s', r, d)

          vs.

         CWM
(s, a) → (ŝ', r̂, d̂)

If

\[ (\hat s',\hat r,\hat d)=(s',r,d), \]

the CWM passes this test.

For a deterministic environment, we can simply require the CWM to pass all observed transitions.

This is very similar to regression testing in software engineering:

replay buffer
     ↓
many test cases
     ↓
candidate CWM
     ↓
pass all tests?

Different CWM methods may add other heuristical objectives, such as physical constraints, planning objectives, or numerical likelihoods for stochastic environments.

Refinement: how to improve the CWM?

Once the evaluation tells us that a CWM is wrong, the next question is:

How should we modify the program?

The simplest approach is iterative repair. We give the LLM the current CWM together with diagnostic information, such as a failed transition, violated constraint, or runtime error:

current CWM
    +
failed example / constraint
    +
expected behavior
       ↓
      LLM
       ↓
refined CWM
       ↓
evaluate again

This gives a single refinement chain:

CWM_0 → CWM_1 → CWM_2 → CWM_3 → ...

However, an LLM repair is not guaranteed to be better. One edit may fix one failure while introducing another, and different repairs may lead to very different hypotheses about how the environment works.

A more general approach is therefore to keep multiple candidate CWMs:

            CWM_0
          /       \
    repair         repair
      ↓              ↓
    CWM_1          CWM_2
    /   \
repair   repair
  ↓        ↓
CWM_3    CWM_4

Each node is a different hypothesis about the environment. After generating new candidates, we:

generate candidates
        ↓
evaluate each CWM
        ↓
assign scores / pass-fail results
        ↓
select a promising candidate
        ↓
refine it again

This turns CWM refinement into a program-search problem:

\[ \text{candidate CWM} \rightarrow \text{evaluation} \rightarrow \text{selection} \rightarrow \text{LLM refinement}. \]

The selection strategy can balance:

  • exploitation: continue refining candidates that already perform well;
  • exploration: try less-explored candidates that may lead to a better solution.

WorldCoder follows this branching view: it maintains multiple candidate programs and uses REx to decide which candidate should receive the next LLM refinement.

So, analogous to neural-network optimization,

Neural network                         CWM

loss / objective                       evaluation objective
       ↓                                      ↓
optimizer                              candidate selection
       ↓                                      ↓
gradient update                        LLM refinement
       ↓                                      ↓
new parameters                         new CWM candidate

The important point is that evaluation tells us which hypotheses are good, while refinement searches for better hypotheses.

For simplicity, we will not focus on the details of tree-search strategies here. The central issue for the following discussion is the evaluation objective used to judge a CWM.

WorldCoder: replay consistency is not enough

WorldCoder evaluates a candidate CWM with two constraints.

The first, \(\phi_1\), is replay consistency, which we have already discussed: \[ \hat T(s,a)=s', \qquad \hat R(c)(s,a,s')=(r,d) \] for every observed transition.

The second, \(\phi_2\), is optimism under uncertainty:

From the initial state and mission, the CWM should contain at least one trajectory that reaches positive reward and terminates successfully.

Thus, WorldCoder searches for \[ \phi_1\land\phi_2. \] Intuitively:

φ1: explain what has already happened
φ2: imagine at least one way to succeed

1.3 Illustration: MiniGrid DoorKey

A representative WorldCoder task is MiniGrid DoorKey:

MiniGrid DoorKey environment

MiniGrid DoorKey. The agent must pick up the key, unlock the door, and reach the goal. Animation from the official MiniGrid project.

The agent receives a natural-language goal:

"use the key to open the door and then get to the goal"

The environment provides a fully observable symbolic state:

[
    Agent(x=1, y=2, carrying=None),
    Key(x=1, y=3, color="yellow"),
    Door(x=2, y=2, color="yellow", state="locked"),
    Goal(x=3, y=3),
]

Through interaction, WorldCoder must infer rules such as:

pick up key
    ↓
Agent.carrying = key

toggle locked door with key
    ↓
Door.state = open

reach goal
    ↓
reward > 0, done = True

These rules are encoded in the synthesized Python world model: \[ \hat T(s_t,a_t)=s_{t+1}, \] which models how the environment changes, and \[ \hat R_c(s_t,a_t,s_{t+1})=(r_{t+1},d_{t+1}), \] which models whether a transition satisfies the given goal.

The planner can then search inside the learned CWM: \[ \text{key} \rightarrow \text{pick up} \rightarrow \text{door} \rightarrow \text{unlock} \rightarrow \text{goal}. \] This example captures the basic WorldCoder setting: the environment dynamics are unknown, but the state is fully observable and symbolic, making it possible for the LLM to recover the underlying rules as executable code.

Another research line: using LLM as an agent

Beyond CWMs, a parallel research direction uses the LLM itself as the agent in interactive environments: the LLM observes the environment, reasons about it, and directly selects actions.

1. LLM directly select action

A representative early example is ReAct, which follows a simple loop:

reason → action → observation → reason → ...

The LLM is directly responsible for action selection.

2. Building world models to help action selection

More recent approaches introduce a harness around the LLM to support longer-term reasoning. A common idea is to let the agent maintain and revise hypotheses—or an explicit world model—about how the environment works.

Depending on the method, this world understanding can be represented in different forms, such as free-form natural language or an executable code world model (CWM).

comparison of schema harness and VISTA harness, from VISTA website

1. VISTA Harness

VISTA harness workflow, from VISTA website

VISTA Harness is an example of this approach. Its hypotheses/world model about the environment (including objects, dynamics, and goals of this environment) is a set of free-form language. As the agent collects new observations, it updates these hypotheses and uses them to guide future reasoning and actions.

VISTA
interaction history
      ↓
hypothesize environment rules
      ↓
free-form natural language
("this object seems to...",
 "clicking X causes Y...", ...)
      ↓
revise hypotheses / reason / act

2. Schema Harness

Schema Harness follows a closely related idea, but represents the inferred world model as a CWM. The program explicitly describes the state representation, transition rules, and goal conditions, and can be executed and verified against interaction history.

Schema
interaction history
      ↓
hypothesize environment rules
      ↓
executable program / CWM
(state + transitions + goal conditions)
      ↓
execute / verify / plan

Advanced topics

WorldCoder assumes a single-agent, fully observable, and deterministic environment. Later CWM works explore how to relax these assumptions.

How to model multi-player environments?

DeepMind's Code World Models for General Game Playing models multiplayer games as extensive-form games. A history \[ h=(a_1,\ldots,a_t) \] records all actions so far and can be viewed as the environment's ground-truth game state. At each history, \(\tau(h)\) specifies which player acts next, \(A(h)\) gives the legal actions. Each player aims to maximize its cumulative reward.

To model randomness, a special chance player represents stochastic events such as card draws or dice rolls

How to model partially observable environments?

When the true state is hidden, there are two natural approaches.

1. Condition directly on interaction history.

PoE-World models future observations from the full history:

\[ p(o_{t+1}\mid o_{1:t},a_{1:t}), \]

avoiding an explicit latent-state estimator.

observation-action history
          ↓
         CWM
          ↓
     p(o_{t+1})

2. Infer a latent state.

DeepMind CWM instead synthesizes both a hidden-state simulator and an inference program:

\[ p_M(s_t\mid o_{1:t},a_{1:t}). \]

The inference program reconstructs a latent state or latent history from observations, which is then used by the CWM.

observation-action history
          ↓
   inference program
          ↓
      latent state
          ↓
         CWM

In the hardest closed-deck setting, the inference program and CWM are synthesized jointly, forming a programmatic encoder-decoder system.

How to model stochastic environments?

One solution is to make the CWM itself probabilistic.

DeepMind CWM uses a particularly clean representation for games: ordinary transition functions remain deterministic, while randomness is represented explicitly by the chance player:

For example, dealing a card is represented as a chance action sampled from a chance-outcome distribution. Thus, stochastic dynamics are reduced to deterministic transitions plus explicit random events.

PoE-World takes a different approach and directly represents a probabilistic world model: \[ p_\theta(o_{t+1}\mid o_{1:t},a_{1:t}) \propto \prod_i p_i^{\mathrm{expert}} (o_{t+1}\mid o_{1:t},a_{1:t})^{\theta_i}. \] Each expert is a small program, while the weighted product of experts defines a distribution over possible next observations. The programs themselves can remain simple and deterministic; PoE-World converts their outputs into distributions and learns the expert weights from interaction data.

Factorizing the CWM

Most early CWMs, including WorldCoder, try to synthesize a monolithic program describing the entire environment. As environments become more complicated, this creates a difficult program-synthesis problem.

PoE-World instead decomposes the world model into many small programmatic experts. Each expert describes one local causal rule, for example:

if player touches skull:
    player dies

if player is on platform and action == RIGHT:
    player.vx = 2

if player touches ladder and action == DOWN:
    player.vy = 4

The complete world model is obtained by combining the predictions of these experts through a weighted product. Importantly, an expert does not necessarily correspond to one object; it represents a local rule or mechanism, which may involve multiple objects.

expert 1 ─┐
expert 2 ─┤
expert 3 ─┼── weighted Product of Experts ──► p(o_{t+1})
   ...    │
expert n ─┘

The LLM synthesizes new experts from observed transitions, while their scalar weights are optimized by maximum likelihood. Experts receiving sufficiently small weights are removed. This gives the system a way to reject a bad local hypothesis without discarding the rest of the CWM.

The trade-off is that composition itself becomes an assumption. PoE-World makes inference tractable by factorizing object attributes and treating them as conditionally independent given the history. Local interactions are easy to express, but strongly coupled mechanisms or global latent variables can entangle many experts and reduce the benefits of modularity.

In short:

Monolithic CWM
environment ──► one large program

Factorized CWM
environment ──► many local programs ──► composition ──► world model

Factorization changes the CWM synthesis problem from writing one increasingly complicated simulator to discovering and composing many reusable causal rules.

Referrences

  1. WorldCoder, a Model-Based LLM Agent: Building World Models by Writing Code and Interacting with the Environment 2024 paper
  2. Code World Models for General Game Playing 2025 paper
  3. PoE-World: Compositional World Modeling with Products of Programmatic Experts 2025 paper
  4. ReAct: Synergizing Reasoning and Acting in Language Models 2022 paper
  5. schema harness 2026 website
  6. VISTA harness 2026 website