PaperlayerView exact revision on arXiv

Source-linked AI summary

Generating Code World Models with Large Language Models Guided by Monte Carlo Tree Search

Nicola Dainese, Matteo Merler, Minttu Alakuijala, Pekka Marttinen

arXiv:2405.15383v2cs.AI

TL;DR

World models built directly with LLMs can be slow, while existing language-conditioned model-based agents have limited language understanding and generalization. The paper introduces code-generated world models and GIF-MCTS, and reports performance across three benchmarks as well as successful downstream planning.

  • Problem

    LLM-based world modeling and language-conditioned model-based RL each face trade-offs in inference speed, language understanding, and generalization, motivating ways to combine their strengths.

    Source: §1 · p. 2

  • Method

    The paper generates Python Code World Models with GIF-MCTS, which iteratively uses generate, improve, and fix actions, and introduces an 18-environment benchmark with descriptions and curated trajectories.

    Source: §1 · p. 2

  • Results

    GIF-MCTS outperforms existing methods on all three benchmarks; synthesized CWMs support planning with inference four to six orders of magnitude faster than direct LLM querying, and match an oracle planner when accurate.

    Source: §1 · p. 3

  • Takeaways & Limitations

    Synthesized CWMs can be used for planning, yielding model-based RL agents with high sample efficiency and inference speed when the models are accurate.

    Source: §1 · p. 3

    The approach assumes deterministic, fully observable environments; stochasticity and partial observability complicate verification, and code-based models may be too rigid when dynamics change.

    Source: §6.1 · p. 10

Abstract

from arXiv · show

In this work we consider Code World Models, world models generated by a Large Language Model (LLM) in the form of Python code for model-based Reinforcement Learning (RL). Calling code instead of LLMs for planning has potential to be more precise, reliable, interpretable, and extremely efficient. However, writing appropriate Code World Models requires the ability to understand complex instructions, to generate exact code with non-trivial logic and to self-debug a long program with feedback from unit tests and environment trajectories. To address these challenges, we propose Generate, Improve and Fix with Monte Carlo Tree Search (GIF-MCTS), a new code generation strategy for LLMs. To test our approach in an offline RL setting, we introduce the Code World Models Benchmark (CWMB), a suite of program synthesis and planning tasks comprised of 18 diverse RL environments paired with corresponding textual descriptions and curated trajectories. GIF-MCTS surpasses all baselines on the CWMB and two other benchmarks, and we show that the Code World Models synthesized with it can be successfully used for planning, resulting in model-based RL agents with greatly improved sample efficiency and inference speed.

1 Introduction

The paper proposes generating world models as code to combine LLM language understanding with fast, precise model-based planning. It introduces GIF-MCTS and a benchmark for evaluating code-generated world models, reporting results across three benchmarks and downstream planning.

1 Introduction

  • The paper motivates model-based RL because task descriptions must convey environmental mechanics—not only instructions—including how actions change observations and whether outcomes meet goals.Natural-language conditioning may help ground observations in prior knowledge and support adaptation.Source: §1 · p. 1
  • The framework iteratively tests candidate CWMs against trajectories from the true environment, feeds prediction feedback to the LLM, and returns the best model for planning.The process ends when all transitions match or the computational budget is exhausted.Source: Figure 1 · p. 2
  • Prior approaches trade off prediction fidelity and language understanding against planning speed, or offer fast planning with limited language understanding and generalization.Directly using an LLM as a world model is also slow at inference and restricted to textual inputs and outputs [Hao et al., 2023].Source: §1 · p. 2
  • The Code World Models Benchmark comprises 18 diverse RL environments, paired with natural-language descriptions and curated trajectories for evaluating code generation across environments of varying complexity.The environments cover discrete and continuous control tasks.Source: §1 · p. 2
  • GIF-MCTS addresses the reasoning, instruction-following, implementation, and debugging demands of synthesizing world-model programs through Generate, Improve, and Fix actions guided by Monte Carlo Tree Search.The method is evaluated on CWMB, APPS Competition, and RTFM.Source: §1 · p. 2
  • GIF-MCTS outperforms existing methods on all three benchmarks, and its synthesized models support planning with inference four to six orders of magnitude faster than directly querying an LLM.When the CWM is accurate, the resulting agent matches an oracle planner with access to the real-world model.Source: §1 · pp. 2–3

2 Related Work

Related work explores code, language models, and search for world modeling and program generation. The paper distinguishes its offline, broader-environment evaluation and code-specific MCTS strategy from nearby approaches.

World models with code

  • Prior code-based world models include AutumnSynth’s functional synthesis with finite-state automata, PDDL action preconditions and effects, and probabilistic-code generation [Wong et al., 2023].PDDL approaches rely on environment predicates and high-level actions that require a low-level language-conditioned controller [Guan et al., 2023; Wong et al., 2024].Source: §2 · p. 3
  • WorldCoder also uses an LLM to generate Python world models, but learns online in two grid-world tasks; this work uses offline data across broader environments, including continuous and discrete spaces.The paper reports state-of-the-art APPS Competition results and superior or on-par CWMB performance to WorldCoder.Source: §2 · p. 3

Code generation with LLMs

  • Prompting methods include Chain of Thought [Wei et al., 2022, Kojima et al., 2022] and tree-based approaches [Yao et al., 2023, Hao et al., 2023]; token-level MCTS becomes impractical for programs hundreds of tokens long [Zhang et al., 2023].Source: §2 · p. 3
  • LATS combines MCTS, self-reflection, and environmental feedback [Zhou et al., 2023], but generates multiple programs from one node, uses a separate reflection prompt, and lacks specialized bug-fixing strategies.Source: §2 · p. 3
  • Parsel decomposes programs into independent pseudocode problems [Zelikman et al., 2023], whereas this work uses sequential MCTS refinement and environmental feedback to generate executable Python for model-based RL.Source: §2 · p. 4

3 Code World Models

The paper formulates Code World Models as executable models of environment transitions and rewards, learned from descriptions and offline trajectories. Its benchmark spans 18 environments with varied difficulty and action and observation spaces.

3 Code World Models

  • The framework assumes deterministic, fully observable environments and a natural-language description detailed enough to infer observation spaces and transition and reward logic.Determinism means the transition is s′ = f(s, a); the authors identify stochasticity and partial observability as limitations.Source: §3 · p. 4
  • In the offline setting, a dataset contains one-step transitions collected under a behavior policy; the authors note it could be lifted by using a CWM and planner to collect more trajectories online.Source: §3 · p. 4
  • CWMB contains 18 RL environments of varying difficulty, including classical control, physics-based PyGame, and MuJoCo tasks, with continuous and discrete action and observation spaces.The environments and documentation are adapted from Gymnasium.Source: §3 · p. 4
  • For each environment, the benchmark provides curated trajectories with both low- and relatively high-scoring behavior, without requiring optimal demonstrations or maximal state-space coverage.For most environments, the collection includes five random-action trajectories and five suboptimal demonstrations above a return threshold.Source: §3 · p. 4
  • The benchmark includes end-user language descriptions, and each trajectory transition serves as an input-output sample for validating generated models.The validation compares predictions with collected transitions.Source: §3 · p. 4
  • The illustrated GIF-MCTS process tests candidate CWMs against environment trajectories, feeding prediction accuracy back to the LLM until transitions match or the computational budget is used.The figure depicts up to three sequential bug-fix attempts per buggy node.Source: Figure 2 · p. 5

4 GIF-MCTS

GIF-MCTS uses sequential MCTS iterations and unit-test feedback to generate, evaluate, and refine complete Python programs for Code World Models. Its code-specialized actions and tree representation support both incremental construction and targeted correction.

4 GIF-MCTS

  • A CWM is a Python Environment class whose step() function jointly implements transition and reward functions; correctness means reproducing all transitions in the training dataset.Accuracy is the uniformly weighted fraction of correctly predicted next-state, reward, and done signals.Source: §4 · p. 5
  • Each MCTS node stores a partial program state and a rollout completion; appending L lines (set to two) forms the child state, while the rollout estimates its expected accuracy.This stores code in multi-line blocks, unlike prior approaches that store full programs or single tokens.Source: §4 · p. 6
  • MCTS selection uses UCT; expansion prompts the LLM for a program, parsing separates state and rollout, evaluation computes accuracy, and backpropagation updates ancestor values.Source: §4 · p. 6
  • GIF-MCTS frames MCTS nodes and actions for long-form code generation with unit tests, adds code-specialized Generate, Improve, and Fix actions, and uses a heuristic to balance exploration and exploitation across actions.The paper reports an ablation study of the action types and empirical improvement from the UCT heuristic.Source: §4 · p. 6

4.1 GIF-MCTS Actions

  • Generate samples full-code continuations from a node’s state to explore possible solutions, while Improve supplies the full program and a failed example with its expected output to revise program logic.Improve also prompts the LLM to explain where the code fails and attempt a fix.Source: §4.1 · p. 6
  • Fix targets syntax or runtime errors by providing the buggy program and error feedback, with temporary node values keeping buggy programs eligible for further repair attempts.Without fixing, programs that fail to execute receive zero reward and may be wastefully abandoned.Source: §4.1 · p. 6

5 Experiments

The experiments compare GIF-MCTS with code-generation baselines on APPS and CWMB, then study planning in a simplified RTFM environment. Additional ablations and qualitative results are reported in appendices.

5.1 Baselines

  • The APPS baseline is Zero-shot CoT; WorldCoder selects prior programs using Beta-distribution samples, then fixes or improves and tests the selected program.The WorldCoder comparison uses the same number of LLM calls as GIF-MCTS.Source: §5.1 · p. 7

5.2 APPS

  • GIF-MCTS outperforms Zero-shot CoT and WorldCoder on APPS Competition and reaches a new state of the art to the authors’ knowledge.Parsel is second-best in accuracy but evaluates exponentially many solutions, whereas GIF-MCTS evaluates only 20 programs.Source: §5.2 · pp. 7–8, Table 1 · p. 8

5.3 Code World Models Benchmark

  • The experiments assess generated-program accuracy and planning performance on CWMB, using MCTS for discrete-action environments and CEM for continuous-action environments.Table 2 reports accuracy and normalized return averaged separately across discrete- and continuous-action environments.Source: §5.3 · p. 8, Table 2 · p. 8
  • GIF-MCTS outperforms WorldCoder across every CWMB environment split and backbone; gains vary by action space, and Llama 3 could not generate executable models for the two hardest environments.GPT-4 successfully generated a model for each benchmark environment.Source: §5.3 · pp. 8–9, Table 2 · p. 8

5.4 Read to Fight Monsters

  • On simplified RTFM, GIF-MCTS significantly outperforms WorldCoder at every setting, but matches ground-truth-simulator performance only when its CWM is perfect.The CWM improves as GIF-MCTS receives more calls; this experiment fixes the manual and keeps monsters stationary.Source: §5.4 · p. 9, Table 3 · p. 9

6 Discussion

The discussion explains gaps between CWM accuracy and planning returns, compares CWMs with offline RL, and identifies assumptions and practical constraints on broader use.

6 Discussion

  • GIF-MCTS generates programs from scratch or partial programs, diversifying its search; a tree-search ablation that only refines complete programs performs worse.Source: §6 · pp. 9–10
  • CWM accuracy exceeds normalized return, matching it only when the model is flawless; missed terminal success rewards can derail planning, and sparse trajectory coverage may also contribute.Source: §6 · p. 10
  • The experiments use 10 trajectories of at most 100 steps, and the authors leave more thorough CWM sample-efficiency experiments to future work.Transitions validate generated code and serve as in-context examples; the curated set should cover edge cases and reward values.Source: §6 · p. 10
  • CWMs compare favorably with CQL on discrete-action environments, while CQL performs better on continuous-action environments, which are harder to model.The authors note that scarce data can lead RL methods, including CQL, to overfit.Source: §6 · p. 10

6.1 Limitations

  • The approach assumes deterministic, fully observable environments; stochasticity and partial observability complicate prediction verification and are left for future work.Source: §6.1 · p. 10
  • CWMs need environment descriptions convertible to Python, may require rewriting as dynamics change, and struggle especially with complex physics-based environments.The authors suggest image-to-text preprocessing, modular rewrites, or external physics tools as possible directions.Source: §6.1 · p. 10
  • GIF-MCTS depends on test cases for code evaluation, which may be difficult to provide in some tasks.Self-generated tests are proposed as a possible alternative but left for future work.Source: §6.1 · p. 10

7 Conclusion

The conclusion presents CWMs as an LLM-based framework for building RL world models and reports code self-debugging and planning demonstrations across environments. It also cautions that generated code should be vetted by a human expert before use.

7 Conclusion

  • Code World Models provide a general framework for using LLMs to build world models for RL agents, with examples of feedback-driven code improvement and downstream planning across environments.Source: §7 · pp. 10–11

A Broader Impact

  • Generated Python code should be checked by a human expert before use; expert-vetted code may offer interpretable models for critical applications.Source: A Broader Impact · p. 16

B Additional GIF-MCTS implementation details

GIF-MCTS uses action-specific tree expansion and value estimates to guide search, with special handling for unexplored and buggy programs. Ablations show that removing any action reduces performance, especially in harder continuous environments.

B Additional GIF-MCTS implementation details

  • Non-buggy nodes offer generate and improve actions, while buggy nodes receive only sequential fix actions.The root has only a generate action because it contains no existing code to improve.Source: B Additional GIF-MCTS implementation… · p. 16
  • The modified UCT selection uses node value, an exploration constant, parent visits, and the count of expanded children with the same action type.The same-action count is intended to discourage trees from growing only horizontally.Source: B Additional GIF-MCTS implementation… · p. 16
  • For unexplored nodes, separate action-specific linear models combine global and local average values, with weights learned by gradient descent.A tuned prior and initial count seed the global average before values are available.Source: B Additional GIF-MCTS implementation… · p. 16
  • Buggy nodes receive a temporary value of 0.99 to encourage a fix attempt, then that value falls to zero after three unsuccessful fixes.Temporary values are excluded from MCTS backtracking; the implementation hyperparameters and their descriptions are reported in Table 4.Source: B Additional GIF-MCTS implementation… · pp. 16–17, Table 4 · p. 17

C Ablation Study on GIF-MCTS

  • Removing any GIF-MCTS action reduces performance, most notably in harder continuous environments, while discrete-environment results have greater statistical uncertainty.The authors identify fixing bugs as the most important action and improving already-working code as the least impactful.Source: C Ablation Study on GIF-MCTS · p. 17

D Qualitative Study

The qualitative study compares action frequencies across the search tree with those along the path to the highest-value node. It finds a generate-action bias on that path and identifies ineffective improve actions as a potential bottleneck.

D Qualitative Study

  • GIF-MCTS often chooses generate actions on the path to its highest-value node, while improve actions—used more than 35% of the time in most cases—much less often lead to that node.Table 7 separates frequencies across the whole tree from those on the root-to-highest-value path.Source: D Qualitative Study · p. 18, Table 7 · p. 18
  • Early generate actions often produce values near the tree’s maximum, while later improve actions may yield marginal gains or fail to improve the program.The same-action penalty in modified UCT can lead to later improve actions; many may be needed before one raises performance.Source: D Qualitative Study · p. 18
  • Improve actions that leave accuracy unchanged or make it worse are a potential bottleneck, and stronger models may be needed to analyze and criticize code.The authors suggest pairing a code-explaining model with a code-generating model.Source: D Qualitative Study · p. 18
  • Fix actions are rare in easier discrete-action problems but become more prominent in harder continuous-action problems.For GPT-4 Turbo, the authors report that only generate actions are needed to synthesize perfect or near-perfect models in the discrete settings described.Source: D Qualitative Study · p. 19

E Data Contamination

The paper examines possible data contamination in its benchmarks but says the concern cannot be dismissed outright. It argues that using the same underlying models for GIF-MCTS and its baselines prevents potential contamination from biasing their reported comparison.

E Data Contamination

  • For APPS, Kattis supplies 691 problems without online reference solutions, yet it has the highest competition-split strict accuracy rate at 0.347.All methods and models have a 0% pass rate on the 41 AtCoder problems whose reference solutions are available online.Source: E Data Contamination · p. 19
  • The authors note that APPS performance correlates with problem source websites, not with reference-solution availability, and suggest source difficulty matters more.This is presented as evidence addressing, but not eliminating, contamination concerns.Source: E Data Contamination · p. 19
  • The authors cannot rule out data contamination outright, but compare GIF-MCTS with baselines using the same underlying models to avoid bias from potential contamination.For CWMB, generated step-function logic remains distinct from reference models, and the modified RTFM environment has no online reference solution.Source: E Data Contamination · p. 19

F Data Quality

The paper tests whether language descriptions can help synthesize accurate models from limited trajectories, and situates its approach among related work on language-based world models and programmatic RL. It also reports that calling a CWM is substantially faster than directly querying an LLM for predictions.

F Data Quality

  • The CWMB dataset includes both low- and high-scoring trajectories, and the authors argue that language descriptions could support generalization beyond observed transitions.They describe this possibility as theoretical and demonstrate it with the RTFM experiment.Source: F Data Quality · pp. 19–20
  • With only 10 failure trajectories and no observed +1 reward, a GPT-4 CWM synthesized with 50 calls reached 100% accuracy on the new dataset and predicted unseen positive transitions.It also scored 100% on the original dataset, matching the model generated with the full range of data.Source: F Data Quality · pp. 19–20

G Additional Related Work

  • The paper describes language-based world models that predict text, image representations, or pixels, alongside programmatic RL approaches that represent policies or models as programs.Related approaches include Dynalang, Voltron, PIRL, LEAPS, and program-generating agents such as Voyager.Source: G Additional Related Work · p. 20

H Comparison of Inference Times

  • Calling a CWM’s Python step function is four orders of magnitude faster than GPT-4 Turbo on the easiest selected environment and seven orders faster on the hardest.None of GPT-4 Turbo’s next-observation predictions in this comparison were accurate.Source: H Comparison of Inference Times · p. 20

I Code World Models Benchmark Details

The appendix details how CWMB environment descriptions and statistics were compiled, gives inference-time and environment-level results, and reports per-environment accuracy and return for CWMs generated with two LLMs.

I Code World Models Benchmark Details

  • CWMB descriptions omit documentation unrelated to environment definitions, and environment descriptions were parsed from Gymnasium v.0.29.1.Reported code lengths and token counts use GPT-4 Turbo CWMs generated with GIF-MCTS at budget 10 as a general indication; they can vary.Source: I Code World Models Benchmark Details · pp. 20–21
  • Table 8 compares GPT-4 and CWM inference times using a sample of 10 transitions from the replay buffer used during GIF-MCTS.Source: Table 8 · p. 21
  • Table 9 gives detailed statistics for each CWMB environment and marks discrete action or observation spaces with their number of choices.Source: Table 9 · p. 21

J Results for Individual Environments

K Comparison with Offline RL

This section compares CWMs with CQL across offline RL tasks, while also showing fragments of generated game-program logic. The comparison frames its aim as validating the CWM framework rather than maximizing offline-RL performance.

K Comparison with Offline RL

  • CWMs are more suited to discrete tasks, while CQL outperforms them in complex physics tasks.Table 12 compares average raw rewards over 10 episodes for random, CQL, GIF-MCTS CWM planners using Llama 3 and GPT-4, and an oracle planner.Source: K Comparison with Offline RL · p. 21, Table 12 · p. 23
  • CWMs remain competitive on Pendulum-v1 and Reacher-v4, and to a lesser extent Ant-v4, Pusher-v4, and HalfCheetah-v4; CQL can overfit immediately, even performing worse than random on Pendulum-v1.The authors attribute CQL’s overfitting as likely due to the small dataset, while noting that few trajectories can suffice to validate a CWM.Source: K Comparison with Offline RL · p. 22
  • The authors prioritize validating the framework over outperforming offline-RL methods, using simple planners with default parameters and leaving performance improvements to future work.One suggested direction is allowing generated code to call a physics simulator in continuous environments.Source: K Comparison with Offline RL · p. 23

O Examples of Generated Programs

  • The code fragments show conditional game-step logic for monsters, inventory, item modifiers, and game completion, then return a board copy, inventory, and valid actions.The excerpts include a `step(self, action)` method and checks involving `pos`, `monster`, `self.inventory`, `content_parts`, and `game_done`.Source: O Examples of Generated Programs · p. 39
  • One excerpt returns `(self.board.copy(), self.inventory)` together with a NumPy array of valid actions.Source: O Examples of Generated Programs · p. 40
Loading 2405.15383v2…