Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Hao Tang, Darren Key, Kevin Ellis

TL;DR

WorldCoder proposes a novel framework for model-based RL where a world model is a Python program synthesized by an LLM from environment interactions. The agent uses an optimism-under-uncertainty objective to drive exploration, enabling sample-efficient learning and transfer across environments by editing code. It demonstrates strong performance on grid-worlds and AlfWorld with reduced computational costs compared to deep RL baselines and ReAct-style agents. The work highlights the benefits of explicit, auditable, and transferable world models, and discusses limitations such as assuming deterministic dynamics and symbolized states.

Abstract

We give a model-based agent that builds a Python program representing its knowledge of the world based on its interactions with the environment. The world model tries to explain its interactions, while also being optimistic about what reward it can achieve. We define this optimism as a logical constraint between a program and a planner. We study our agent on gridworlds, and on task planning, finding our approach is more sample-efficient compared to deep RL, more compute-efficient compared to ReAct-style agents, and that it can transfer its knowledge across environments by editing its code.

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

TL;DR

WorldCoder proposes a novel framework for model-based RL where a world model is a Python program synthesized by an LLM from environment interactions. The agent uses an optimism-under-uncertainty objective to drive exploration, enabling sample-efficient learning and transfer across environments by editing code. It demonstrates strong performance on grid-worlds and AlfWorld with reduced computational costs compared to deep RL baselines and ReAct-style agents. The work highlights the benefits of explicit, auditable, and transferable world models, and discusses limitations such as assuming deterministic dynamics and symbolized states.

Abstract

We give a model-based agent that builds a Python program representing its knowledge of the world based on its interactions with the environment. The world model tries to explain its interactions, while also being optimistic about what reward it can achieve. We define this optimism as a logical constraint between a program and a planner. We study our agent on gridworlds, and on task planning, finding our approach is more sample-efficient compared to deep RL, more compute-efficient compared to ReAct-style agents, and that it can transfer its knowledge across environments by editing its code.
Paper Structure (31 sections, 3 theorems, 11 equations, 5 figures, 3 algorithms)

This paper contains 31 sections, 3 theorems, 11 equations, 5 figures, 3 algorithms.

Table of Contents

  1. Introduction
  2. Methods
  3. Problem statement and core representations
  4. The world model learning problem
  5. Understanding optimism under uncertainty
  6. Overall architecture
  7. Program Synthesis via Refinement
  8. Experimental Results
  9. AlfWorld.
  10. Related Work
  11. Limitations and Open Directions
  12. Acknowledgements.
  13. Theoretical analysis of Sample Efficiency using optimism under uncertainty ($\phi_2$)
  14. Observation.
  15. Example learning trajectory using ($\phi_1\wedge\phi_2$) MiniGridUnlockPickup
  16. ...and 16 more sections

Key Result

Theorem 2.4

Assume an episodic MDP $(S, A, T, R, \gamma)$. Assume an agent acting according to an optimal planner operating over world model $(\hat{T},\hat{R})\in \mathcal{T}\times \mathcal{R}$ satisfying $\phi_1\wedge\phi_2$, and that the true MDP is in the agent's model class: $T\in \mathcal{T}$ and $R\in \ma

Figures (5)

  • Figure 1: Overall agent architecture. The agent also inputs a goal in natural language
  • Figure 3: (A) Sokoban domain (per-step reward of -0.1 ellided from figure). (B) Learning curves. ReAct has the same pretrained knowledge of Sokoban but cannot effectively play the game. (C) Our method has different asymptotic LLM cost compared to prior LLM agents, which consume LLM calls/tokens at every action. (D) Deep RL takes >1 million steps to learn 2-box Sokoban. (E) Nonstandard Sokoban with teleport gates
  • Figure 4: (A) Minigrid environments ordered into a curriculum that tests different kinds of transfer learning. (B) Transfer learning performance, compared with (C) performance when solving each environment independently. Appendix Fig. \ref{['fig:ppo-minigrid-results']}: deep RL comparison.
  • Figure 5: AlfWorld environments and tasks. Each learning curve shows average reward at each episode (solid line) and how many steps the episode took (dashed), averaged over 3 seeds. Curves annotated with arrows and text explaining what was learned at each episode. Optimism objective is necessary for any non-zero performance.
  • Figure 6: Performance of PPO in MiniGrid environments.

Theorems & Definitions (8)

  • Definition 2.1
  • Definition 2.2
  • Definition 2.3
  • Theorem 2.4
  • Definition A.1
  • Definition A.2
  • Lemma A.3
  • Theorem A.4