WorldCoder, a Model-Based LLM Agent: Building World Models by Writing Code and Interacting with the Environment

1 Introduction

• WorldCoder synthesizes a Python program from past experience to model transitions, then plans using that learned world model.The LLM builds a simulation rather than simulating the world or performing the planning itself.Source: §1 · pp. 1–2
• Code can reuse prior knowledge, support sample-efficient transfer across tasks, and make the agent’s knowledge auditable because programs are human-interpretable.World models can also support novel goals by changing the reward function.Source: §1 · p. 1
• Program search and efficient exploration are the central challenges; the proposed objective favors models whose planner expects rewarding states, particularly under uncertainty.Curriculum learning reuses successful program pieces to make synthesis across environments more tractable.Source: §1 · p. 1
• The architecture is presented as more sample-efficient than deep RL and more compute-efficient than prior LLM agents, while supporting transferable, interpretable models.Source: §1 · p. 2
• The optimism objective generates goal-driven exploration and can reduce environment interactions needed to obtain reward by orders of magnitude.The paper reports theoretical and empirical support for this objective.Source: §1 · p. 2
• The transfer analysis finds that code can be quickly adapted to new world dynamics in gridworld and robot task-planning domains.Source: §1 · p. 2

2.1 Problem statement and core representations

• Goals are natural-language contexts in a deterministic episodic MDP with fixed dynamics across goals; Python separates transition code from context-conditioned reward code.The MDP terminates upon reaching the goal.Source: §2.1 · pp. 2–3

2.3 Understanding optimism under uncertainty

• Optimism prefers data-consistent models whose planner can reach reward, prompting exploratory goals and allowing new natural-language goals to trigger reward-model updates.For a new context, the reward model is debugged to cover it while allowing a goal state to be reached.Source: §2.2 · p. 3, §2.3 · pp. 3–4

2.2 The world model learning problem

• Candidate programs must replay observed transitions and rewards and support a positive-reward plan from recorded starts; the LLM repairs candidates against inconsistent data.The replayed experiences include state transitions, rewards, context, and whether an episode ended.Source: §2.2 · p. 3

2.3 Understanding optimism under uncertainty

• Under realizability and optimal planning, at most D_{S,A,T} × (K_{T×R} + 1) actions suffice to reach positive reward.The bound assumes the true transition and reward functions are in the agent’s model class.Source: §2.3 · p. 4

2.4 Overall architecture

• After initial random data collection and synthesis, the agent plans with depth-limited value iteration in simple domains or MCTS in more complex ones.The architecture first initializes an interaction dataset through random exploration.Source: §2.4 · p. 5

2.5 Program Synthesis via Refinement

• LLM refinement repairs old programs against unexplained transitions for transfer; REx selects candidates by constraint progress while balancing rarely refined programs against promising ones.The heuristic measures replay-buffer consistency with both learning constraints.Source: §2.5 · p. 5

Sokoban

• 50 actions suffice for a human-understandable Sokoban model that generalizes to more boxes; deep RL takes >1 million steps on 2-box Sokoban.WorldCoder does not solve very hard levels, such as those with 5+ boxes; the authors attribute this to planning difficulty.Source: Figure 3 · p. 6, §3 · p. 6
• 15% ± 8% of basic levels are solved by ReAct, while WorldCoder front-loads 400k LLM tokens and then acts without later LLM queries.ReAct receives the same pretrained Sokoban knowledge but cannot effectively play the game.Source: Figure 3 · p. 6, §3 · pp. 6–7
• With teleport gates that deactivate when blocked by a box, WorldCoder learns the altered dynamics and uses teleporters to solve levels faster.The authors introduce these extra rules to test whether pretrained Sokoban knowledge alone explains the agent’s performance.Source: §3 · p. 7

MiniGrid

• Transfer reduces exploration in new MiniGrid environments; without transfer, optimism is needed to solve harder environments and generalize zero-shot to new goals.The curriculum distinguishes learning new object dynamics from following new natural-language goals.Source: Figure 4 · p. 7, §3 · p. 7
• A prompt-for-reward ablation struggles to follow new instructions, showing the importance of optimism in correcting LLM-generated reward-function mistakes.The comparison includes transfer from environment 2 to environment 3.Source: §3 · pp. 7–8

AlfWorld

• On AlfWorld, the optimism objective is necessary for nonzero performance; the learned model contains 250+ lines of code and typically earns reward after around 20 exploratory steps.The model learns object pickup and placement, then appliance use; later episodes make small model corrections.Source: Figure 5 · p. 8, §3 · p. 8

World models

• World models underpin model-based RL and robotic task planning, but practical models have required large training sets or careful hand-engineering; neurosymbolic models can learn from raw perception at the expense of transfer and sample complexity.The paper suggests combining these approaches may bring together their strengths.Source: §4 · p. 8

LLMs and world models

• Unlike work that learns a rich but frozen world model from many examples or generates planning operators from language-described affordances, WorldCoder builds a simple model on-the-fly through environmental interaction.The authors describe language-based and interaction-based approaches as potentially complementary.Source: §4 · p. 9

LLMs for decision-making

• ReAct [77] and related agents use an implicit neural world model, whereas WorldCoder builds an explicit one.ReAct prompts an LLM to reason step-by-step and predict an action.Source: §4 · p. 9

Programs as policies

• Program policies can ignore world details irrelevant to action, but they do not readily generalize to new goals as world models used by planners can.Recent LLM work uses general-purpose languages such as Python to express policies as programs.Source: §4 · p. 9

Programs as world models

• WorldCoder extends program-synthesis world models to Python, uses LLMs as general-purpose synthesizers, and supports goal-directed behavior and transfer across environments and goals.The authors say this avoids the more intricate synthesis algorithms used in prior work.Source: §4 · p. 9

Exploration and optimism

• Optimism under uncertainty is formulated as a logical constraint between a program and a planner, for use with LLM-guided discrete program search.RLLang uses source code to specify parts of policies and world models, and could be combined with WorldCoder to integrate prior knowledge.Source: §4 · p. 9

Deterministic dynamics

• The current method assumes deterministic dynamics; the authors suggest probabilistic-program synthesis as a possible extension.They describe this limitation as nontrivial but potentially surmountable.Source: §5 · p. 9

State representation

• The approach gains sample efficiency and transferability by representing knowledge as code, but requires symbolic, discrete states rather than messy, continuous real-world states.Pretrained object detectors, multimodal models [30], and neurosymbolic programs [64] are suggested as possible ways to bridge perception and symbol processing.Source: §5 · pp. 9–10

Growing and transferring knowledge

• Knowledge growth and transfer currently use a basic mechanism; the authors identify reusable libraries of subroutines and classes as a future direction.Such components could be shared across environments and goals.Source: §5 · p. 10

Empirical illustration

• At most 100 actions let WorldCoder with optimism learn the correct UnlockPickup world model from scratch, while PPO failed to gain enough positive-reward supervision after 3 × 10^8 actions and WorldCoder without optimism also failed.The authors report this as evidence that the optimism objective is more sample-efficient than traditional data consistency alone.Source: A Theoretical analysis of Sample… · p. 16

Sample-complexity bound

• The maximum actions to achieve a positive-reward goal are bounded by D_{S,A,T} × (K_{T×R} + 1), assuming an optimal planner, an episodic MDP, and a true MDP within the model class.The bound is polynomial in the state-space diameter and the model space's logical dimensionality.Source: A Theoretical analysis of Sample… · p. 16

Counterexample-driven learning

• With a model satisfying data consistency and optimism, planning either reaches the goal or finds a counterexample; unexplained counterexamples progressively shrink the set of potentially correct models.The analysis bounds the needed counterexamples by K and the actions to find each by the state-space diameter.Source: A Theoretical analysis of Sample… · p. 16

Exploration efficiency

• For deterministic environments, an optimal planner finds a counterexample within D_S actions under optimism, whereas random exploration can require exponentially more actions in the worst case.The bound does not assume free resets to all possible states.Source: A Theoretical analysis of Sample… · p. 16

Logical dimensionality

• The bound uses logical dimensionality K, the maximum size of a mutually independent dataset; the authors note this is a loose upper bound and can be much smaller than the solution-space size.For linear models, n independent data points suffice to characterize the correct solution in R^n.Source: A Theoretical analysis of Sample… · p. 16

A.1 Example learning trajectory using (ϕ1 ∧ϕ2) MiniGridUnlockPickup

• Optimism lets the agent imagine action and tool functionality, explore using that model, and correct it with newly collected data.Source: §A.1 · pp. 17–18
• Under data-consistency alone, random exploration had figured out rotation but had no feasible move yet.Source: §A.1 · pp. 17–18
• Early imagined plans ignored the locked door and attempted pickup at the agent’s location; a later plan correctly included pickup of an object in front.Source: §A.1 · pp. 19, 21, 22
• Data-consistent updates established that the agent could not step onto walls, locked doors, or keys, though its route to the goal remained unknown.Source: §A.1 · pp. 20, 22
• The goal-driven model succeeded after learning that pickup is impossible while carrying an object and improving its drop behavior.The final model also included checking for obstacles between the agent and the target item.Source: §A.1 · p. 24

• Facing items are pickable unless they are doors or walls; picking one sets it as the agent’s carried item and removes it from the state.Source: §A.1 · p. 26
• Turn-left and turn-right actions toggle the door in front of the agent.Source: §A.1 · p. 26
• The agent moves to a target position only when it contains no wall, box, ball, lava, key, or closed door.Source: §A.1 · p. 26

B Planning with Monte Carlo Tree Search and the BM25-based Heuristic

• BM25 guides planning in text-based environments, while MCTS trades off exploration and exploitation in deterministic settings with sparse rewards.The heuristic encourages trajectories closer to the user-specified goal; MCTS lets the planner explore alternatives after promising paths fail for too long.Source: B Planning with Monte Carlo Tree… · p. 26
• The heuristic steers the planner toward trajectories that are closer to the goal specified by the user.Source: B Planning with Monte Carlo Tree… · p. 26
• MCTS allows exploration of other options after seemingly promising trajectories fail to make progress for too long.Source: B Planning with Monte Carlo Tree… · p. 26

B.1 Monte Carlo Tree Search for Deterministic Environments with Sparse Rewards

• Algorithm 2 searches within a computation budget, selects nodes with a tree policy, expands untried actions, and backs up heuristic values.Source: §B.1 · p. 27
• UCT-based MCTS balances exploration and exploitation by selecting parts of the search space and evaluating trajectories that share the selected partial trajectory's prefix.The standard MCTS description uses simulations to evaluate these continuations.Source: §B.1 · pp. 26–27
• The method removes rollouts and substitutes a heuristic for deterministic sparse-reward environments, redirecting computation to smarter online tree expansion.The authors argue that rollouts are not cost-effective when there is no randomness or partial reward: returns remain zero until a goal is found.Source: §B.1 · p. 27

B.2 BM25 as a Heuristic for Planning in Text-based Environments

• The heuristic estimates a selected state's promise by comparing the user-specified goal text with the text of the partial trajectory leading to that state.Source: §B.2 · p. 27
• The heuristic assumes that greater similarity between a partial trajectory and the goal indicates a more goal-relevant state; for example, matching terms such as “alarmclock,” “desk,” and “hot” make a trajectory seem promising for putting a hot alarm clock in a desk.Source: §B.2 · p. 28
• The authors use BM25 for efficiency and maintain an online trajectory corpus during search rather than a static pre-collected corpus.Source: §B.2 · p. 28
• BM25's saturating term-frequency formulation significantly outperforms linear TF-IDF in preliminary experiments.Without saturation, repetitive long trajectories can receive excessive scores and trap the planner in local minima until its budget is exhausted.Source: §B.2 · p. 28

C.1 PPO for MiniGrid

• PPO is much more sample-inefficient than the agent, learning valid policies in most environments after 104 −105 steps but failing to solve UnlockPickup within 3 × 108 steps.The comparison uses Figure 6 and tuned RL Baselines3 Zoo hyperparameters, with the same symbolic memory state as the agent.Source: Figure 6 · p. 28, §C.1 · p. 29
• The online BM25 heuristic tracks trajectory counts and average trajectory length, then parses mission text into tokens for scoring.Its defaults are k1 = 1.5 and b = 0.75.Source: §C.1 · p. 29
• PPO is evaluated as a deep-RL baseline using tuned RL Baselines3 Zoo hyperparameters and the agent’s symbolic memory state.The passage reports valid policies in most environments after 104 −105 steps; it does not specify that these are perfect policies.Source: §C.1 · p. 29

C.2 PPO for Sokoban

• The PPO baseline uses Stable Baselines3 and gym-sokoban with 256 parallel environments, batch_size 2048, and a 50-step horizon.The remaining listed hyperparameters follow Stable Baselines3 defaults.Source: §C.2 · p. 29
• The network input is a (3, 7, 7) representation of the 7-by-7 Sokoban grid, with each grid item represented as an RGB pixel.Source: §C.2 · p. 29
• The policy uses three convolutional layers with 16, 32, and 64 filters, followed by a linear layer with nine action outputs.All layers use 2-by-2 kernels, stride 1, and ReLU activations.Source: §C.2 · p. 29

C.3 DreamerV3 for Sokoban

• The DreamerV3 baseline uses the simple training script provided by the open-source DreamerV3 library.Source: §C.3 · p. 29
• The baseline is run with the gym-sokoban library.Source: §C.3 · p. 29
• The baseline uses default hyperparameters and one environment.Source: §C.3 · p. 29

D Example of synthesized world models for Sokoban

• The transition function maps the four movement actions to coordinate changes and leaves the player in place when a wall blocks the move.Source: D Example of synthesized world models… · p. 30
• When a player pushes a box, the transition function moves both unless a wall or another box blocks the box’s destination.Source: D Example of synthesized world models… · p. 30
• The reward function gives a positive reward when a box newly reaches a target and a negative reward when a box leaves one.It also checks whether all boxes are on targets.Source: D Example of synthesized world models… · pp. 30–31

• Entities store location, containment, and world-state and affordance attributes such as hot, coolable, openable, and pickupable.Source: E Example of synthesized world model… · p. 31
• Typed action classes encode commands including goto, open, close, examine, use, pickup, and put, with relevant entity names as arguments.Source: E Example of synthesized world model… · p. 32
• Pickup requires an object to be pickupable and at the agent’s location, while placement requires the agent to hold the object and refuses a closed openable receptacle.Source: E Example of synthesized world model… · pp. 33–34
• Transitions copy the state before dispatching actions; cooling a coolable object opens a closed receptacle, releases the agent’s hold, and marks it cool and not hot.Source: E Example of synthesized world model… · pp. 35–36
• A put action earns 1 and ends the episode only if an egg is found in the fridge and is hot; otherwise the function returns 0 and False.Source: E Example of synthesized world model… · p. 36

F Prompts

• The prompts expose collected replay-buffer data and code synthesized by earlier LLM calls as dynamic context, alongside task instructions.The appendix distinguishes dynamic information from the main instruction.Source: F Prompts · p. 37

F.1 Initializing the transition function

• Transition initialization asks the LLM to model seven uniformly randomly sampled replay-buffer experiences as a transition function.The experiences are state, action, and next-state tuples, with natural-language summaries of state changes.Source: §F.1 · p. 37
• The transition prompt requires runnable Python code implementing the template’s main transition function to return the next state from state and action inputs.It allows helper functions and provides example interactions to guide the implementation.Source: §F.1 · pp. 40–41

F.2 Initializing the reward function

• Reward initialization asks the LLM to model seven uniformly randomly sampled replay-buffer experiences for mission c as a reward-and-done function.The experiences include state, action, next state, reward, and done values.Source: §F.2 · p. 41
• The reward prompt requires runnable Python code implementing reward_func to return reward and done from state, action, and next-state inputs.Its example experiences concern the mission “use the key to open the door and then get to the goal.”Source: §F.2 · pp. 41, 44, 45

F.3 Refining the transition function

• Transition refinement supplies a failed experience, successful examples, and the partial code’s wrong prediction or execution error to guide code correction.The revised code is expected to model all experiences and return the next state from state and action inputs.Source: §F.3 · p. 45
• The transition prompt illustrates a turn-left failure: the expected agent direction changes, but the implementation returns the unchanged state.The prompt asks the LLM to explain and fix this discrepancy.Source: §F.3 · pp. 48–49
• Transition refinement contrasts failed and successfully modeled experiences, asks the LLM to diagnose code–environment differences, and requires a runnable transition(state, action) function.The prompt asks which code caused the mismatch and how to fix it.Source: §F.3 · pp. 47, 49

F.4 Refine the reward function

• Reward refinement gives LLMs a failed data point, successful examples, and the partial code’s wrong prediction or execution error as revision input.The revised code must run on (state, action, next_state) and return (reward, done).Source: §F.4 · p. 49
• The reward-refinement examples test whether the code distinguishes a no-op toggle’s correct 0.0 reward from its predicted -0.1 reward.The prompt asks the LLM to explain the discrepancy and revise its reward function.Source: §F.4 · p. 52

F.5 Generating reward functions for new goals

• For a new goal, the prompt asks the LLM to implement a reward function using sample code synthesized for previous goals and the new mission description.The requested function maps (state, action, next_state) to (reward, done); the example mission is “pick up the yellow box.”Source: §F.5 · p. 52
• The new-goal prompt includes an example mission to pick up a grey box and a separate prompt for a purple-box mission, alongside their code examples.These examples accompany the request to adapt reward-function code to a new goal.Source: §F.5 · pp. 52–53
• The green-box reward example assigns reward 1.0 and ends the episode when the next state has the agent carrying a green box.Source: §F.5 · p. 55
• For a new “pick up the yellow box” mission, the prompt asks the LLM to analyze when reward is positive and when the episode ends, then implement the reward function using the supplied template.The code must run on (mission, state, action, next_state) and return (reward, done).Source: §F.5 · p. 56

F.6 Refine to satisfy optimism under uncertainty

• The prompt supplies valid actions and defines goal achievement as reward > 0 and done=True.The listed actions are turn right, nothing, move forward, turn left, toggle, drop, and pick up.Source: §F.6 · pp. 56, 58
• The refinement must provide complete, compilable, runnable transition(state, action) and reward_func(state, action, next_state) functions.The transition returns the next state; the reward function returns (reward, done).Source: §F.6 · pp. 56, 58
• The example transition model assumes the agent can toggle regardless of what is in front, because the provided experiences do not rule out that ability.Source: §F.6 · p. 57
• The refinement example starts with a world model that fails the “use the key to open the door and then get to the goal” mission from a state containing a yellow key and locked yellow door.The depicted state also contains a goal at (3, 3).Source: §F.6 · p. 58
• The prompt asks the LLM to explain the failure, identify necessary subgoals, check whether the model permits them, and improve the code while retaining its original logic.The model is described as partially correct: it models some experiences but fails to model logic needed to achieve the goal in another environment.Source: §F.6 · pp. 56, 58