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

Authors: nicola  dainese
matteo  merler
minttu  alakuijala
pekka  marttinen
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 21:34:47

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.

Main paper: 
 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] .Figure 1 : Overview of the Code World Models (CWM) framework. Given the description of an environment and a task, we use an LLM guided by the GIF-MCTS method to iteratively generate and refine a candidate CWM. The candidate's correctness is evaluated by checking if it correctly predicts a set of trajectories collected from the true environment. If the model cannot fully predict all transitions, the fraction of correct predictions and other information are given as feedback to the LLM and the cycle repeats. After matching all transitions or having used up a computational budget, the best CWM is returned and used to solve the task via model-based planning.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. 
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 WorldCoder3 [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.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. 
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 Figure 2 : Example of a GIF-MCTS tree for generating a CWM. Starting from the root of the tree, every action taken corresponds to 1) prompting the LLM to either generate, improve or fix a CWM, 2) parsing the LLM completion, and 3) evaluating the CWM's correctness using the available environment trajectories as unit tests (presented as a percentage inside the nodes). On buggy nodes, we allow only fix actions for up to f sequential attempts and replace the actual value with a temporary one, represented in red. In healthy nodes we allow only generate and improve actions. All action prompts are exemplified on the right. The number of total fix f attempts is a model hyperparameter, set to three in this Figure and for our method.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:EQUATIONand 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:EQUATION)where 1 is the indicator function (equals to one if the pair is matching, zero otherwise) and ŝ′ i , ri and di 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.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). 
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.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.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 . [Zelikman et al., 2023] code-davinci-002 N/A 25.50 pass@any Zero-shot CoT * [Kojima et al., 2022] Llama 3 70B 23.2±1.3 pass@20 WorldCoder * [Tang et al., 2024b] Llama 3 70B 25.1±1.4 pass@20 GIF-MCTS (ours) Llama 3 70B 28.3±1.4 pass@20 * Our re-implementation.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 solutions4 , while GIF-MCTS outperforms it by evaluating only 20 different programs.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:EQUATIONwhere 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).Table 2 : CWMB: main results. For each method, we report the CWM accuracy and the normalized return R, averaged separately across environments with discrete and continuous action spaces. Budget indicates the number of LLM calls. For each metric, we report the mean value across environments (and for the return, also across 10 episodes) with its error. For Llama 3, we report an average of three different random seeds for additional statistical significance. [Tang et al., 2024b] .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.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. GPT-4 Turbo GIF-MCTS (ours) 50 1.00 ± 0.00 1.00 ± 0.00 WorldCoder * 50 0.64 ± 0.02 -0.06 ± 0.12 * Our re-implementation of [Tang et al., 2024b] .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. 
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.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.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. 
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. 
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. 
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. 
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. 
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. 
 
 Appendix: 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.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:UCT(node i ) = v i + C • ln N i n a=ai + ϵ ,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=ai 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=ai 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:v i = w G • v G + w L • v L w G + w L ,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 documentation5 for an accurate description of each LLM parameter. 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.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 Table 6 : CWMB results: ablation study. We compare the full GIF-MCTS method against three ablated variants, each leaving out one of the three action types. For each method, we report the CWM accuracy and the normalized return R, averaged separately across environments with discrete and continuous action spaces. For each metric we report the mean value across environments (and for the return, also across 10 episodes) with its error. 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.Table 7 : Qualitative Analysis. We report a qualitative study for the frequency with which GIF-MCTS chooses each type of action on average. The first section of the table is considering the whole tree, while the second section (path quantities) only consider the path from the root node to the node with the highest value (where the code used as the environment was generated). 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.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.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.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 Verma et al. [ , 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 VAE. 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.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.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 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 .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, Table 10 : CWMB results. Individual results for each environment in the CWMB using Llama 3 (we report the results for the first seed only).Action Space GIF-MCTS WorldCoder 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. 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.Accuracy (↑) R(↑) Accuracy (↑) R(↑) CartPole-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.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)EQUATION)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 . 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-of-Sequence 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.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> Figure 3 : Prompt on the APPS benchmark 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> {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. <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. -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. -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. For the CWMB we extract the description for each environment directly from the Gymnasium source code7 . 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 environment8 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.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 |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. For reference, the official implementation of the environment can be found at the official Gymnasium GitHub repository of the Farama Foundation (at gymnasium.envs.classic_control.cartpole).1 import numpy as np For reference, the official implementation of the environment can be found at the official Gymnasium GitHub repository of the Farama Foundation (at gymnasium.envs.classic_control.cartpole).1 import numpy as np • 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).We release our code at https://github.com/nicoladainese96/code-world-models.Due to the timing of our experiments, which were performed in April and May 2024, we replicate the results from the first version of the WorldCoder paper, which can be found at https://arxiv.org/abs/2402.12275v1. The authors have since developed a slightly different algorithm for code generation, which was published after we finalized our experiments. The original code generation algorithm based on Thompson Sampling, which we call WorldCoder in this work, was later published inTang et al. [2024a].Results reported for Parsel use 8 pseudo-codes per problem, each implementing n sub-functions (with n being problem-dependent) 16 times and then evaluating up to 8 * 16 n sub-functions combinations against APPS unit tests and keeping the best result.https://huggingface.co/docs/transformers/main_classes/text_generationhttps://pypi.org/project/tiktoken/https://github.com/Farama-Foundation/Gymnasium/tree/main/gymnasium/envshttps://github.com/Farama-Foundation/Gymnasium/blob/main/gymnasium/envs/classic_ control/cartpole.py'gleaming': 'poison', 'fanatical': 'cold', 'mysterious': 'poison', 


Decision: Many thanks to the authors for submitting this work and engaging in discussion with reviewers. I appreciate that the authors introduced an innovative MCTS approach for using LLM as an agent in an RL-based setup. Specifically, the authors introduced Code World Models to learn RL world models formulated as code/programs to facilitate LLM agents. The authors reported comprehensive experimental results and ablation analysis to demonstrate the effectiveness of the proposed method. The authors also contributed a new benchmark, including applicable environments for optimizing and evaluating RL agents.

As noted by reviewers, there are some minor issues: (1) the tasks for evaluating code world models are limited to simple domains; it is still an opening question of how effectively Code World Models might work in more complex and practical scenarios.  (2) secondly, some presentation issues should be addressed e.g. Figure 2 and the offline setup,  and claims that need to be adjusted (as pointed out by Reviewer 3vW3). 

In conclusion, I recommend accepting this paper. I encourage the authors to review all the above-mentioned points and other feedback from reviewers and revise their paper accordingly.

Reviews: review: paper_summary: paper_summary: This 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.
strength_weakness: strength_weakness: **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.

Specifically, 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.

**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.

References

[1] Yang, John, et al. "Swe-agent: Agent-computer interfaces enable automated software engineering." arXiv preprint arXiv:2405.15793 (2024).

[2] Zhang, Shun, et al. "Planning with large language models for code generation." arXiv preprint arXiv:2303.05510 (2023).
questions: questions: Can 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)?

Offline 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?
limitations: limitations: The limitations are clearly discussed in the paper.
paper_summary: paper_summary: This 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.

To me, the contributions of this paper are:
- The GIF-MCTS method for effective iterative program synthesis wrt available test cases, which is separately evaluated on the APPS code generation benchmark.
- The concept of code world models (as outlined above) which shows promise in efficiently obtaining world models from observational data
- A benchmark which covers common environments used in Deep RL that enables the community to propose further advances
strength_weakness: strength_weakness: I 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

For 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.
questions: questions: I 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.
limitations: limitations: I found the paper's discussion of its limitations to be upfront and comprehensive and have nothing to add.
paper_summary: paper_summary: The 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. 

The 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.
strength_weakness: strength_weakness: The paper has 4 weaknesses that are worth mentioning, so the authors can react to them in the rebuttal.

**1. Clarity of the paper could be much improved.**

When 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.

I 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.

**2. The paper overclaims**

> we improve the trade-off between exploration and exploitation in the UCT formula used for action selection.

What was probably meant was that the paper uses a heuristic that achieves better empirical results. There are no theoretical improvements in the trade-off.

> Calling code instead of LLMs for planning has the advantages of being precise, reliable, interpretable, and extremely efficient.

None 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.

**3. Lack of discussion on data contamination**

With 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?

The 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)?

**4. Statistical significance is unclear**

It 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.
questions: questions: 1. 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.
2. 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?
3. Please comment on the data contamination issue mentioned above and clarify what is the research question behind this work.
4. Please comment on the statistical metrics and empirical design of the experiments.
limitations: limitations: In weaknesses, please see the parts related to overclaiming and data contamination.
paper_summary: paper_summary: This 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.
strength_weakness: strength_weakness: **Dataset collection**
The 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? 

**The effect of dataset quality**
It 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? 

**Online setting**
This 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.

**Access to physics engines as a tool**
The 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.

**Comparison to offline RL methods**
This 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.

**Related work: programmatic RL**
The 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:
- Programmatically Interpretable Reinforcement Learning (ICML 2018)
- Imitation-projected programmatic reinforcement learning (NeurIPS 2019)
- Programmatic Reinforcement Learning without Oracles (ICLR 2022)
- Learning to synthesize programs as interpretable and generalizable policies (NeurIPS 2021)
- Hierarchical programmatic reinforcement learning via learning to compose programs (ICML 2023)
- Show me the way! Bilevel search for synthesizing programmatic strategies (AAAI 2023)
- Reclaiming the source of programmatic policies: Programmatic versus latent spaces (ICLR 2024)

**Clarity: offline setup**
When 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.
questions: questions: See above
limitations: limitations: Yes

Comments: title: Thanks for your response
Thanks for the response
Re: Rebuttal by Authors
comment: Thank you for clarifying the issues raised in my review. 

I have increased the overall score of my review.
Thanks for confirming that my understanding is correct. Overall I believe learning a world model represented by code is an interesting and promising idea, which is also validated empirically. I believe the paper would benefit from improved clarity and the new results discussed in the rebuttal.

I have raised the score to 6 (Weak Accept).
Thank you for your careful review and for taking the time to clarify your understanding. We are glad that the confusion was cleared away and we can confirm that your points are accurate. 

We will clarify that GIF-MCTS is a tree search-based code generation algorithm and that we use APPS to evaluate its performance on the general code synthesis task, but that APPS and CWMs are a related but separate application and that the method is framed differently for APPS (for example the dataset D becomes the collection of unit tests). 

We hope that this helps to explain the novelty of the paper: in our view, the central contribution is the concept of a Code World Model and the possible use cases for RL, while GIF-MCTS is a tree search algorithm we designed specifically for the task of synthesizing better CWMs, with the evaluation on APPS serving the purpose of generally evaluating the method on a widely used benchmark.

Please let us know if there are any concerns that remained unaddressed and thank you again for your time and effort.
Thank you for your thoughtful consideration and for increasing your score. We are committed to revising the paper as promised, ensuring that the offline setup is clearly stated in the abstract and introduction, and including the necessary comparisons to offline RL methods. We appreciate your feedback and will make the necessary improvements.
Thank you for the responses and clarifications. The clarification on pre-training does clear up my confusion. I thought the training data set $D$ (line 148) is used to train a language model to generate world models. That was a misunderstanding.

I wanted to make sure thatm I am following the logic of the paper correctly. Could you confirm if the following points are accurate?

- GIF-MCTS is a new tree search algorithm evaluated on the APPS dataset, where it outperforms baseline algorithms like WorldCoder.
- GIF-MCTS is then applied to Code World Model learning and evaluated on the Code World Models Benchmark and Read to Fight Monsters. The results show that GIF-MCTS learns a better world model than WorldCoder, and planning with the GIF-MCTS-learned world models achieves better returns.

If these points are correct, it would be helpful to clarify them in the paper. Initially, I thought we needed to learn a world model for APPS, which confused me.
Thank you for the rebuttal, which addresses some of my concerns. I am increasing my score to 5 (borderline accept), counting on the authors will revise the paper according to the promises, e.g., being upfront about the offline setup in the abstract and the introduction, including the comparisons to offline RL methods, etc.