Title: Generating Code World Models with Large Language Models Guided by Monte Carlo Tree Search
Venue: NeurIPS 2024
Accepted: yes
Keywords: Large Language Models

code generation

MCTS

model-based reinforcement learning
Abstract: 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.

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

Abstract

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 ability to model the world is essential for goal-oriented intelligent agents [Ha and Schmidhuber, 2018]. When faced with a novel environment, the agent must quickly understand its mechanics to achieve its goal, for example by building an internal representation of the world and planning with it. In this context, natural language conditioning can be useful for grounding current observations in past knowledge and improving the agent's understanding of the world. Therefore, communicating information about a new task to the agent in natural language is particularly promising, and multiple works explore instruction-following agents [Jang et al., 2022, Ahn et al., 2022]. However, not all important information can be communicated in the form of imperative instructions. Many key facts required to solve a task involve understanding observations, predicting outcomes of different actions and determining whether those outcomes align with the agent's goals. Thus, systems capable of leveraging additional descriptive information, such as model-based Reinforcement Learning (RL) agents, have a greater potential for fast and efficient adaptation via natural language [Lin et al., 2024].

Matteo Merler ∗

Pekka Marttinen

Code generation

with GIF-MCTS

Candidate

Can predict

correctly?

Environment and task

description

Feedback on

wrong predictions

class Environment:

def __init__(self):

se

lf.cart_position =

...

Trajectories collected

from the environment

No

Planning with the

Code World Model

Yes

Large Language Models (LLMs) have revolutionized the field of Natural Language Processing, and offer great opportunities for world modeling, thanks to their internet-scale knowledge, reasoning, and instruction-following abilities. However, it is not clear how to best combine LLMs and world models. One option is multi-modal systems such as text-to-video models [Gupta et al., 2023], which present the highest prediction fidelity, language understanding and out-of-distribution generalization for generation tasks, yet they are too slow to be called repeatedly in a planning loop due to their high inference cost. On the other hand, language-conditioned model-based RL agents [Dainese et al., 2023, Lin et al., 2024] are typically fast at planning and easily trainable. However, they cannot conveniently incorporate LLMs because of their specialised architectures and as such have poor language understanding and generalization capabilities. Other works, such as [Hao et al., 2023], perform planning using an LLM as a world model directly, but they are slow for inference and restricted to textual inputs and outputs, limiting their applicability in RL.

In this study we propose to model the world with code, rather than directly predicting the future with an LLM, which is known to be costly, slow and unreliable. In contrast, code is precise, fast, reliable and interpretable. We thus introduce Code World Models (CWMs), a novel approach to generate RL world models by writing Python code with an LLM, for which a high-level overview can be seen in Figure 1. The concept of CWMs has been independently and contemporaneously proposed by Tang et al. [2024b]; however, our method is technically distinct (Section 2) and scales to more complex world models (Section 5). Alongside this paradigm, we introduce the Code World Models Benchmark (CWMB), consisting of 18 diverse RL environments for discrete and continuous control, paired with corresponding natural language descriptions and curated trajectories. This benchmark aims to facilitate the accurate synthesis of Code World Models through learning from the provided data and evaluate different code generation methods across environments of varying complexity.

Synthesizing programs for world models requires complex reasoning, precise instruction following, accurate implementation of the environment dynamics and reward functions, as well as coding skills for debugging and refining long programs using unit tests. To meet these challenges we propose Generate, Improve and Fix with Monte Carlo Tree Search (GIF-MCTS), a new code generation method based on Monte Carlo Tree Search (MCTS, Kocsis and Szepesvári [2006]) for LLMs, especially suited for generating Code World Models. 2 We evaluate the performance of our method on three benchmarks: the new CWMB, the Competition split on APPS [Hendrycks et al., 2021], a popular and challenging coding benchmark, and RTFM [Zhong et al., 2020], a language-conditioned grid-world, showcasing environments with varying characteristics and complexity. GIF-MCTS

outperforms existing methods on all three benchmarks. Moreover, we demonstrate successful planning in several environments using the synthesized CWMs. This results in model-based RL agents with exceptional sample efficiency and inference speed (from four to six orders of magnitude faster compared to directly querying an LLM as a world model, as shown in Appendix H), while, provided the CWM is accurate, matching the performance of an oracle planner with access to the real-world model. Finally, we discuss the limitations and challenges to overcome to make Code World Models more broadly applicable.

2 Related Work

World models with code. Code is a promising choice for predictive world models thanks to its fast inference, exact syntax and interpretable behavior. However, code alone often struggles to cover the entire scope of the environment's dynamics and previous works often uses different techniques to build a full world model. AutumnSynth [Das et al., 2021] uses a custom programming language named Autumn and integrates a functional synthesis step with a synthesized finite-state automata to model any latent variable. Another popular choice is the Planning Domain Definition Language (PDDL) [Ghallab et al., 1998], which expresses actions as a set of preconditions and effects on the environment. However, PDDL approaches, as in the works by Guan et al. [2023] and Wong et al. [2024], are reliant on having access to predicates about the environment and plan in terms of high-level language actions, which need a low-level language-conditioned controller to be carried out. LLMs have also been used to generate a model based on probabilistic code [Wong et al., 2023].

Most similar to our approach, the concurrently proposed WorldCoder 3 [Tang et al., 2024b] also leverages LLMs to generate a Python-based world model. WorldCoder chooses a program to refine from a working set of programs using the classical Thompson Sampling bandit algorithm [Thompson, 1933, Katehakis and Veinott, 1987], informed by a Beta prior, to iteratively learn a world model from gathered experience. Tang et al. focus on learning world models from online interactions with the environment in two grid-world tasks and on transferring knowledge across variants of the same task. We instead consider a broader selection of environments, propose to learn from offline data, and handle continuous state and action spaces in addition to discrete worlds. Furthermore, we rigorously benchmark and ablate our code generation method, GIF-MCTS, achieving state-of-the-art results on the Competition split of the APPS coding benchmark, and obtain superior or on par performance to WorldCoder on CWMB.

Code generation with LLMs. Current state-of-the-art code generation methods all employ LLMs. While improvements to this task can come from both advancements in the LLMs' coding abilities and enhancements in prompting strategies to guide LLM decoding, the latter is the most relevant to our work. A host of prompting techniques have shown how to leverage the In-Context Learning (ICL) [Brown et al., 2020] abilities of LLMs to enhance a model's reasoning skills, and, as a result, the quality of generated programs. Perhaps the most influential of these is Chain of Thought (CoT) [Wei et al., 2022, Kojima et al., 2022], which leverages in-context examples to encourage intermediate reasoning steps. Tree-like approaches based on the CoT method have also been presented [Yao et al., 2023, Hao et al., 2023]. The work by Zhang et al. [2023] proposes to guide the LLM generation with an MCTS method based on the feedback from unit tests. However, the method considers every token decoded by the LLM as an action in the MCTS tree, which becomes impractical when we have hundreds of tokens per program.

Most similar to our method, LATS [Zhou et al., 2023] uses an MCTS-based generation strategy that incorporates both self-reflection [Madaan et al., 2023, Shinn et al., 2023, Gou et al., 2024] and feedback from the environment. While LATS is broadly applicable to reasoning tasks, it has limitations in code-specific applications like ours. For instance, it generates n programs simultaneously from the same node, rather than sequentially, which does not fully exploit the sequential nature of MCTS. Additionally, it uses a separate prompt to reflect on incorrect code predictions, whereas we integrate self-reflection within the generation prompt. Furthermore, LATS lacks specialized prompts and strategies for fixing buggy programs.

Previous research has also focused on pseudocode-based reasoning, such as Parsel [Zelikman et al., 2023], which uses a custom pseudocode language to decompose the program into independent problems that can be solved separately. In contrast, we focus on the sequential refinement of solutions using a variant of MCTS and the environment's feedback to produce directly executable Python code that can be leveraged in model-based RL.

We refer the reader to Appendix G for further discussion on works that build language-conditioned world models but do not use code and on works that use programs as policies in RL.

3 Code World Models

In this Section, we first introduce the Code World Models framework and then the proposed Code World Models Benchmark.

Code World Models framework. Following the model-based Reinforcement Learning problem setting, we consider an environment represented by a Markov Decision Process with state space S , action space A , a transition function p ( s ′ | a, s ) , and a scalar reward function R ( s, a, s ′ ) , with s, s ′ ∈ S indicating respectively the current and next state, and a ∈ A being the action taken from the current state. The task of a world model is to accurately represent p and R . We make the following assumptions: 1) the environments are deterministic and fully observable, and 2) we are provided with a natural language description of the environment, which is detailed enough to infer the observation space as well as the logic of the transition and reward functions.

The first assumption implies a deterministic transition function s ′ = f ( s, a ) , rather than a probabilistic one as in the general case; we address this limitation in Section 6.1. The second assumption is akin to the situation where a human would be provided with an explanation, or a tutorial, about a task that they need to solve, in order to facilitate the learning process. Crucially, in a model-based scenario, we only need explanations about how the environment works, rather than requiring instructions about what to do in order to solve the task. Furthermore, we place ourselves in an offline RL scenario [Levine et al., 2020], assuming that a dataset D of n one time-step transitions { ( s, a, r, s ′ , d ) i } i =1 ,...,n , where d stands for the episode termination or done signal, is available, collected with some behavioural policy π B ( a | s ) in the environment of interest. However, this last assumption could be lifted, by using the Code World Model with a suitable planning algorithm to collect more trajectories from the environment, turning the algorithm into online RL, as done in Tang et al. [2024b].

Code World Models Benchmark. To comprehensively test world model generation for a variety of environments, we define a novel benchmark consisting of 18 RL environments of varying difficulty. We focus on commonly used environments of particular relevance to the RL community: classical control, physics-based PyGame environments and MuJoCo tasks. The environments' Python implementations as well as their documentation are adapted from the Gymnasium library [Towers et al., 2024]. The environments included in the resulting Code World Models Benchmark (CWMB) feature a mix of continuous and discrete action and observation spaces (more details in Appendix I).

For each environment, we collect a training dataset D of past trajectories. We curate D so that it includes at least some low-scoring and some relatively high-scoring behavior. However, we neither attempt to maximally cover the state space nor do we require optimal demonstrations. We aim to show that relatively low annotation effort is required to build CWMs: for the majority of environments, we collect just 5 trajectories equivalent to taking random actions and a further 5 suboptimal demonstrations exceeding some return threshold. As part of the benchmark, each transition ( s, a, r, s ′ , d ) in each resulting trajectory is used as an input-output sample to validate the generated models. The benchmark further includes a language description of each environment, derived from the documentation written for Gymnasium's end users (an example is included in Appendix N.3). A further discussion on how the quality of the collected dataset affects the performance of our method can be found in Appendix F.

4 GIF-MCTS

In this Section, we first specify the format of the Code World Models that we consider in this work and how we evaluate their accuracy. We then present Generate, Improve and Fix with Monte Carlo Tree Search (GIF-MCTS), a novel approach to leverage LLMs for code generation via multiple

25%

0%

0%

62%

62%

58%

100%

LLM

import numpy as np

class Environment:

def __init__(self):

se

lf.cart_position =

np.random.uniform(-0.1, 0.1)

self.cart_velocity =

np.random.uniform(-0.1, 0.1)

self.mass_cart = 1.0

self.mass_pole = 0.1

def step(self, action):

...

return next_state, reward, done

Previous

state

Added

lines

Used as

rollout

import numpy as np

class Environment:

def __init__(self):

se

lf.cart_position =

np.random.uniform(-0.1, 0.1)

self.cart_velocity =

np.random.uniform(-0.1, 0.1)

self.pole_angle =

np.random.uniform(-0.1, 0.1)

self.pole_angular_velocity =

np.random.uniform(-0.1, 0.1)

...

import numpy as np

class Environment:

def __init__(self):

se

lf.cart_position =

np.random.uniform(

-0.2, 0.2

)

...

def step(self, action):

...

return next_state, reward, done

f1

class Environment:

def __init__(self):

se

lf.cart_position =

np.random.uniform(-0.1, 0.1)

...

def step(self, action):

...

return next_state, reward, done

LLM

LLM

import numpy as np

class Environment:

def __init__(self):

se

lf.cart_position =

np.random.uniform(-0.1, 0.1)

...

def step(self, action):

...

return next_state, reward, done

Prediction failed:

State and

Action: ...

Predicted

output:...

Ground truth:

...

g1

g2

f1

g1

i1

i2

Generate new lines

Fix bugs

Improve predictions

import numpy as np

class Environment:

def __init__(self):

se

lf.cart_position =

np.random.uniform(-0.1, 0.1)

self.cart_velocity =

np.random.uniform(-0.1, 0.1)

NameError: name

'np' is not defined

LEGEND

99%

66%

62%

Actual value

backpropagated

after bug is fixed

LLM call and

code evaluation

buggy

node

remaining

'fix' attempts

Temporary value

assigned to

buggy nodes

sequential attempts in the presence of feedback, specifically tailored to the needs of building Code World Models.

We formulate the task of synthesizing a Code World Model as that of writing a Python Environment class with a step() function that jointly implements the transition and reward functions:

and consider a Code World Model correctly synthesized if it correctly reproduces all transitions in D . We additionally define the accuracy A of the Code World Model as the fraction of correctly predicted transitions (weighted uniformly on next state, reward and done signals) from the training dataset D , or in other words:

where 1 is the indicator function (equals to one if the pair is matching, zero otherwise) and ˆ s ′ i , ˆ r i and ˆ d i are the model's predictions.

GIF-MCTS takes as input the description of an environment, an LLM, environment trajectories and builds a tree to construct the code for the environment. Nodes in the tree are programs and edges are actions. Each action taken from a parent node produces a new complete program, which is split into a state part and a rollout part and stored in a child node. The child node's state is formed from the parent's state by appending L additional lines of code (we set L = 2 in our work), while the rollout is the remaining part of the program, and represents one possible completion of the state, needed to evaluate (i.e., run) the code. This is a novel formulation of the state of a node, as we store

in the states partial programs in blocks of multiple lines, whereas previous work either stores only full programs [Zhou et al., 2023], or single tokens [Zhang et al., 2023]. The state represents the main flow of information from parent to child, while the rollout is used to estimate the expected accuracy of the child's state.

As in the standard MCTS algorithm, we perform multiple sequential iterations consisting of the following phases: selection, expansion, evaluation and value backpropagation. During the selection phase, starting from the root node, we use the Upper Confidence Bound for Trees (UCT) formula [Kocsis and Szepesvári, 2006] to select which action to take. If the corresponding node has never been expanded, we enter the expansion phase, otherwise we continue to apply the UCT formula to the actions of the new node. At expansion phase, we call the LLM to produce a program according to the type of action selected, parse the resulting program into the state and the rollout parts, and store both in the newly expanded node. We then compute the accuracy, defined above, using the rollout (evaluation phase), store the resulting value in the node, and backpropagate it to its ancestors. An example of a GIF-MCTS tree and the corresponding actions can be found in Figure 2.

With GIF-MCTS, we make the following contributions: 1) we present a novel framing of MCTS nodes and actions for long-form code generation in the presence of unit tests, 2) we propose three action types, specialised for code, whose added value we demonstrate through an ablation study, and 3) we propose a heuristic that empirically improves the trade-off between exploration and exploitation in the UCT formula used for action selection, balancing both explored and unexplored actions, and different action types (Appendix B). All these factors make GIF-MCTS specifically suitable for generating world models. Next we present the three action types (generate new lines, improve predictions and fix bugs) used in GIF-MCTS. We point the reader to the Appendix for the full action prompts, the remaining implementation details, and for the ablation study on the importance of the three action types.

4.1 GIF-MCTS Actions

Generate new lines. The goal of the generate action is to leverage the stochastic sampling ability of the LLM by generating varying continuations for a single code snippet in different branches of the tree, to fully explore the underlying space of possible solutions. The action prompt asks the LLM to generate the full code required to solve the task starting from the code stored in the node's state .

Improve predictions. Generating code in sequential blocks of lines can be too rigid if subtle or interdependent changes need to be made to the full program in order to pass more test cases and increase the reward. With the improve action, the LLM is prompted with the full program ( state plus rollout ) from the parent node, as well as one input example where the code did not behave as intended, along with the expected output. In the case of a Code World Model, this can be a wrongly predicted transition, with the input state and action taken by the agent, the ground-truth next state, and the model's predicted next state. The improve prompt also asks the LLM to produce a Chain-of-Thought explanation about where the current code is failing, and to attempt to fix the logic. The inclusion of both generate and improve actions allows GIF-MCTS to combine the advantages of block-wise incremental generation with the flexibility to backtrack and edit the whole program if needed.

Fix bugs. The code obtained with a generate or improve action will sometimes not be able to execute due to a syntax or runtime error, and will thus receive a reward of 0, strongly discouraging further exploration of the node. This can be wasteful, as sometimes the newly generated program can have sound logic and would receive a good reward if its bug(s) were removed. The fix action is tasked with resolving these bugs: the model is given the full program from the parent that encountered a bug along with feedback about the error and is asked to produce a fixed version of the code, aided by a Chain-of-Thought reasoning structure. To ensure that buggy nodes are chosen by the UCT formula, we assign them with temporary value until either the bug is fixed or no more attempts are allowed (see Appendix B for additional details).

5 Experiments

In this Section, we first describe the baseline code generation methods we compare against and then present empirical results on the APPS benchmark, the proposed CWMB and perform an additional

study on the RTFM environment. Additional ablations and qualitative results on GIF-MCTS are presented in Appendices C and D.

5.1 Baselines

The first baseline, denoted as Zero-shot CoT and used only for the experiments on APPS, adapts the work by Kojima et al. [2022] to code generation by appending "Let's think step by step." to the prompt and then parsing out from the completion only the code part. To report pass@20, we generate 20 independent completions for each problem, submit each of them, and count a problem as completed if at least one solution is correct.

The second baseline adapts the work by Tang et al. [2024b] to make as fair a comparison as possible. The WorldCoder algorithm calls the LLM with our generate prompt to produce an initial program, then for each remaining iteration we 1) select one of the previous programs as explained below, 2) refine it by calling the LLM with our fix prompt if the code has a bug, or our improve prompt otherwise, and 3) evaluate the resulting program against the unit tests. Each program ρ is associated with a Beta distribution B ( α, β ) with initial parameters α = 1 + C ∗ r ( ρ ) and β = 1 + C (1 -r ( ρ )) , which are updated every time the program is selected. Here r ( ρ ) stands for the fraction of unit tests passed (same metric used in the evaluation phase of GIF-MCTS) and C is a constant set to 5, as in the original work. To select the next program to be refined, one sample is drawn from each Beta distribution and the program with the highest score is selected. In all experiments, we use the same amount of calls of GIF-MCTS.

5.2 APPS

We assess the overall performance of GIF-MCTS for generic code synthesis in the presence of public unit tests on the APPS benchmark [Hendrycks et al., 2021], which consists of 10,000 Python coding problems in three categories of increasing difficulty: Introductory , Interview and Competition . We focus our evaluation on the hardest, Competition level test set comprised of 1000 problems, as it most closely reflects the challenges found in synthesizing CWMs: the problems tend to have a longer description, follow a specific format for the input and output, and include challenging logic. Early experiments on HumanEval [Chen et al., 2021], another popular coding benchmark, did not show a clear correlation between a model's performance on the benchmark and its ability to generate CWMs, as HumanEval problems are typically easier and solvable with much shorter code snippets.

As GIF-MCTS requires a reward signal from the environment, we make use of the suite of unit tests provided by APPS to evaluate the accuracy of a generated program. However, we note that the ground truth result from these tests is provided to GIF-MCTS with the improve action, and as such the model could simply memorize all possible results and return them without actually solving the problem. To avoid this, while we use all unit tests for computing the reward function, we only use samples from the first half as input-output examples for the improve action. In general, we use at least a fraction of the provided unit tests to evaluate every program generated during the GIF-MCTS loop, so our approach is only eligible for the pass@ B metric, where B is the budget for the number of LLM calls used during the synthesis process. We leave extending the approach for pass@1 eligibility using self-generated unit tests [Chen et al., 2023] for future work. We report the strict accuracy rate (the fraction of problems on which all test cases are solved) on APPS for GIF-MCTS and other baselines in Table 1.

Results. GIF-MCTS outperforms strong previous baselines on the APPS competition split, reaching a new state of the art to the best of our knowledge. While part of this can be due to advances in the underlying model, the comparisons with Zero-shot CoT and WorldCoder show improved performance over either prior method. GIF-MCTS is also markedly more sample efficient compared to established baselines; Parsel achieves the second best accuracy, but evaluates an exponentially growing number of solutions 4 , while GIF-MCTS outperforms it by evaluating only 20 different programs.

5.3 Code World Models Benchmark

We evaluate our proposed GIF-MCTS approach and the WorldCoder baseline on the CWMB (introduced in Section 3). In this setting, we are interested in both the accuracy of the generated CWM, as well as its performance when actually employed by a planning algorithm. We use as accuracy the same metric used in the evaluation phase of GIF-MCTS (Section 4). To measure the performance of planning with the CWM, we define the normalized return R of a CWM as:

where R ( π CWM ) represents the return obtained when using the CWM as the internal model for the planner, R ( π true ) is the return gathered with the true environment as the model while using the same planner (oracle planner), and R ( π rand ) is the return from a random policy. This metric is positive when the performance of the CWM planner is above that of a random policy and reaches one when the return approaches the value from the oracle planner. We report results for the CWMB in Table 2. As the planner, we use a vanilla MCTS implementation for the environments with discrete actions and a Cross Entropy Method (CEM) planner [Rubinstein, 1997] for the ones with continuous action spaces (full details of the two planning algorithms are reported in Appendix L).

Results. Overall, GIF-MCTS outperforms WorldCoder for all environment splits and backbone models. For Llama 3, the most significant gains are made on the environments with discrete actions, while for GPT-4 on those with continuous actions. We speculate that, on discrete environments, Llama 3 makes better use of the budget with GIF-MCTS than with WorldCoder, whereas GPT-4 saturates its performance in both cases. On the other hand, on the harder environments with continuous actions, Llama 3 hits a performance ceiling in both cases, while GPT-4 leads to higher improvements with our method. For example, Llama 3 was unable to generate a fully executable CWM (with either method) for the two hardest environments, Humanoid-v4 and HumanoidStandup-v4, due to their complexity and large observation space, while GPT-4 successfully generated a model for each environment in the benchmark.

5.4 Read to Fight Monsters

We perform an additional experiment on the Read to Fight Monsters (RTFM) grid-world environment, first introduced by Zhong et al. [2020] for testing grounded language understanding in RL. Every

episode presents two monsters belonging to two teams, and two items, each effective on a specific monster. The environment provides the agent with a written descriptions of the task dynamics (also called manual), describing monsters' weaknesses and membership to teams, and a goal (which team of monsters to defeat). Crucially, the agent needs to perform multi-step reasoning between such information and the current state of the environment to figure out a plan of action (for more details we refer to the original work by Zhong et al. [2020]). We consider a version of the environment where we fix the input manual, meaning all relationships between items and monsters are fixed across episodes, and we don't allow the monsters to move, as their patterns are stochastic. This isolates the natural language understanding component of the task, while we leave to future work to demonstrate the applicability of the CWM framework to the full RTFM task.

We report the results on the simplified RTFM environment in Table 3, using MCTS as a planner for computing the normalized returns. We further experiment with a higher number of LLM calls for GPT-4 Turbo, matching the one used for Llama 3, as we couldn't do this on the full CWMB due to budget concerns.

Results. GIF-MCTS outperforms WorldCoder under all settings by a significant margin in terms of accuracy, but the generated CWM is only able to match the performance of the ground-truth simulator when the program is perfect. This highlights the necessity of completely accurate predictions, as further discussed in Section 6, while also providing empirical validation for the scaling properties of the approach: as GIF-MCTS is allowed more calls, it manages to refine the CWM it generated with a lower budget. As this version of the RTFM environment has never been published, this experiment can also alleviate concerns that the final CWM was memorized by the LLM during pre-training. We present and discuss further evidence against the significance of data contamination in Appendix E.

6 Discussion

In this section, we first discuss some takeaways from the empirical results and then elaborate on some of the limitations for our method.

GIF-MCTS vs. WorldCoder. We believe that GIF-MCTS outperforms WorldCoder because it produces a more diverse set of programs. WorldCoder initially generates a single program from scratch and then samples and refines a complete program in each iteration. In contrast, GIF-MCTS can generate multiple programs either from scratch or from partial programs by taking the generate new lines action at the root node or subsequent nodes. This approach better explores the solution space, leading to improved performance. Our ablation study No Generate action in Table 6 of the Appendix supports this finding. This study uses a tree search like GIF-MCTS but always refines a complete program, similar to WorldCoder, and results in lower performance compared to our method.

Accuracy-Return Gap. We observe empirically from Table 2 that the CWM accuracy is always higher than its normalized return, and the two metrics match only when the CWM is flawless. This is often due to the incorrect prediction of terminal states: these are rarer in the replay buffer, especially states that terminate with a success/positive reward. This can cause the planning algorithm to fail, as it is missing the reward signal. Part of the performance gap could also be due to sparse coverage of the environment by the collected trajectories. Individual results for each environment elaborating

on this are included in Appendix J. Future work could explore retrieving and combining different CWMs that complement each other to improve the performance on important edge cases.

Sample Efficiency. Generating a CWM requires far less interaction with the environment than traditional model-based approaches. As the gathered transitions are only used to validate the program and as in-context examples, a small curated set (enough to cover possible edge cases and different reward values) is enough to properly validate the generated code. In our experiments we only gather 10 trajectories made up of at most 100 steps as the offline dataset, while benchmarks specifically designed to challenge for sample efficiency [Bellemare et al., 2013] require agents to use at most 100k frames, which is two orders of magnitude higher. We leave more thorough experiments on sample efficiency for CWM agents to future work.

Comparison with Offline RL. We expect CWMs to hold advantages over classical RL methods in regimes with scarce data and environments that can be easily described by language and modeled with code. We report in Appendix K a preliminary comparison on the CWMB of the return achieved with our CWMs or with a SOTA offline RL method, Conservative Q-Learning (CQL) [Kumar et al., 2020], trained on the same amount of trajectories used for synthesizing the CWMs. We find that CWMs compare favourably against CQL on environments with discrete action spaces, while CQL's performance is superior on the continuous action space environments, which are harder to model. RL methods, including CQL, would likely benefit from more experience, as they overfit with scarce data.

6.1 Limitations

Code World Models. The CWMs framework is an exciting direction for model-based planning, but we still rely on limiting assumptions of deterministic and fully observable environments. Both stochasticity and partial observability would pose challenges, especially on the verification of the CWM prediction, as there is no set result for a given input. We leave extending the approach to account for both stochastic and partially observable environments to future work.

Another potential issue is providing a description of the environment that can be reasonably converted to a Python function (e.g. a manual documenting key variables) when such a description is not available (e.g. when the environment is defined with image observations). Previous work has begun to tackle this issue [Migimatsu and Bohg, 2022] and preprocessing techniques such as image-to-text models [Ren et al., 2024] could be used to address this problem in future work.

Code-based models may also be too rigid when the environment requires adapting to changing dynamics, which would imply rewriting the CWM on the fly. A possible solution could be breaking down the CWM into smaller functions that can be re-written individually by an LLM, to account for some changes in the environment, or modeling variable factors as arguments to the step function. CWMs struggle especially on complex physics-based environments; thus a promising direction could also be allowing programs generated by GIF-MCTS to make use of external tools and libraries, such as physics simulators.

GIF-MCTS. We have validated the GIF-MCTS approach as an efficient code synthesis method, with the key limiting assumption of having available test cases to evaluate code, which could be difficult to provide in certain tasks. In those cases, it would be possible to use self-generated test cases [Chen et al., 2023], but since this does not reflect the CWM setting we leave this for future work.

7 Conclusion

We present Code World Models, a general framework to leverage LLMs to build world models for RL agents. We further show that GIF-MCTS is a strong code synthesis method, able to successfully integrate external feedback to self-debug and improve code, demonstrating examples of world modeling and downstream planning for a range of environments. We are confident that the Code World Models approach will lead to the development of fast, interpretable and sample efficient model-based RL agents, exploiting the strengths provided by increasingly powerful LLMs, without directly predicting the environment dynamics with them. We are hopeful that improvements to both the underlying LLM backbone and refinements to the code generation method itself will result in powerful Code World Models for even more complex environments than those treated in this work.

Acknowledgments and Disclosure of Funding

This work was supported by the Research Council of Finland (Flagship programme: Finnish Center for Artificial Intelligence FCAI, and grants 352986, 358246) and EU (H2020 grant 101016775 and NextGenerationEU). We acknowledge CSC for awarding this project access to the LUMI supercomputer, owned by the EuroHPC Joint Undertaking, hosted by CSC (Finland) and the LUMI consortium through Finland.

References

Michael Ahn, Anthony Brohan, Noah Brown, Yevgen Chebotar, Omar Cortes, Byron David, Chelsea Finn, Keerthana Gopalakrishnan, Karol Hausman, Alexander Herzog, et al. Do as i can, not as i say: Grounding language in robotic affordances. In Conference on Robot Learning , 2022.

Abdus Salam Azad, Edward Kim, Qiancheng Wu, Kimin Lee, Ion Stoica, Pieter Abbeel, Alberto Sangiovanni-Vincentelli, and Sanjit A Seshia. Programmatic modeling and generation of real-time strategic soccer environments for reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 36, pages 6028-6036, 2022.

Osbert Bastani, Yewen Pu, and Armando Solar-Lezama. Verifiable reinforcement learning via policy extraction. In Advances in Neural Information Processing Systems , volume 31, 2018.

Marc G. Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environment: An evaluation platform for general agents. Journal of Artificial Intelligence Research , 47: 253-279, 2013.

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, et al. Language models are few-shot learners. In Advances in Neural Information Processing Systems , volume 33, pages 1877-1901, 2020.

Jake Bruce, Michael D Dennis, Ashley Edwards, Jack Parker-Holder, Yuge Shi, Edward Hughes, Matthew Lai, Aditi Mavalankar, Richie Steigerwald, Chris Apps, Yusuf Aytar, et al. Genie: Generative interactive environments. In Proceedings of the 41st International Conference on Machine Learning , volume 235 of Proceedings of Machine Learning Research , pages 4603-4623. PMLR, 21-27 Jul 2024.

Tales Henrique Carvalho, Kenneth Tjhia, and Levi Lelis. Reclaiming the source of programmatic policies: Programmatic versus latent spaces. In The Twelfth International Conference on Learning Representations , 2024.

Bei Chen, Fengji Zhang, Anh Nguyen, Daoguang Zan, Zeqi Lin, Jian-Guang Lou, and Weizhu Chen. Codet: Code generation with generated tests. In The Eleventh International Conference on Learning Representations , 2023.

Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde De Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. Evaluating large language models trained on code. arXiv preprint arXiv:2107.03374 , 2021.

Nicola Dainese, Pekka Marttinen, and Alexander Ilin. Reader: Model-based language-instructed reinforcement learning. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing , pages 16583-16599, Singapore, December 2023. Association for Computational Linguistics.

Ria Das, Joshua B. Tenenbaum, Armando Solar-Lezama, and Zenna Tavares. Autumnsynth: Synthesis of reactive programs with structured latent state. In Advances in Programming Languages and Neurosymbolic Systems Workshop , 2021.

Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219 , 2020.

Malik Ghallab, Adele Howe, Craig Knoblock, Drew McDermott, Ashwin Ram, Manuela Veloso, Daniel Weld, and David Wilkins. PDDL, The Planning Domain Definition Language, 1998.

Zhibin Gou, Zhihong Shao, Yeyun Gong, yelong shen, Yujiu Yang, Nan Duan, and Weizhu Chen. CRITIC: Large language models can self-correct with tool-interactive critiquing. In The Twelfth International Conference on Learning Representations , 2024.

Lin Guan, Karthik Valmeekam, Sarath Sreedharan, and Subbarao Kambhampati. Leveraging pretrained large language models to construct and utilize world models for model-based task planning. In Advances in Neural Information Processing Systems , volume 36, pages 79081-79094, 2023.

Agrim Gupta, Lijun Yu, Kihyuk Sohn, Xiuye Gu, Meera Hahn, Li Fei-Fei, Irfan Essa, Lu Jiang, and José Lezama. Photorealistic video generation with diffusion models. arXiv preprint arXiv:2312.06662 , 2023.

David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In Advances in Neural Information Processing Systems , volume 31, 2018.

Danijar Hafner, Timothy P Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In International Conference on Learning Representations , 2021.

Shibo Hao, Yi Gu, Haodi Ma, Joshua Jiahua Hong, Zhen Wang, Daisy Zhe Wang, and Zhiting Hu. Reasoning with language model is planning with world model. In The 2023 Conference on Empirical Methods in Natural Language Processing , 2023.

Dan Hendrycks, Steven Basart, Saurav Kadavath, Mantas Mazeika, Akul Arora, Ethan Guo, Collin Burns, Samir Puranik, Horace He, Dawn Song, and Jacob Steinhardt. Measuring coding challenge competence with apps. In Proceedings of the Neural Information Processing Systems Track on Datasets and Benchmarks , volume 1, 2021.

Jeevana Priya Inala, Osbert Bastani, Zenna Tavares, and Armando Solar-Lezama. Synthesizing programmatic policies that inductively generalize. In 8th International Conference on Learning Representations , 2020.

Eric Jang, Alex Irpan, Mohi Khansari, Daniel Kappler, Frederik Ebert, Corey Lynch, Sergey Levine, and Chelsea Finn. Bc-z: Zero-shot task generalization with robotic imitation learning. In Conference on Robot Learning , pages 991-1002. PMLR, 2022.

Siddharth Karamcheti, Suraj Nair, Annie S. Chen, Thomas Kollar, Chelsea Finn, Dorsa Sadigh, and Percy Liang. Language-driven representation learning for robotics. In Robotics: Science and Systems (RSS) , 2023.

Michael N. Katehakis and Arthur F. Veinott. The multi-armed bandit problem: Decomposition and computation. Mathematics of Operations Research , 12(2):262-268, 1987.

Levente Kocsis and Csaba Szepesvári. Bandit based monte-carlo planning. In Proceedings of the 17th European Conference on Machine Learning , ECML'06, page 282-293, Berlin, Heidelberg, 2006. Springer-Verlag.

Takeshi Kojima, Shixiang (Shane) Gu, Machel Reid, Yutaka Matsuo, and Yusuke Iwasawa. Large language models are zero-shot reasoners. In Advances in Neural Information Processing Systems , volume 35, pages 22199-22213, 2022.

Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. In Advances in Neural Information Processing Systems , volume 33, pages 1179-1191, 2020.

Hung Le, Yue Wang, Akhilesh Deepak Gotmare, Silvio Savarese, and Steven Chu Hong Hoi. Coderl: Mastering code generation through pretrained models and deep reinforcement learning. In Advances in Neural Information Processing Systems , volume 35, pages 21314-21328, 2022.

Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643 , 2020.

Jacky Liang, Wenlong Huang, Fei Xia, Peng Xu, Karol Hausman, Brian Ichter, Pete Florence, and Andy Zeng. Code as policies: Language model programs for embodied control. In 2023 IEEE International Conference on Robotics and Automation (ICRA) , pages 9493-9500. IEEE, 2023.

Jessy Lin, Yuqing Du, Olivia Watkins, Danijar Hafner, Pieter Abbeel, Dan Klein, and Anca Dragan. Learning to model the world with language. In Proceedings of the 41st International Conference on Machine Learning , volume 235 of Proceedings of Machine Learning Research , pages 2999230017. PMLR, 21-27 Jul 2024.

Guan-Ting Liu, En-Pei Hu, Pu-Jen Cheng, Hung-Yi Lee, and Shao-Hua Sun. Hierarchical programmatic reinforcement learning via learning to compose programs. In Proceedings of the 40th International Conference on Machine Learning , volume 202 of Proceedings of Machine Learning Research , pages 21672-21697. PMLR, 23-29 Jul 2023.

Hao Liu, Wilson Yan, Matei Zaharia, and Pieter Abbeel. World model on million-length video and language with ringattention. arXiv preprint arXiv:2402.08268 , 2024.

Aman Madaan, Niket Tandon, Prakhar Gupta, Skyler Hallinan, Luyu Gao, Sarah Wiegreffe, Uri Alon, Nouha Dziri, Shrimai Prabhumoye, Yiming Yang, Shashank Gupta, Bodhisattwa Prasad Majumder, Katherine Hermann, Sean Welleck, Amir Yazdanbakhsh, and Peter Clark. Self-refine: Iterative refinement with self-feedback. In Advances in Neural Information Processing Systems , volume 36, pages 46534-46594, 2023.

Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample-efficient world models. In The Eleventh International Conference on Learning Representations , 2023.

Toki Migimatsu and Jeannette Bohg. Grounding predicates through actions. In 2022 IEEE International Conference on Robotics and Automation (ICRA) , pages 3498-3504. IEEE, 2022.

Theo X Olausson, Jeevana Priya Inala, Chenglong Wang, Jianfeng Gao, and Armando Solar-Lezama. Is self-repair a silver bullet for code generation? In The Twelfth International Conference on Learning Representations , 2023.

Wenjie Qiu and He Zhu. Programmatic reinforcement learning without oracles. In The Tenth International Conference on Learning Representations , 2022.

Tianhe Ren, Shilong Liu, Ailing Zeng, Jing Lin, Kunchang Li, He Cao, Jiayu Chen, Xinyu Huang, Yukang Chen, Feng Yan, et al. Grounded sam: Assembling open-world models for diverse visual tasks. arXiv preprint arXiv:2401.14159 , 2024.

Reuven Y Rubinstein. Optimization of computer simulation models with rare events. European Journal of Operational Research , 99(1):89-112, 1997.

Noah Shinn, Federico Cassano, Ashwin Gopinath, Karthik Narasimhan, and Shunyu Yao. Reflexion: language agents with verbal reinforcement learning. In Advances in Neural Information Processing Systems , volume 36, pages 8634-8652, 2023.

Tom Silver, Kelsey R Allen, Alex K Lew, Leslie Pack Kaelbling, and Josh Tenenbaum. Few-shot bayesian imitation learning with logical program policies. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 34, pages 10251-10258, 2020.

Ishika Singh, Valts Blukis, Arsalan Mousavian, Ankit Goyal, Danfei Xu, Jonathan Tremblay, Dieter Fox, Jesse Thomason, and Animesh Garg. Progprompt: Generating situated robot task plans using large language models. In 2023 IEEE International Conference on Robotics and Automation (ICRA) , pages 11523-11530. IEEE, 2023.

Richard S. Sutton. Dyna, an integrated architecture for learning, planning, and reacting. SIGART Bulletin , 2(4):160-163, July 1991.

Hao Tang, Keya Hu, Jin Peng Zhou, Sicheng Zhong, Wei-Long Zheng, Xujie Si, and Kevin Ellis. Code repair with llms gives an exploration-exploitation tradeoff. arXiv preprint arXiv:2405.17503 , 2024a.

Hao Tang, Darren Key, and Kevin Ellis. Worldcoder, a model-based llm agent: Building world models by writing code and interacting with the environment. arXiv preprint arXiv:2402.12275v1 , 2024b.

William R Thompson. On the likelihood that one unknown probability exceeds another in view of the evidence of two samples. Biometrika , 25(3-4):285-294, 1933.

Mark Towers, Ariel Kwiatkowski, Jordan Terry, John U Balis, Gianluca De Cola, Tristan Deleu, Manuel Goulao, Andreas Kallinteris, Markus Krimmel, Arjun KG, et al. Gymnasium: A standard interface for reinforcement learning environments. arXiv preprint arXiv:2407.17032 , 2024.

Dweep Trivedi, Jesse Zhang, Shao-Hua Sun, and Joseph J Lim. Learning to synthesize programs as interpretable and generalizable policies. In Advances in Neural Information Processing Systems , volume 34, pages 25146-25163, 2021.

Pedro A Tsividis, Joao Loula, Jake Burga, Nathan Foss, Andres Campero, Thomas Pouncy, Samuel J Gershman, and Joshua B Tenenbaum. Human-level reinforcement learning through theory-based modeling, exploration, and planning. arXiv preprint arXiv:2107.12544 , 2021.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Ł ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems , volume 30, 2017.

Abhinav Verma, Vijayaraghavan Murali, Rishabh Singh, Pushmeet Kohli, and Swarat Chaudhuri. Programmatically interpretable reinforcement learning. In International Conference on Machine Learning , pages 5045-5054. PMLR, 2018.

Abhinav Verma, Hoang Le, Yisong Yue, and Swarat Chaudhuri. Imitation-projected programmatic reinforcement learning. In Advances in Neural Information Processing Systems , volume 32, 2019.

Guanzhi Wang, Yuqi Xie, Yunfan Jiang, Ajay Mandlekar, Chaowei Xiao, Yuke Zhu, Linxi Fan, and Anima Anandkumar. Voyager: An open-ended embodied agent with large language models. arXiv preprint arXiv:2305.16291 , 2023.

Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, brian ichter, Fei Xia, Ed H. Chi, Quoc V Le, and Denny Zhou. Chain of thought prompting elicits reasoning in large language models. In Advances in Neural Information Processing Systems , 2022.

Lionel Wong, Gabriel Grand, Alexander K Lew, Noah D Goodman, Vikash K Mansinghka, Jacob Andreas, and Joshua B Tenenbaum. From word models to world models: Translating from natural language to the probabilistic language of thought. arXiv preprint arXiv:2306.12672 , 2023.

Lionel Wong, Jiayuan Mao, Pratyusha Sharma, Zachary S Siegel, Jiahai Feng, Noa Korneev, Joshua B. Tenenbaum, and Jacob Andreas. Learning grounded action abstractions from language. In The Twelfth International Conference on Learning Representations , 2024.

Sherry Yang, Yilun Du, Seyed Kamyar Seyed Ghasemipour, Jonathan Tompson, Leslie Pack Kaelbling, Dale Schuurmans, and Pieter Abbeel. Learning interactive real-world simulators. In The Twelfth International Conference on Learning Representations , 2024.

Shunyu Yao, Dian Yu, Jeffrey Zhao, Izhak Shafran, Tom Griffiths, Yuan Cao, and Karthik Narasimhan. Tree of thoughts: Deliberate problem solving with large language models. In Advances in Neural Information Processing Systems , volume 36, pages 11809-11822, 2023.

Eric Zelikman, Qian Huang, Gabriel Poesia, Noah Goodman, and Nick Haber. Parsel: Algorithmic reasoning with language models by composing decompositions. In Advances in Neural Information Processing Systems , volume 36, pages 31466-31523, 2023.

Alex Zhang, Khanh Nguyen, Jens Tuyls, Albert Lin, and Karthik Narasimhan. Language-guided world models: A model-based approach to ai control. arXiv preprint arXiv:2402.01695 , 2024.

Shun Zhang, Zhenfang Chen, Yikang Shen, Mingyu Ding, Joshua B. Tenenbaum, and Chuang Gan. Planning with large language models for code generation. In The Eleventh International Conference on Learning Representations , 2023.

Victor Zhong, Tim Rocktäschel, and Edward Grefenstette. RTFM: Generalising to new environment dynamics via reading. In International Conference on Learning Representations , 2020.

Andy Zhou, Kai Yan, Michal Shlapentokh-Rothman, Haohan Wang, and Yu-Xiong Wang. Language agent tree search unifies reasoning acting and planning in language models. arXiv preprint arXiv:2310.04406 , 2023.

A Broader Impact

The CWM framework enables LLMs to generate world models for model-based Reinforcement Learning, which could potentially be employed for planning with a real agent. As the code generated by the LLM is untrusted, it should always be checked by a human expert before it is used under any circumstances. Alternatively, as CWMs are represented with Python code, this also allows for interpretable world models, which could be safer for critical applications after being vetted by an expert.

B Additional GIF-MCTS implementation details

Choice of Actions If a node doesn't contain a bug, new generate and improve actions should always be available (with the exception of the root node, which will only have a new generate action, since there is no pre-existing code to improve). After an action is expanded, we add a new action of the same type to the parent node, so that the tree can have a variable number of nodes at any level. By contrast, a buggy node will only ever have a single fix action available, and no new fix actions will be added to the parent, enforcing the fixes to be applied sequentially (as there is no need to expand the tree horizontally in a buggy node). To select actions, we follow a modified variant of the Upper Confidence Bound for Trees (UCT) formula [Kocsis and Szepesvári, 2006] as follows:

where v i is the value of the node, C is a constant parameter used to balance exploration (empirically set to 0.1), N i is the number of visits to the node's parent and n a = a i is the number of expanded children with the same action type (relative to the parent). This last parameter is required to avoid trees that only grow horizontally due to the added actions: if a single action is chosen too many times from the same parent, the n a = a i term will cause the exploration value for new nodes for the same action to keep decreasing and therefore encourage more exploration.

Value Estimation for Unexplored Nodes. Nodes that have not yet been visited are missing their value, which prevents the application of the UCT formula. To circumvent this, we employ a simple linear model, trained during the overall search, to predict the value of unexplored nodes. This estimate is specific to an action type, so that each has a separate classifier, and further differentiates local and global values. We define the global value v G as the average of all values of the nodes with the same action type at any level of the tree and the local value v L as the average of all expanded children with the same action type. The linear model then simply learns to predict the value v i of a given action as a balanced sum of the two values, normalized between zero and one, with the following formula:

where the w G and w L parameters are learned during the search using gradient descent.

Initially, the global average v G will also be empty, which would cause the first values to be ill-defined. To mitigate this, we initialize the global average with a prior value which we tune empirically. To ensure a single unlucky generation does not prematurely downweight an action type, this prior is further assigned an initial count, used to weight the prior when computing the average (effectively acting as if there were n nodes already discovered with the prior value).

Value Estimation for Buggy Nodes. As mentioned in Sec. 4, buggy nodes will get a reward of 0 and would thus never be explored. To allow the fix action to be chosen, we assign a temporary value to the buggy node (which is effectively the parent of the fix action nodes). This can be chosen arbitrarily to trade-off attempting to fix buggy nodes (exploration) and focusing on other already functioning branches (exploitation). In our implementation, we initially set this value to 0.99, effectively forcing the model to attempt fixing a buggy node at least once. Naturally, a program can have more than one bug which could require the method taking multiple fix actions. To account for this, if the outcome of a fix action is still a bug, we gradually linearly decrease the temporary value of the parent until it reaches zero after a certain number of allowed fixes f , which we set to three. After f unsuccessful fixes, the temporary value is set to zero, which strongly discourages the buggy parent node from

being selected again. Otherwise, the value of the buggy parent and the fix children are set to the value received by the newly fixed program. It is also important to note that the temporary values are excluded from the backtracking step of the MCTS algorithm, to avoid skewing the ancestors' values.

Hyperparameters We report all hyperparameters used for GIF-MCTS as well as their description in Table 4, while hyperparameters related to the backbone LLM are reported in Table 5. We refer to the Huggingface documentation 5 for an accurate description of each LLM parameter.

C Ablation Study on GIF-MCTS

We perform an ablation study to validate the individual contribution of each action type of GIF-MCTS. We run the MCTS procedure on CWMB with only two out of the three actions available and compare the accuracy with the full method in Table 6. Note that for the Fix and Improve MCTS variant, one generate action is applied at the root node to obtain an initial program, which the algorithm expands from with the available budget. All ablations are performed using Llama 3 70B. For budget constraints, we run a single random seed for each ablation and compare with a single GIF-MCTS run with the same random seed.

Results. The performance of the method drops after the removal of each action, most significantly in the harder set of continuous environments (while there is more statistical uncertainty for the discrete environments). Fixing bugs appears to be the most important action: it is much more efficient to try fixing a bug aided by external feedback compared to blindly generating the same code snippet until bug-free. As the complexity of the environment grows, it might also become increasingly challenging to generate a fully functioning program from the start. On the other hand, improve seems to be the least impactful: this makes sense, as intuitively improving a code snippet that already works is has less room for improvement.

D Qualitative Study

To investigate the specific effectiveness of each individual type of action, we analyze the trees produced by GIF-MCTS and report some statistics of interest in Table 7. We specifically focus on the

difference in the overall distribution of action types in the tree as a whole compared to the actions chosen on the path that led to the best result, which can be used to find specific biases towards a specific action.

From the results, the method presents a pretty clear bias towards the generate action at the expense of the improve action on the optimal path. While the model tries to improve its previous code reasonably often (more than 35% of the times in most cases) the percentage of these actions that actually led to the best node drops significantly in the optimal path, which could imply that generate actions are the most effective.

With a closer inspection into the trees themselves, we find that often there is an initial set of generate actions that already result in values that are close to the maximum found by the tree, and then later improve actions are chosen thanks to the same-action penalty term in the modified UCT formula, which can result in marginal increases (as they are only refining code that is already promising) or fail to improve the previous program (as the logic might be hard to extrapolate). As such, many improve actions are needed in order to find a sample that is actually increasing the performance, while generate actions have the advantage of being chosen at the beginning, where it is possibly easier to find good programs.

Still, the fact that many improve actions are taken that result in either the same value as the previous node or at times even in worse accuracy is a potential bottleneck for the method, which seems to corroborate recent evidence [Olausson et al., 2023] showing that LLMs are often unable to provide proper feedback on their own code generations. Stronger models might thus be needed to specifically analyze and criticize the code (e.g. one model specialized in explaining code which provides feedback to another one specialized in generating it).

There is also a clear difference between the set of easier discrete action space problems, for which the percentage of fix actions is very low (with GPT-4 Turbo only needing generates in order to synthesize perfect or near-perfect models, as shown in Table 11) and the harder continuous action space problems, where fixing bugs becomes much more prominent.

E Data Contamination

With any experiment involving LLMs there is a concern about data contamination: the model's pre-training corpus could have included the original implementation for the various programs we are trying to generate, which means that hypothetically the model could simply be memorizing them and repeating them. To alleviate these concerns, we analyze each experiment individually:

For the APPS benchmark, the programming problems we used are sourced from three main websites. The benchmark authors managed to crawl reference solutions for only two of these sites ( AtCoder and Codeforces , which include 264 and 41 problems respectively). This means that for the third website, Kattis , which makes up a majority of the benchmark with 691 problems, no reference solution can be found online (and thus likely also not in the training corpus for the LLMs).

Performance across all methods and models in the competition split is correlated with the source websites of the problems, but not with the availability of the solutions: the highest results are obtained from Kattis (0.347 strict accuracy rate), the only site where solutions are not available online. Notably, all methods and models achieve a 0% pass rate for the 41 problems from AtCoder, for which reference solutions are available online. This suggests that the difficulty of the various sources is more important than the reference solution.

While we observe that some parts of the generated CWMB environments recall implementations available online (e.g., constants' values in the CartPole environment), the logic of the step function remains distinct from the reference model. Furthermore, the MuJoCo-based environments used the simulator in the official implementation, which is not available in our setting, so the code is necessarily different. Examples of generated CWMs along with their ground-truth implementations can be found in Appendix O for a more thorough comparison.

As we use a modified version of the RTFM environment (with fixed manuals and no stochasticity), there is no reference solution for it online, which provides evidence that our solution is not merely retrieving information from the LLM's training data.

Generally speaking, there is of course no way to outright dismiss these concerns. However, our method is compared to baselines using the same underlying models, ensuring that the superior performance reported for GIF-MCTS is not biased by potential data contamination.

F Data Quality

As part of the CWMB, for each environment the collected dataset D contains both low-scoring and high-scoring trajectories. As discussed in Section 3, this is fairly standard practice for offline RL, as the general assumption is that in the real world large datasets can be collected from a very diverse ensemble of sources. While it would be expected that at least one example for all possible outcomes is required in order for the world model to be precise and comprehensive, our approach can in principle learn a fair model even in hard environments when provided with only a few random trajectories by leveraging the language description provided to the LLM when generating the program. This could theoretically be used to generalize the rules of the environment outside of the observed transitions: the model does not need to see what happens if it can read about it.

We performed an additional experiment on RTFM: we collected 10 trajectories all resulting in failures, so that a reward of +1 is never observed. In other words, this is a worse version of the same buffer used for the main experiment, which by construction carries less information. We synthesized a CWM with GIF-MCTS and GPT-4 using 50 calls, which in the original experiment resulted in a perfect model (Section 5.4). The resulting CWM is 100% accurate on the newly collected dataset and even correctly predicts a reward of +1 for positive transitions, which are not included in the dataset, thanks to the language description. When tested on the original dataset D from the CWMB (which

contains both positive and negative rewards), the model still scores 100% accuracy, on par with the model generated with the full range of data.

G Additional Related Work

We expand in the following the Related Work section, covering the works that try to build world models with language and those who explored using programs to express RL policies.

World Models with Language. Model-based RL methods are built around learning a predictive model of the environment to inform the agent's decisions [Sutton, 1991]. A recently growing body of research is focusing on building world models that can include information in natural language, as opposed to approaches using only vision or full state observations [Hafner et al., 2021]. Dynalang [Lin et al., 2024] predicts the future text and image representation of the environment with an encoderdecoder architecture with a joint input of previous frames and text, while Zhang et al. [2024] formulate the modeling task as an autoregressive prediction task performed by a Transformer [Vaswani et al., 2017]. Voltron [Karamcheti et al., 2023] also uses an encoder-decoder model for language-driven representation learning for robotics. Other promising avenues include predicting the pixels in the next image observation [Yang et al., 2024, Bruce et al., 2024, Micheli et al., 2023, Liu et al., 2024].

Programmatic RL. Verma et al. [2018, 2019] first introduced Programmatically Interpretable RL (PIRL), which focuses on representing RL policies as interpretable and verifiable programs by first learning an oracle policy with deep RL and then distilling a program with a domain specific language that can model tree-like programs. Similarly, Bastani et al. [2018] focus on extracting decision trees from an oracle policy with imitation learning and Inala et al. [2020] use finite-state automata, which can also include advanced control structures such as loops, with Silver et al. [2020] similarly using a language with a token that can perform loops. The need for an oracle was later removed by Qiu and Zhu [2022] by directly optimizing differentiable programs. Later, Trivedi et al. [2021] introduce LEAPS, which uses a Variational Auto-Encoder (VAE) to embed programs into a latent space and search new programs in the latent space, further extended by Liu et al. [2023] with the use of Hierarchical RL that composes simple programs together in order to generalize to out of distribution codes not seen by the V AE. However, Carvalho et al. [2024] has recently shown that the latent space is actually harder for optimization algorithms, and that simply performing the search in the program space leads to better results. Azad et al. [2022] instead proposed using a similar domain specific language to build a world model, with a similar approach presented by EMPA [Tsividis et al., 2021]. As these methods all use traditional program synthesis methods to generate their code, recent works have also looked into using LLMs to generate RL policies. Liang et al. [2023] uses Python code to interface with APIs and generate a robotic policy, with a similar approach concurrently introduced by Singh et al. [2023]. Voyager [Wang et al., 2023] generates an incrementally growing skill library using JavaScript code to play Minecraft.

H Comparison of Inference Times

We further demonstrate the efficiency of CWMs compared to directly using an LLM as the world model in Table 8. On a selection of three environments from the CWMB we ask GPT-4 Turbo to directly predict the next observation of the environment given its description and some in-context examples of the task, and compare the inference time with calling the step function of the CWM. Calling the Python program is four orders of magnitude quicker for the easiest environment and seven orders of magnitude quicker for the hardest environment. We additionally observe that none of the predictions made by GPT-4 Turbo were accurate.

I Code World Models Benchmark Details

We include a detailed list of statistics for each environment in the CWMB in Table 9. Notice that when creating the descriptions from the Gymnasium docstrings, we left out documentation sections that do not relate to the environment definition itself, such as versioning information, Gymnasium-related arguments, and external references, from these descriptions. For the reported number of tokens we

choose OpenAI's open source tiktoken tokenizer 6 . The code lines and code tokens are reported from the corresponding CWM generated by GPT-4 Turbo using GIF-MCTS with a budget of 10. This is meant to be a general indication of how long a typical implementation of the environment would be, but can of course vary. All environment descriptions were parsed from Gymnasium v.0.29.1.

J Results for Individual Environments

We report the individual accuracy and return for each environment in the CWM when using Llama 3 in Table 10 and when using GPT-4 Turbo in Table 11.

K Comparison with Offline RL

We compare the overall performance of a SOTA offline RL method, Conservative Q-Learning (CQL) [Kumar et al., 2020], against a planning agent using the synthesized CWM with our method. We report in Table 12 the average raw reward obtained over 10 episodes for a random policy, CQL, planning agents with the CWM obtained by GIF-MCTS (ours) respectively with Llama 3 and GPT-4, and a planning agent with oracle access to the true environment. CQL was trained with 10 epochs for 100 steps per epoch (1000 total) using the same dataset D used to learn our CWMs. We chose 1000 steps to match the data to gradient steps ratio from the original CQL paper. Since our replay buffers are much smaller (the original paper worked with D4RL [Fu et al., 2020], which provides 1M transitions per task), we started to observe severe overfitting for CQL with more training steps.

Overall, there is a balance between CQL and CWMs, with CWMs being more suited to discrete tasks and CQL outperforming CWMs in complex physics tasks, where our method struggles. However,

CWMs also reach competitive results in some of these harder environments, such as Pendulum-v1 , Reacher-v4 and to a lesser extent Ant-v4 , Pusher-v4 and HalfCheetah-v4 , even without direct access to the original physics simulator. Particularly in these tasks, but also in general, we observe severe overfitting happening in CQL almost immediately (for example, CQL performs worse than random in Pendulum-v1 ), likely due to the small size of the provided dataset. As mentioned previously, sample efficiency is one of the main promises of the CWM approach, as very few trajectories are needed to validate the model, whereas traditional methods are typically designed to work best with large amounts of data.

* N/A for CQL indicates a failed run, while for GIF-MCTS it indicates a failure in synthesizing a syntactically correct CWM.

It is also worth noting that outperforming state-of-the-art methods for offline RL was not the principal goal we set out to achieve with our work, and as such many aspects are not specifically tuned for performance. For instance, we chose very simple planners with default parameters in order to collect the rewards with the synthesized CWMs, to study the performance of the models in the simplest possible setting. In general, our main objective is to validate the effectiveness of the framework, and we leave improvements that can show increased performance over offline RL methods (for instance, allowing the generated code to call a physics simulator in the continuous environments) to future work, now that the effectiveness of the method has been proven.

L Planning algorithms details

In this section we report all the parameters used in our implementations of Monte Carlo Tree Search (MCTS) [Kocsis and Szepesvári, 2006] and Cross Entropy Method (CEM) [Rubinstein, 1997], together with a brief explanation of the meaning of those parameters within the context of the two algorithms.

MCTS. At each time-step, we run I mcts simulations with MCTS to select the best action to play. At every simulation, starting from the root node, we select one action via the Upper-Confidence Bound formula for Trees (UCT)

where v i is the estimated value of node i , C is the exploration constant, N i is the visit count of the parent of node i , n i is the visit count of node i and ϵ is a factor offsetting the visit count. Once we select an unexplored action at one of the nodes, we expand the node that the action leads to

and perform a rollout with a random policy for up to max_actions to estimate its value. The value backpropagation is done as in standard MCTS and we use a discount factor of γ . The values of all parameters are reported in Table 13.

CEM. In this case, assuming deterministic environments, we plan directly for the next T cem timesteps, meaning that we choose the actions for up to T cem steps ahead, using the CEM algorithm. At every iteration we sample N cem action plans from a zero-mean Gaussian with dimensions T cem ×A and standard deviation for each dimension given by half the maximum absolute value between the upper and lower bounds for that action dimension (as it's always the case that each continuous action dimension is bounded in a box in the CWMB environments). The action plans are then clipped in the legal ranges of the action space and scored by their return as rollouts in the environment, starting from the current state. We then select the top K cem action plans (elites samples), fit the Gaussian parameters to them and repeat. At the last iteration, we return the top scoring action plan. All parameters are reported in Table 14.

M Computational Resources

In the following section we report as accurately as possible the computational resources used in this work. On the high level, the bulk of the computational costs, performed on an AMD cluster, was comprised of the experiments with Llama 3 on APPS, reported in Table 1. The reported experiments require running 3 times Llama 3 on 1000 problems, 20 times each, receiving approximately 1000 tokens in input and producing 1500 tokens in output (as the model is not good in using the End-ofSequence token to stop earlier). We split the runs in 100 array jobs, each taking approximately 15 hours and requiring 4 AMD MI250x each, for an estimated total of 18000 GPU hours.

Experiments on the CWMB were composed of 18 problems for which we ran our method, one baseline and 3 ablations, which should be roughly equivalent to a single experiment with 100 APPS problems, or 10 jobs of 15 hours with 4 GPUs, for a total of 600 GPU hours. The single experiment performed on RTFM with three different configurations also fits into this budget.

However, many more preliminary attempts were taken, so the full computational budget was of 31.800 GPU hours and a similar amount of CPU hours.

Furthermore, we have paid approximately $62.3 in OpenAI calls to GPT-3.5 Turbo (used only for prototyping) and GPT-4 Turbo (used with a budget of 10 calls on the CWMB experiments in Table 2, with 50 calls in some instances (Table 3) and for other preliminary experiments with GIF-MCTS).

Finally, all environment returns for planning were performed on a single consumer CPU in a few hours.

N Prompts

In this section we report the main prompts used for GIF-MCTS. These prompts are also shared by our WorldCoder implementation, while we avoid reporting explicitly the prompts used for Zero-shot CoT, as they are simply the problem description followed by "Let's think step by step".

N.1 APPS Prompts

<system>

You are an experienced Python developer. You will be provided with an incomplete code snippet from a Python program. The task this program is supposed to perform is described in the following user prompt. Your task is to complete the code snippet by writing the missing code so that the program performs the task as expected without any errors. You will be rewarded based on the number of test cases your code passes.

</system>

<user>

{PROB_DESCRIPTION}

Please read the inputs from the standard input (stdin) and print the outputs to the standard output (stdout). Output your code solution with the following format: ''python [your code] '' </user>

<assistant>

''python

{CODE_SO_FAR}

</assistant>

<system>

You are an experienced Python developer. You will be provided with an incorrect code snippet from a Python program. The task this program is supposed to perform is described in the following user prompt. Your task is to rewrite the program so that it performs the task as expected without any errors. You will be rewarded based on the number of test cases your code passes.

</system>

<user>

{PROB_DESCRIPTION}

Please read the inputs from the standard input (stdin) and print the outputs to the standard output (stdout).

First, write an explanation of the difference between the ground-truth output and the program's

output in the example provided. Secondly, point out the part of the code responsible for the

incorrect prediction and why its logic is erroneous. Third, suggest a concrete, actionable fix

for it. Finally fix the program in its entirety following the suggestion. The expected output is

in the format:

## Error explanation

[your explanation of the error]

## Error location and wrong logic

[where the error comes from and why]

## Fix suggestion

[how to fix the error]

## Correct code

''python

[your code]

''

## Incorrect code

You are provided with the following code snippet to fix.

''python

{CODE}

''

The code additionally makes a wrong prediction about this input.

## Input

## Ground-truth output

{OUTPUT}

## Code incorrect outputs

{PREDICTION}

</user>

<assistant>

## Error explanation

</assistant>

<system>

You are an experienced Python developer. You will be provided with an incorrect Python program. The task this program is supposed to perform is described in the following user prompt. Your task is to rewrite the program so that it performs the task as expected without any errors. You will be rewarded based on the number of test cases your code passes.

</system>

<user>

{PROB_DESCRIPTION}

Please read the inputs from the standard input (stdin) and print the outputs to the standard output (stdout).

First, write an explanation of the error and point out the part of the code responsible for the error and why its logic is erroneous. Second, suggest how you would fix the error, reasoning about the problem. Finally fix the program in its entirety following the suggestion. The expected output is in the format:

## Error explanation

[your explanation of the error]

## Fix suggestion

[how to fix the error]

## Correct code

''python

[your code]

''

## Incorrect code

You are provided with the following code snippet to fix.

''python

{CODE}

''

{ERROR}

</user>

<assistant>

## Error explanation

</assistant>

N.2 CWMBPrompts

<system>

You are an experienced Python developer. You will be provided with an incomplete code snippet from a Python program. The task this program is supposed to perform is described in the following user prompt. Your task is to complete the code snippet by writing the missing code so that the program performs the task as expected without any errors. You will be rewarded based on the number of test cases your code passes.

</system>

<user>

{ENV_DESCRIPTION}

## Class Definition

The class should be called "Environment". It should have at least:

- an __init__ function to set up the Environment, which defines all the variables described in the above documentation, plus any additional variables needed to maintain the environment state or to implement its functionality.

- a set_state function to set a custom value for the environment and its internal representation (you can assume that when "set_state" is used, the task is not done and internal variables should be set as a consequence). set_state takes a single argument as input: a state observation from the observation space defined above.

- a step function to predict a step in the environment. The input parameters for the step function are:

- An action, which must be contained in the action space described above.

The outputs required by the step function are:

- An observation, which must be contained in the observation space described above.

- The reward for taking the action, as described in the reward definition above.

- A boolean variable indicating if the episode is done.

## Important Notes

Only produce the environment class, containing the __init__, set_state and step functions and any additional functions you may need to complete this task. Do not write an example of how to use the class or anything else. Be careful about edge cases. Make sure to write all the required functions and that they have the exact names as specified in the task description. Missing or incorrectly named functions will not pass the tests and will result in a score of 0. It is of VITAL importance that you do not leave undefined any function, but implement each of them completely.

</user>

<assistant>

''python

{CODE_SO_FAR}

</assistant>

Figure 6: Prompt on the CWMB for the generate action.

<system> You are an experienced Python developer. You will be provided with an incorrect code snippet from a Python program. The task this program is supposed to perform is described in the following user prompt. Your task is to rewrite the program so that it performs the task as expected without any errors. You will be rewarded based on the number of test cases your code passes. </system>

<user> {ENV_DESCRIPTION}

## Class Definition

The class should be called "Environment". It should have at least:

- an __init__ function to set up the Environment, which defines all the variables described in

the above documentation, plus any additional variables needed to maintain the environment

state or to implement its functionality.

- a set_state function to set a custom value for the environment and its internal representation (you can assume that when "set_state" is used, the task is not done and internal variables should be set as a consequence). set_state takes a single argument as input: a state observation from the observation space defined above.

- a step function to predict a step in the environment. The input parameters for the step function

are:

- An action, which must be contained in the action space described above.

The outputs required by the step function are:

- An observation, which must be contained in the observation space described above.

- The reward for taking the action, as described in the reward definition above.

- A boolean variable indicating if the episode is done.

## Important Notes

Only produce the environment class, containing the __init__, set_state and step functions and any additional functions you may need to complete this task. Do not write an example of how to use the class or anything else. Be careful about edge cases. Make sure to write all the required functions and that they have the exact names as specified in the task description. Missing or incorrectly named functions will not pass the tests and will result in a score of 0. It is of VITAL importance that you do not leave undefined any function, but implement each of them completely.

First, write an explanation of the difference between the ground-truth transition and the step

function's outputs in the example provided. Second, point out the part of the code responsible

for the incorrect prediction and why its logic is erroneous. Third, suggest a concrete, actionable

fix for it. Finally, fix the program in its entirety following the suggestion. The expected output

is in the format:

## Error explanation

[your explanation of the error]

## Error location and wrong logic

[where the error comes from and why]

## Fix suggestion

[how to fix the error]

## Correct code

''python [your code] ''

## Incorrect code

You are provided with the following code snippet to fix.

''python {CODE} ''

The code additionally makes a wrong prediction about this input.

## Input

## Ground-truth output

{OUTPUT}

## Code incorrect outputs

{PREDICTION} </user>

<assistant> ## Error explanation </assistant>

<system>

You are an experienced Python developer. You will be provided with an incorrect Python program. The task this program is supposed to perform is described in the following user prompt. Your task is to rewrite the program so that it performs the task as expected without any errors. You will be rewarded based on the number of test cases your code passes.

</system>

<user>

{ENV_DESCRIPTION}

## Class Definition

The class should be called "Environment". It should have at least:

- an __init__ function to set up the Environment, which defines all the variables described in the above documentation, plus any additional variables needed to maintain the environment state or to implement its functionality.

- a set_state function to set a custom value for the environment and its internal representation (you can assume that when "set_state" is used, the task is not done and internal variables should be set as a consequence). set_state takes a single argument as input: a state observation from the observation space defined above.

- a step function to predict a step in the environment. The input parameters for the step function are:

- An action, which must be contained in the action space described above.

The outputs required by the step function are:

- An observation, which must be contained in the observation space described above.

- The reward for taking the action, as described in the reward definition above.

- A boolean variable indicating if the episode is done.

## Important Notes

Only produce the environment class, containing the __init__, set_state and step functions and any additional functions you may need to complete this task. Do not write an example of how to use the class or anything else. Be careful about edge cases. Make sure to write all the required functions and that they have the exact names as specified in the task description. Missing or incorrectly named functions will not pass the tests and will result in a score of 0. It is of VITAL importance that you do not leave undefined any function, but implement each of them completely.

First, write an explanation of the error and point out the part of the code responsible for the error and why its logic is erroneous. Second, suggest how you would fix the error, reasoning about the problem. Finally fix the program in its entirety following the suggestion. The expected output is in the format:

## Error explanation

[your explanation of the error]

## Fix suggestion

[how to fix the error]

## Correct code

''python

[your code]

''

## Incorrect code

You are provided with the following code snippet to fix.

''python

{CODE}

''

{ERROR}

</user>

<assistant>

## Error explanation

</assistant>

N.3 Sample Environment Descriptions

For the CWMB we extract the description for each environment directly from the Gymnasium source code 7 . We clean the description string found for each environment to remove irrelevant information (Arguments, Vectorized Environment, Version History, metadata) as well as manually remove mentions of external links or sources that may provide the LLM with an implementation of the environment. An example description for the CartPole-v1 environment 8 can be seen in Figure 9.

## Description

A pole is attached by an un-actuated joint to a cart, which moves along a frictionless track. The pendulum is placed upright on the cart and the goal is to balance the pole by applying forces in the left and right direction on the cart.

## Action Space

The action is a 'ndarray' with shape '(1,)' which can take values '0, 1' indicating the direction of the fixed force the cart is pushed with.

0: Push cart to the left - 1: Push cart to the right

**Note**: The velocity that is reduced or increased by the applied force is not fixed and it depends on the angle the pole is pointing. The center of gravity of the pole varies the amount of energy needed to move the cart underneath it

## Observation Space

The observation is a 'ndarray' with shape '(4,)' with the values corresponding to the following

positions and velocities:

| Num | Observation | Min | Max |

|--|--------|-------|-------|

| 0 | Cart Position | -4.8 | 4.8 |

| 1 | Cart Velocity | -Inf | Inf |

| 2 | Pole Angle | -0.418 rad (-24°) | 0.418 rad (24°) |

| 3 | Pole Angular Velocity | -Inf | Inf |

**Note:** While the ranges above denote the possible values for observation space of each element, it is not reflective of the allowed values of the state space in an unterminated episode. Particularly: - The cart x-position (index 0) can be take values between '(-4.8, 4.8)', but the episode terminates if the cart leaves the '(-2.4, 2.4)' range. - The pole angle can be observed between '(-.418, .418)' radians (or **±24°**), but the episode terminates if the pole angle is not in the range '(-.2095, .2095)' (or **±12°**) ## Rewards

Since the goal is to keep the pole upright for as long as possible, a reward of '+1' for every step taken, including the termination step, is allotted. The threshold for rewards is 500 for v1 and 200 for v0.

## Starting State

All observations are assigned a uniformly random value in '(-0.05, 0.05)'

## Episode End

The episode ends if any one of the following occurs:

1. Termination: Pole Angle is greater than ±12° 2. Termination: Cart Position is greater than ±2.4 (center of the cart reaches the edge of the display) 3. Truncation: Episode length is greater than 500 (200 for v0)

O Examples of Generated Programs

We report examples of Code World Models generated by both Llama 3 and GPT-4 Turbo. We include CartPole-v1 as an example of a discrete environment, Ant-v4 as an example of a continuous environment and RTFM.

Figure 16: Code World Model of the Ant-v4 environment generated by GIF-MCTS with GPT-4. For reference, the official implementation of the environment can be found at the official Gymnasium GitHub repository of the Farama Foundation (at gymnasium.envs.mujoco.Ant_v4 ).

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: All claims in Abstract and Introduction are backed by experiments found in the Experiments Section 5 or in the Appendix, and clearly make reference to them.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: Limitations discussed in the Limitations Section 6.1.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [NA]

Justification: No theoretical results are presented in this work.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We comprehensively report, either in the main text of in the appendix, all parameters, prompts and experimental details to reproduce our results.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: Anonymized code will be provided in zip file together with the submission and released with URL referenced in the paper upon acceptance of this work.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: All details of our experiments are specified. We will also release, together with the code, the data collected in RL environments, as they present stochasticity in terms of initial conditions.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: Errors are reported on all the main experimental tables. All errors are errors of the mean value, not standard deviations.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: All computational resources have been accounted for in Section M of the appendix.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: The authors confirm that the research conducted in the paper conforms, in every respect, with the NeurIPS Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: We report a discussion on potential positive and negative impacts in Appendix A.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: To the best of our knowledge, no risk is posed by this work.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: All benchmarks, Reinforcement Learning environments and models used have been properly cited and were open to use for research purposes.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: We do introduce a new benchmark derived from existing openly available resources, detailing how to re-create it and reporting its details in appendix. The opensourcing of the code for this work will contain instructions on how to run the benchmark.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper explores a novel and promising direction of model-based RL, proposing to represent the dynamic model and the reward model using code world models, which are Python programs that can be executed and rollout environments. With such code world models, we can learn a policy to maximize the return predicted by the models. To achieve this, the paper proposes a framework, Generate, Improve and Fix with Monte Carlo Tree Search (GIF-MCTS), that can synthesize code world models from a pre-collected dataset of environment interactions. This paper also introduces a benchmark for evaluating the performance of code world models. The experiments show that the proposed framework GIF-MCTS can reliably model environments and allow for learning policies that achieve reasonable performance. I believe this work explores an interesting and promising research direction. Yet, I am concerned with the assumptions of offline datasets, the missing comparisons to offline RL methods, and the insufficient description of related work. Hence, I am slightly leaning toward rejection at this moment, but I am willing to increase my score if my concerns are addressed in the author's rebuttal.\n\nSTRENGTHS:\n**Motivation and intuition**\n- The motivation for modeling environments using codes is interesting and convincing, and this work presents an effective framework to achieve it.\n\n**Technical contribution**\n- Given that WorldCode has explored the online setup, this work explores the offline setup, which offers an unique perspective to the code world model idea.\n\n**Clarity**\n- The overall writing is clear. The authors utilize figures well to illustrate the ideas.\n\nWEAKNESSES:\n**Dataset collection**\nThe authors stated that the dataset \"includes at least some low-scoring and some relatively high-scoring behavior.\" I am wondering how we can do this without accessing true reward functions. Given a new environment that we have not seen before, how could we collect such a dataset to learn a code world model? \n\n**The effect of dataset quality**\nIt is unclear how the quality of the dataset would affect the code world model performance. Can we learn a reasonable code world model with just trajectories collected by a random policy? \n\n**Online setting**\nThis paper focuses on an offline setup, where the code world model learns from a pre-collected offline dataset. I am curious how we can extend the framework so that it can improve the code world model and the policy constructed from the offline dataset, just like how the research studies how to fine-tune policies learned using offline RL algorithms.\n\n**Access to physics engines as a tool**\nThe proposed framework writes a code to represent a world model for each environment from scratch. It could be difficult to precisely simulate some interaction that requires modeling physics. It would be interesting to see if granting access to physics engines, e.g., calling MuJoCo, could generalize to physics-intensive environments.\n\n**Comparison to offline RL methods**\nThis work mainly compares the proposed framework with WorldCoder, which also builds a world model. I am curious about how the proposed framework compares to offline RL methods, which simply learn a policy from the dataset to maximize the return.\n\n**Related work: programmatic RL**\nThe related work section is not comprehensive. Highly relevant programmatic RL works, which represent policies as programs and interact with environments by executing the programs, are missed from the section, such as:\n- Programmatically Interpretable Reinforcement Learning (ICML 2018)\n- Imitation-projected programmatic reinforcement learning (NeurIPS 2019)\n- Programmatic Reinforcement Learning without Oracles (ICLR 2022)\n- Learning to synthesize programs as interpretable and generalizable policies (NeurIPS 2021)\n- Hierarchical programmatic reinforcement learning via learning to compose programs (ICML 2023)\n- Show me the way! Bilevel search for synthesizing programmatic strategies (AAAI 2023)\n- Reclaiming the source of programmatic policies: Programmatic versus latent spaces (ICLR 2024)\n\n**Clarity: offline setup**\nWhen I read the introduction and the abstract, I thought it focused on an online model-based RL setting. Then, I realized this is an offline setup while reading Section 3. I suggest the authors revise the introduction and the abstract to align the expectations.\n\nQUESTIONS:\nSee above\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors present a search algorithm based on Monte Carlo tree search to synthesize programs with LLMs. The contribution includes formulating a search space that is compatible with the functionalities of a LLM: the actions in the search tree include generating new lines of code, fixing bugs, and improving current implementation. \n\nThe system is evaluated on the synthesis of programs for solving competition-level problems and programs encoding world models. As baselines, the experiments include recent works such as CodeRL, Parsel, and WorldCoder.\n\nSTRENGTHS:\nAfter reading the paper once, I believe its main contribution is to show the power of tree search in the context of synthesis with LLMs. A simple MCTS-based algorithm already produces programs able to attain better quality metrics than other methods not using search.\n\nThe experiments are nicely done from an ablation perspective: I quite enjoyed Table 6 in the appendix, where different versions of GIF-MCTS are evaluated.\n\nWEAKNESSES:\nThe paper has 4 weaknesses that are worth mentioning, so the authors can react to them in the rebuttal.\n\n**1. Clarity of the paper could be much improved.**\n\nWhen reading Figure 2 and the description around it, I got confused with the use of the word \"rollout.\" I think it is being used to mean two different things. The first is a portion of the code that a node represents. The second is the LLM call that generates that portion.\n\nI did not find in the paper the specification in terms of how many lines are generated in each \"generate new lines\" action and how many are generated in the rollouts.\n\n**2. The paper overclaims**\n\n> we improve the trade-off between exploration and exploitation in the UCT formula used for action selection.\n\nWhat was probably meant was that the paper uses a heuristic that achieves better empirical results. There are no theoretical improvements in the trade-off.\n\n> Calling code instead of LLMs for planning has the advantages of being precise, reliable, interpretable, and extremely efficient.\n\nNone of these properties are actually properly evaluated in the paper. This might be a bit of nitpicking, but code can also be imprecise, unreliable, uninterpretable, and extremely inefficient. Anyone working as a programmer has already experienced all these negative properties of code.\n\n**3. Lack of discussion on data contamination**\n\nWith the exception of the RTFM benchmark, nothing is said or discussed about possible data contamination. Is the search algorithm helping the LLM remember what is in its training data or is the overall solution doing something else? Are the models for RL tasks available online as code? Are the solutions to the programming problems available online?\n\nThe lack of discussion about data contamination brings me to the question of what is the research question behind this work. Is it a retrieval question (i.e., can the system recover information used in the LLM's training data) or a reasoning question (i.e., can the system write programs that solve problems)?\n\n**4. Statistical significance is unclear**\n\nIt is not clear whether the search was run only once or many times. The averages presented seem to be over a number of problems and episodes, but not on the number of times the entire system was executed.\n\nQUESTIONS:\n1. Please explain the issue around the use of the word \"rollout\" and how the number of lines of code is determined in each call of the \"generate new lines\" actions.\n2. How about running the fix bug action a number of times whenever a state representing a buggy program is generated? Instead of allowing MCTS to decide when to debug, why not debug right away?\n3. Please comment on the data contamination issue mentioned above and clarify what is the research question behind this work.\n4. Please comment on the statistical metrics and empirical design of the experiments.\n\nLIMITATIONS:\nIn weaknesses, please see the parts related to overclaiming and data contamination.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper primarily investigates the application of LLMs to synthesizing world models for reinforcement learning environments, where the world model is expressed as a Python program implementing the state transition and reward functions. Concretely, starting from available environment simulation code (e.g., a benchmark environment using the MuJoCo physics simulator), the authors first obtain a number of trajectories. An off-the-shelf LLM is then prompted with to reimplement the environment, starting from a textual description (and without explicit access to the original code), and the collected trajectories are used as test cases to verify correctness. This is a challenging code synthesis problem for which the authors propose an iterative algorithm to refine and improve the generated program with multiple LLM calls, formulated as monte-carlo tree search.\n\nTo me, the contributions of this paper are:\n- The GIF-MCTS method for effective iterative program synthesis wrt available test cases, which is separately evaluated on the APPS code generation benchmark.\n- The concept of code world models (as outlined above) which shows promise in efficiently obtaining world models from observational data\n- A benchmark which covers common environments used in Deep RL that enables the community to propose further advances\n\nSTRENGTHS:\nI found the main idea of the paper to be quite original and interesting, i.e., generating the code for a world model rather than utilizing LLMs directly. The authors show that this works quite well in environments with discrete action space: they often obtain perfect or near-perfect prediction accuracy as well as rewards approaching those achieved by planning with an oracle world model. In continuous environments, where the original implementations rely on a physics simulation engine, the results are less convincing -- this is understandable to me as I would expect LLMs (or humans, for that matter) having trouble reimplementing those without access to the original or a substitute simulation engine. Whether code world models are a promising direction to ultimately tackle planning in complex or real-life environments remains to be seen, of course; however, I would expect paper at hand and the proposed benchmark to inspire further works in this relevant direction.\n\nI also liked the separate evaluation of their MCTS code synthesis method on APPS where it produces strong results, and it's worth noting that there are full papers proposing similar iterative prompting for code generation. I hope the authors make this part of the paper available for easy re-use.\n\nThe paper itself is well-written and, while it is pretty packed, I found that I could easily locate and understand the method and details on problem setup and evaluations.\n\nWEAKNESSES:\nI would appreciate if the authors included examples of LLM-generated environment implementations in the appendix. I would be interested in how the continuous action space environments are being tackled, or to check for overlap with the original environment implementations (which\n\nFor the planning results it would be good to list the true rewards obtained along with the rewards of a random policy and SOTA model-free/model-based results from the RL literature. While beating SOTA results here is not the focus of the paper, it would provide a helpful perspective for readers familiar with those environments.\n\nQUESTIONS:\nI am not totally clear on the LLM sample budget for the evaluation on APPS. From L258 ff: \"B is the budget for the number of LLM calls\", Table 1 lists pass@20, but L618 in the appendix mentions 50 Llama 3 calls per problem. L618 also mentions 1000 problems, whereas the competition split on which you report results is smaller.\n\nLIMITATIONS:\nI found the paper's discussion of its limitations to be upfront and comprehensive and have nothing to add.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes Generate, Improve, and Fix with Monte Carlo Tree Search (GIF-MCTS) for generating Code World Models (CWMs) using Large Language Models (LLMs). The authors present code representations of reinforcement learning (RL) environments, enabling the application of LLM algorithms for code generation across various domains. They introduce the Code World Models Benchmark (CWMB), which comprises 18 diverse RL environments to evaluate their approach. GIF-MCTS demonstrates better performance compared to the baselines of APPS, CWMB, and a grid-world environment.\n\nSTRENGTHS:\nThe paper is well-written and effectively demonstrates that MCTS can facilitate agentic behaviors and problem-solving across multiple domains. The motivation and the rationale are clear, as we expect the idea of agentic behavior of searching with diverse actions can help solve more complex problems.\n\nThe appendix includes extensive ablative studies that offer a detailed analysis of the method's components.\n\nWEAKNESSES:\n**Reliance on pre-training.** I didn’t fully understand why learning the world model is necessary. MCTS is an inference-time algorithm. The reliance on pre-training for the large language models might be a limitation.\n\nSpecifically, does GIF-MCTS need to predict all transitions well in the offline dataset to perform well on a task? For coding, the agent doesn’t need to accurately predict how a human solves a problem. It’s sufficient if the agent can find a correct solution, possibly using MCTS.\n\n**Novelty.** There seems to be limited novelty in this framework. Related works have explored agentic behaviors of using “implementing a function”, “fixing bug” as actions, like SWE-agent [1]. Using tree search and sequential decision-making are also explored in the literature [2]. The novelty seems to be a combination of both.\n\nReferences\n\n[1] Yang, John, et al. \"Swe-agent: Agent-computer interfaces enable automated software engineering.\" arXiv preprint arXiv:2405.15793 (2024).\n\n[2] Zhang, Shun, et al. \"Planning with large language models for code generation.\" arXiv preprint arXiv:2303.05510 (2023).\n\nQUESTIONS:\nCan the authors clarify why GIF-MCTS is better than WorldCoder? Is it because of the advantage of the tree search algorithm (while WorldCoder is mainly a sampling algorithm)?\n\nOffline pre-training seems to be expensive. For domains like code generation, is pre-training necessary? For new domains that the agent may not have seen in the pre-training data, does it help to use in-context learning to show the agent some example trajectories?\n\nLIMITATIONS:\nThe limitations are clearly discussed in the paper.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}