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 2024 poster (2024)
Keywords: Large Language Models
code generation
MCTS
model-based reinforcement learning
Primary area: natural_language_processing
TL;DR: We propose to model RL environments with code written by an LLM, propose a method to improve code generation for this task and show how to plan with code world models.

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

Generating Code World Models with Large Language Models Guided by Monte Carlo Tree Search
Nicola Dainese"
Matteo Merler'
Department of Computer Science
Department of Computer Science
Aalto University
Aalto University matteo. www.mmeaalto......mmmmmm .i
nicola. doooolaadiieeeeaaattoo.cc
Minttu Alakuijala
Pekka Marttinen
Department of Computer Science
Department of Computer Science
Aalto University
Aalto University
minttu. alakuijala@@atto .i
pekka. www.maatnttaaatoo.....oo
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 Reinforce- ment 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 instruc- tions, to generate exact code with non-trivia logic and to self-debug long program with feedback from unit tests and environment trajectories. To address these chal- lenges, we propose Generate, Improve and Fix with Monte Carlo Tree Search (GIF-MCTS). a new code generation strategy for L.LMs. 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 envi- ronments paired with corresponding textual descriptions and curated trajectories. GIF-MCTS surpasses all baselines 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- model-based RL agents with greatly improved sample efficiency and inference speed.
Introduction
The ability to model the world is essential for goal-oriented intelligent agents [Ha and Schmidhuber, 2018]. When faced with novel environment, the agent must quickly understand its mechanics achieve its goal, for example by building an internal representation of the world and planning with 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 new task the agent in natural language particularly promising, and multiple works explore instruction-follont important information can be communicated in the form of imperative instructions. Many key facts required to solve task involve understanding observations, predicting outcomes of different actions and determining whether those outcomes align with the agent's s goals. Thus, systems capable of leveraging additional descriptive information, such model- Reinforcement Learning (RL) agents, have greater potential for fast and efficient adaptation via natural language [Lin al., 2024].
agents [Jang et al., 2022, Ahn et al., 2022]. However, not all
Asterisk indicates equal contribution.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Validate Code World Model
Trajectoriss collected from the ervironment
clans invironeanss ded Lnde (ante): sel peli-oor__politeee
predict correctly?
Yes
Planning with the Code World Model
Code generation with GIF-MCTS
Candidate Code World Model
description
Feedback
wrong predictions
Large Language Model
Overview of the Code World Models (CWM) framework. Given the description of an environment and a task, and refine a candidate CWM. The candidate's i correctness is evaluated by checking if it correctly predicts of trajectories collected from the true environment. If the model cannot fully prediet 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.
Figure 1:
an LLM guided by the GIF-MCTS method to iteratively generate
Large Language Models (LLMs) have revolutionized the field of Natural Language Processing, and offer great opportunities for world modeling, thanks their internet-scale knowledge, reasoning, and instruction- following abilities. However, it is not clear how to best combine LLMSs and world models. One option 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 be called repeatedly in planning loop due to their model- bbaed RL agents [Dainese et al.,
corruuctiooooolooiinng
high inference cost. On the other hand, ianguage-conditoons 2023, Lin 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 LLM restricted textual inputs and outputs, limiting their applicability in RL.
world model directly, but they are slow for inference and
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 more complex world models (Section 5). Alongside this paradigm, we introduce the Code World Models Benchmark (CWMB), consisting of 18 diverse RI. 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), 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.? We evaluate the performance of our method on three benchmarks: the new CWMB, the Competition split on APPS [Hendrycks et al., 2021], popular and challenging coding benchmark, and RTFM [Zhong et al., 2020], language-conditioned grid-world, showcasing environments with varying characteristics and complexity. GIF-MCTS
We release code http://shiiaccmiiioooooasssoesoo

---

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 L.MM as world model, as shown Appendix H), while, provided the CWM is accurate, matching the performance of an oracle planner with access to the real-world model. Finally, we discuss the limitations and challenges to overcome to make Code World Models more broadly applicable.
2 Related Work
World models with code. Code is 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 's dynamics and previous works often uses different techniques to build a full world model. AutumnSynth [Das et al., 2021] uses 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, 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 low-level wanguage-conditiona controller to be carried out. LLMs have also been used to generate a model based on probabilistic code [Wong et al., 2023].
Most similar to our approach, the concurrently proposed WorldCoder' [Tang et al., 2024b] also leverages LLMs to generate a Python-based world model. WorldCoder chooses program to refine from a working set of programs using the classical Thompson Sampling bandit algorithm [Thoompoon, 1933, Katchakis and Veinott, 1987], informed by a Beta prior, to iteratively learn a world model from gathered experience. Tang et al. focus 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 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 code generation method, GIF-MCTS, achieving state-of-the-an 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 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 work. A host of prompting techniques have shown how to leverage the In-Context Learning (ICL) [Brown et al., 2020] abilitics of LLMs to enhance a model's '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 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, has limita- tions in code-specifi. applications like ours. For instance, it generates programs simultaneously from the same node, rather than sequentially, which does not fully exploit the sequential nature of MCTS. Additionally, it uses 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.
*Due 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 found at htass//axxivoor//ab/////00 12275vv.. 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 in Tang et al. [2024a].
3

---

Previous research has also focused on pseudocode-based reasoning, such as Parsel [Zelikman et al., 2023], which uses a custom pseudocode language to decompose the program into independent problems that can be solved separately. In contrast, we focus on the sequential refinement of solutions using a variant of MCTS and the environment's s feedhack to produce directly executable Python code that can be leveraged in model-based RI..
We refer the reader to Appendix G for further discussion on works that build language-condttioeed world models but do not use code and on works that use programs as policies in RL.
3 Code World Models
In this Section, we first introduce the Code World Models framework and then the proposed Code World Models Benchmark.
Code World Models framework. Following the model-based Reinforcement Learning problem setting, we consider an environment represented by a Markov Decision Process with state space S, action space A, a transition function p(s' la, s), and a scalar reward function R(s, a, s'), with 8, E S indicating respectively the current and next state, and a E A being the action taken from the current state. The task of world model is to accurately represent p and R. We make the following assumptions: 1) the environments are deterministic and fully observable, and 2) we are provided with a natural language description of the environment, which detailed enough to infer the observation space as well as the logic of the transition and reward functions.
The first assumption implies deterministic transition function s' ='= f(s, a), rather than probabilistic one as in the general case; we address this limitation in Section 6.1. The second assumption is akin to the situation where human would be provided with an explanation, 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 al., 2020]. assuming that a dataset D of one time-step transitions ((s, a, s', d)x) 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 suitable planning algorithm to collect more trajectories from the environment, turning the algorithm into online RI, as done Tang et al. [2024b].
tutorial, about task that they
Code World Models Benchmark. To comprehensively test world model generation for a variety of environments, we define 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 MuloCo tasks. The environments' Python imple- mentations as well as their documentation are adapted from the Gymnasium library [TTowers 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 1).
For each environment, collect a training dataset DD of past trajectories. We curate 'D so that includes at least some low-scoring and some relatively high-scoring behavior. However, we neither attempt maximally cover the state space do require optimal demonstrations. We aim to show that relatively low annotation effort is required to build CWMs: for the majority of environments, 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 language description of each environment, derived from the documentation written for Gymnasium's '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.
collect just 5 trajectories equivalent to taking random actions and further
4 GIF-MCTS
In this Section, we first specify the format of the Code World Models that we consider in this work and how we evaluate their accuracy. We then present Generate, Improve and Fix with Monte Carlo Tree Search (GIF-MCTS), novel approach to leverage LLMs for code generation via multiple

---

Generate lines
SEE (ea2d):
Enpaet cespy
Previous
g1
g2
LLM
ap. caaden costnm--t.4 4.4, 1.4)
sbll
Temparary valuo assigned bugzy nodes
Addad
ires
25%
0%
99%
det steloriky
Used si
rolout
Fix bugs
clase Enxieocaents fef Jedt (we22):
62%
0%
66%
cest pesttion
Actusl value
cacdoo www.eeeeeeee.e........ 0.17
backpropageted
f1
cRss Vedeeasy
bug
stepCrelf sectende
cac.o........................... cww.aatteee 0.17 sell 2224 acguine veloeisy
state, rewird,
62%
NaEABEEOES
'ap"
g1
i2
i1
Improve predictions
Irediction facled:
58%
62%
100%
Sensenenects
def 1e0t, Ceedd3: seld predtiee
.G02Ed SEutE:
LEGEND buacy/
stepCenDe actSoe3:
remaining "ix attompts
fef setterl:
LLM
STATO, CEHECO,
codo evaluation
mode
rewerd, Core
Figure 2: Example of GIF-MCTS tree for generating 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, allow only fix actions for up to ʃ 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 attempts is model hyperparameter, set to three in this Figure and for our method.
sequential attempts in the presence of feedback, specifically tailored the needs of building Code World Models.
We formulate the task of synthesizing Code World Model as that of writing Python Environment class with a step() function that jointly implements the transition and reward functions:
(8' f. d) code bode__enviroment.. wttpiieettttttttttttttee
(1)
and consider a Code World Model correctly synthesized it correctly reproduces all transitions 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:
Σ (1165,511
1 1 1r, Al 5 Idd.
A=
(2)
where 1 the indicator function (equals one if the pair is matching, zero otherwise) and 8/, A and d, 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 parent node produces new complete program, which is split into a state part and rollowt part and stored in a child node. The child node's state is formed from the parent's s state by appending I. additional lines of code (we set I, = 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.c., run) the code. This is novel formulation of the state of node, store
5

---

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] select which action to take. 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 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 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 heuristic that empirically improves the trade-off between exploration and exploitation in the UCT formula used for action selection, balancing both explored and unexplored actions, and different action types (Appendix B). All these factors make GIF-MCTS specifically suitable for generating world models. Next we present the three action types (generate new lines, improve predictions and fix bugs) used in GIF-MCTS. We point the reader to the Appendix for the full action prompts, the remaining implementation details, and for the ablation study on the importance of the three action types.
4.1 GIF-MCTS Actions
Generate new lines. The goal of the generate action is to leverage the stochastic sampling ability of the LLM by generating varying continuations for 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 Code World Model, this can be a wrongly predicted transition, with the input state and action taken by the agent, the ground -ruth next state, and the model's 's predicted next state. The improve prompt also asks the LLM to produce a Chain-of- 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 biock-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 syntax runtime error, and will thus receive a reward of O strongly discouraging further exploration of the node, This can be wasteful, as sometimes the newly generated program can have sound logic and would receive 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 fixed version of the code, aided by Chain-of- Thought reasoning structure. To ensure that buggy nodes are chosen by the UCT formula, assign them with temporary value until either the bug is fixed more attempts are allowed (see Appendix B for additional details).
5 Experiments
In this Section, we first describe the baseline code generation methods we compare against and then present empirical results the APPS benchmark, the proposed CWMB and perform an additional

---

study on the RTFM environment. Additional ablations and qualitative results on GIF-MCTS are presented in Appendices C and D.
5.1 Baselines
The first baseline, denoted as Zero-shot CoTT' and used only for the experiments on APPS, adapts the work by Kojima 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 al. [2024b] to make as fair comparison possible. The WorldCoder algorithm calls the LLM with our generate prompt to produce an initial program, then for each remaining iteration we select one of the previous programs as explained below, 2) refine 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 p associated with Beta distribution B(a, (3) with initial parameters
1 C * r(p) and 3 = 1 C(1 -r())) which are updated every time the program selected. Here r(p) stands for the fraction of unit tests passed (same metric used in the evaluation phase of GIF-MCTS) and C is a constant set to 5, as in the original work. To select the next program to be refined, one sample is drawn from each Beta distribution and the program with the highest score is selected. In all experiments, we use the same amount of calls of GIF-MCTS.
5.2 APPS
We assess the overall performance of GIF-MCTS for generic code synthesis in the presence of public unit tests on the APPS benchmark [Hendrycks 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 set comprised of 1000 problems, as most closely reflects the challenges found in synthesizing CWMs: the problems tend to have a longer description, follow specific format for the input and output, and include challenging logic. Early experiments HumanEval [Chen et al., 2021], another popular coding benchmark, did not show clear correlation between a model's performance on the benchmark and its ability to generate CWMs, HumanEval problems typically easier and solvable with much shorter code snippets.
As GIF-MCTS requires 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 provided 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 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 L.L.M calls used during the synthesis process. We leave extending the approach for pass I cligibility using self- geff-genertted unit tests [Chen et al., 2023] for future work. We report the strict accuracy rate (the fraction problems on which all test cases are solved) on APPS for GIF-MCTS and other baselines in Table I.
Table I: APPS competition results: comparison of methods. We report the percentage of problems with all unit tests passed (Strict Accuracy). For our experiments, we also include the error of the mean on the percentage.
Method
Model
Size
Strict Accuracy (%) Evaluation Strategy
CodeRL [Le et al., 2022]
CodeTS
770M
17.90
pass 1000
Parsel [Zelikman et al., 2023]
code-davinci-002
N/A
25.50
pass @ any
Zero-shot CoT [Kojima et al., 2022]
Ilama 3
70B
23.2+
pass @ 20
WorldCoder [Tang ct al., 2024b]
Llama 3
70B
25.1+:1.4
pass @
CUF- (ours)
Llama 3
70B
28.311.4
pass @ 20
Our re-implementationn.
7

---

Results. GIF-MCTS outperforms strong previous baselines on the APPS competition split, reaching a new state of the the best of our knowledge. While part of this can be due advances in the underlying model, the comparisons with Zero-shot CoTT and WorldCoder show improved performance over either prior method. GIF-MCTS also markedly more sample efficient compared to established baselines; Parsel achieves the second best accuracy. but evaluates an exponentially growing number of solutions", while GIF-MCTS outperforms it by evaluating only 20 different programs.
5.3 Code World Models Benchmark
We evaluate our proposed GIF-MCTS approach and the WorldCoder baseline on the CWMB (intro- duced in Section 3). In this setting, we are interested both the accuracy of the generated CWM, as well as its performance when actually employed by 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:
R(CWM) = R(Tewm) R (Tadd)
R((mme) R((Trand)
where R(TCwM) represents the return obtained when using the CWM as the internal model for the planner, R( ((TTn)) ) is the return gathered with the true environment as the model while using the same planner (oracle planner), and R(m)ad) 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 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, 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.
Discrete Action Space
Continuous Action Space
Model
Method
Budget
Accuracy (1)
R(t)
Accuracy (1)
R(t))
GIF-MCTS (ours)
50
0.84+00.3
0.76++0.03
0,35-0.08
0.22++0.01
Llama 3 70B (3 seeds)
WorldCoder
50
0.79=+0.04
0.60=#0.04
0.32=0.03
0.19:10.01
GIF-MCTS (ours)
10
0.91:10.08
0.81
0.40
0.2611.01
GPT-4 Turbo (1 seed)
WorldCoder
10
0.87+00.99
0.79-+0.06
0.24-0006
0.20++000
* Our re-implementation of [TTan 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 fully executable CWM (with either method) for the two hardest environments, Humanoid-ve v4 and HumanoidStandup-we y4, due to their complexity and large observation space, while GPT-4 successfully generated model for each environment in the benchmark.
5.4 Read to Fight Monsters
We perform 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
Ressult reported for Parsel use 8 pseudo-code per problem, each implementing n sub- functions (with n being problem- dependenn))) 16 times and then evaluating up to 8 * 16" sub-functions combinations against APPS unit tests and keeping the best result.
8

---

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 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 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 '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 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 't do this on the full CWMB due to budget concerns.
Table 3: RTFM results. For each method and computational budget (LLM calls), we report the CWM accuracy and the normalized return R (computed across 10 episodes), with their errors.
Model
Method
Budget
Accuracy (t)
R(t)
GIF-MCTS (ours)
50
0.58 + 0.02
-0.11 + 0.12
Llama 3 70B
WorldCoder *
50
0.23 + 0.01
-0.11 + 0.12
GIF-MCTS (ours)
10
0.71 + 0.01
0.31 0.19
GPT-4 Turbo
WorldCoder
10
0.33 + 0.01
0.22 + 0.18
GIF-MCTS (ours)
50
1.00 + 0.00
1.00 + 0.00
Turbo
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 lower budget. As this version of the RTFM environment has never been published, this experiment can also alleviate concerns that the final CWM memorized by the LLM during pre- pre-training. We present and discuss further evidence against the significance of data contamination in Appendix E.
6 Discussion
In this section, we first discuss some takeaways from the empirical results and then elaborate on some of the limitations for our method.
GIF-MCTS vs. WorldCoder. We believe that GIF-MCTS outperforms WorldCoder because produces a more diverse set of programs. WorldCoder initially generates single program from scratch and then samples and refines a complete program 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.
Recuracy-Retura 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- approaches. As the gathered transitions only used 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 future work.
Comparison with Offline RI. We expect CWMS to hold advantages over classical RI. methods in regimes with scarce data and environments that can be easily described by language and modeled with code. We report in Appendix K preliminary comparison on the CWMB of the return achieved with our CWMs or with a SOTA offline RL method, Conservative Q-Learning (CQL) [Kumar et al., 2020], trained on the same amount of trajectories used for synthesizing the CWMs. We find that CWMs compare favourably against CQL on environments with discrete action spaces, while CQL's performance is superior on the continuous action space environments, which are harder to model. RL methods, including CQL, would likely benefit from more experience, as they overfit with scarce data.
6.1 Limitations
Code World Models. The CWMs framework is an exciting direction for model-based planning, but we still rely on limiting assumptions of deterministic and fully observable environments. Both stochasticity and partial observability would pose challenges, especially on the verification of the CWM prediction, as there is no set result for a given input. We leave extending the approach to account for both stochastic and partially observable environments to future work.
Another potential issue is providing a description of the environment that can be reasonably converted to a Python function (e.g. manual documenting key variables) when such description is not available (e.g. when the environment is defined with image observations). Previous work has begun 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- re-written individually by LLM, to account for some changes in the environment, or modeling variable factors as arguments to the step function. CWMs struggle especially complex physics-based environments; thus promising direction could also be allowing programs generated by GIF-MCTS to make use of external tools and libraries, such as physics simulators.
GIF-MCTS. We have validated the GIF-MCTS approach as an efficient code synthesis method, with the key limiting assumption of having available test cases evaluate code, which could be difficult to provide in certain tasks. In those cases, it would be possible to use self-generated test cases [Chen et al., 2023], but since this does not reflect the CWM setting we leave this for future work.
7 Conclusion
We present Code World Models, a general framework to leverage LLMs to build world models for RL agents. We further show that GIF-MCTS is 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 L.L 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.
10

---

Acknowledgments and Disclosure of Funding
This work was supported by the Research Council of Finland (Flagship programme: Finnish Center for Artificial Intelligence FCAI, and grants 352986, 358246) and EU (H2020 grant 101016775 and NextGenerationEU). We acknowledge CSC for awarding this project access to the LUMI supercomputer, owned by the EuroHPC Joint Undertaking, hosted by CSC (Finland) and the LUMI consortium through Finland.
References
Michael Ahn, Anthony Brohan, Noah Brown, Yevgen Chebotar, Omar Cortes, Byron David, Chelsea Finn, Keerthana Gopalakrishnan, Karol Hausman, Alexander Herzog, et al. Do as can, not as i say: Grounding language in robotic affordances. Conference on Robot Learning, 2022.
Abdus Salam Azad, Edward Kim, Qiancheng Wu, Kimin Lee, Ion Stoica, Pieter Abbeel, Alberto Sangiovanni-Vincentee and Sanjit A Seshia. Programmatic modeling and generation of real-time strategic soccer environments for reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 6028-6053 2022.
Osbert Bastani, Yewen Pu, and Armando Solar-Lezama, Verifiable reinforcement learning via policy extraction. In Advances in Neural Information Processing Systems, volume 31, 2018.
Mare G. Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environ- ment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47: 253-279, 2013.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, et al. Language models are few-shot learners. In Advances in Neural Information Processing Systems, volume 33, pages 1877-1901, 2020.
Jake Bruce, Michael D Dennis, Ashley Edwards, Jack Parker-Holder, Yuge Shi, Edward Hughes, Matthew Lai, Aditi Mavalankar, Richie Steigerwald, Chris Apps, Yusuf Aytar, et al. Genie: Generative interactive environments. In Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 4603-4623. PMLR, 21-27 Jul 2024.
Tales Henrique Carvalho, Kenneth Tjhia, and Levi Lelis. Reclaiming the source of programmatic policies: Programmatic versus latent spaces. In The Twelfth International Conference on Learning Representations, 2024.
Bei Chen, Fengji Zhang, Anh Nguyen, Daoguang Zan, Zeqi Lin, Jian-Guang Lou, and Weizhu Chen. Codet: Code generation with generated tests. In The Eleventh International Conference on Learning Representations, 2023.
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde De Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. Evaluating large language models trained on code. arXiv preprint arXiv:2107.0333 2021.
Nicola Dainese, Pekka Marttinen, and Alexander Ilin. Reader: Model-based language-instructed reinforcement learning. In Proceedings of the 2023 Conference on Empirical Methods in Nat- ural Language Processing, pages 16583-16599, Singapore, December 2023. Association for Computational Linguistics.
Ria Das, Joshua B. Tenenbaum, Armando Solar-Lezama, and Zenna Tavares. Autumnsynth: Synthesis of reactive programs with structured latent state. In Advances in Programming Languages and Neurosymbolic Systems Workshop, 2021.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint vrXiv:2004.072111 2020.
Malik Ghallab, Adele Howe, Craig Knoblock, Drew McDermott, Ashwin Ram, Manuela Veloso, Daniel Weld, and David Wilkins. PDDL, The Planning Domain Definition Language, 1998.
11

---

Zhibin Gou, Zhihong Shao, Yeyun Gong, yelong shen, Yujuu Yang, Nan Duan, and Weizhu Chen. CRITIC: Large language models can self-correct with tool coll-interactive critiquing, In The Twelfth International Conference Learning Representations, 2024.
Lin Guan, Karthik Valmeekam, Sarath Sreedharan, and Subbarao Kambhampati. Leveraging pre- trained large language models to construct and utilize world models for model-based task planning. In Advances in Neural Information Processing Systems, volume 36, pages 79081-79099 2023.
Agrim Gupta, Lijun Yu, Kihyuk Sohn, Xiuye Gu, Meera Hahn, Li Fei-Fei, Irfan Essa, Lu Jiang, and José Lezama. Photorealistic video generation with diffusion models. arXiv preprint arXiv:2312.0666 12.6662. 2023.
David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In Advances in Neural Information Processing Systems, volume 31, 2018.
Danijar Hafner, Timothy P LLillicap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In International Conference on Learning Representations, 2021.
Shibo Hao, Yi Gu, Haodi Ma, Joshua Jiahua Hong, Zhen Wang, Daisy Zhe Wang. and Zhiting Hu. Reasoning with language model is planning with world model. In The 2023 Conference on Empirical Methods in Natural Language Processing. 2023.
Dan Hendryeks, Steven Basart, Saurav Kadavath, Mantas Mazeika, Akul Arora, Ethan Guo, Collin Burns, Samir Puranik, Horace He, Dawn Song, and Jacob Steinhardt. Measuring coding challenge competence with apps. In Proceedings of the Neural Information Processing Systems Track on Datasets and Benchmurks, volume 1, 2021.
Jeevana Priya Inala, Osbert Bastani, Zenna Tavares, and Anmando Solar-Lezama. Synthesizing programmatic policies that inductively generalize. In 8th International Conference on Learning Representations, 2020.
Eric Jang, Alex Irpan, Mohi Khansari, Daniel Kappler, Frederik Ebert, Corey Lynch, Sergey Levine, and Chelsea Finn. Bc-z: Zero-shot task generalization with robotic imitation learning. In Conference on Robot Learning, pages 991-1002. PMLR, 2022.
Siddharth Karamcheti, Suraj Nair, Annie S. Chen, Thomas Kollar, Chelsea Finn, Dorsa Sadigh, and Percy Liang. Language-driven representation learning for robotics. In Robotics: Science and Systems (RSS), 2023.
Michael N. Katehakis and Arthur Veinott. The multi-armed bandit problem: Decomposition and computation. Mathematics of Operations Research, 12(2):262-26)) 268. 1987.
Levente Koesis and Csaba Szepesvári. Bandit based monte-carlo planning. In Proceedings of the 17th European Conference on Machine Learning, ECML.'06, page 282-293, Berlin, Heidelberg, 2006. Springer-Verlag.
Takeshi Kojima, Shixiang (Shane) Gu, Machel Reid, Yutaka Matsuo, and Yusuke Iwasawa. Large language models are zero-shot reasoners. In Advances in Neural Information Processing Systems, volume 35, pages 22199-2211 22213, 2022.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. In Advances in Neural Information Processing Systems, volume 33, pages 1179-1191, 2020.
Hung Le, Yue Wang, Akhilesh Deepak Gotmare, Silvio Savarese, and Steven Chu Hong Hoi Coderl: Mastering code generation through pretrained models and deep reinforcement learning. In Advances in Neural Information Processing Systems, volume 35, pages 21314-21328, 2022.
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.0 2020.
Jacky Liang, Wenlong Huang, Fei Xia, Peng Xu, Karol Hausman, Brian Ichter, Pete Florence, and Andy Zeng. Code as policies: Language model programs for embodied control. 2023 IEEE International Conference Robotics and Automation (ICRA), pages 9493-9500. IEEE, 2023.
12

---

Jessy Lin, Yuqing Du, Olivia Watkins, Danijar Hafner, Pieter Abbeel, Dan Klein, and Anca Dragan. Learning to model the world with language. In Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 29992- 30017. PMIR, 21-27 Jul 2024.
Guan-Ting Liu, En-Pei Hu, Pu-Jen Cheng, Hung- Yi Lee, and Shao-Hua Sun. Hierarchical pro- grammatic reinforcement learning via learning to compose programs. Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 21672-21697. PMLR, 23-29 Jul 2023.
Hao Liu, Wilson Yan, Matei Zaharia, and Pieter Abbeel. World model on million-length video and language with ringattention. arXiv preprint arXiv:2402.08268 2024.
Aman Madaan, Niket Tandon, Prakhar Gupta, Skyler Hallinan, Luyu Gao, Sarah Wiegreffe, Uri Alon, Nouha Dziri, Shrimai Prabhumoye, Yiming Yang. Shashank Gupta, Bodhisattwa Prasad Majumder, Katherine Hermann, Sean Welleck, Amir Yazdanbakhsh, and Peter Clark. Self-refine: Iterative refinement with self-feedback. In Advances in Neural Information Processing Systems, volume 36, pages 46534-46594, 2023.
Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample-efficient world models. In The Eleventh International Conference on Learning Representations, 2023.
Toki Migimatsu and Jeannette Bohg. Grounding predicates through actions. In 2022 IEEE Interna- tional Conference on Robotics and Automation (ICRA), pages 3498-3504, IEEE, 2022.
Theo X Olausson, Jeevana Priya Inala, Chenglong Wang, Jianfeng Gao, and Armando Solar-Lezama. Is self-repair silver bullet for code generation? In The Twelfth International Conference Learning Representations, 2023.
Wenjie Qiu and He Zhu. Programmatic reinforcement learning without oracles. In The Tenth International Conference on Learning Representations, 2022.
Tianhe Ren, Shilong Liu, Ailing Zeng, Jing Lin, Kunchang Li, He Cao, Jiayu Chen, Xinyu Huang, Yukang Chen, Feng Yan, et al. Grounded sam: Assembling open-world models for diverse visual 2024.
tasks. arXiv preprint arXiv:2401.144
Reuven Rubinstein. Optimization of computer simulation models with rare events. European Journal of Operational Research, 99(1):99-112, 1997.
Noah Shinn, Federico Cassano, Ashwin Gopinath, Karthik Narasimhan, and Shunyu Yao. Reflexion: language agents with verbal reinforcement learning. In Advances in Neural Information Processing Systems, volume 36, pages 8634-8652, 2023.
Tom Silver, Kelsey R Allen, Alex K Lew, Leslie Pack Kaelbling, and Josh Tenenbaum. Few-shot bayesian imitation learning with logical program policies. Proceedings of the AAAI Conference on Artificial Inielligence, volume 34, pages 10251-10258. 2020.
Ishika Singh, Valts Blukis, Arsalan Mousavian, Ankit Goyal, Danfei Xu, Jonathan Tremblay, Dieter Fox, Jesse Thomason, and Animesh Garg. Progprompt: Generating situated robot task plans using large language models. In 2023 IEEE International Conference on Robotics and Automation (ICRA), pages 11523-11530. IEEE, 2023.
Richard S. Sutton. Dyna, an integrated architecture for learning, planning, and reacting. SIGART Bulletin, 2(4): July 1991.
Hao Tang, Keya Hu, Jin Peng Zhou, Sicheng Zhong, Wei-Long Zheng, Xujie Si, and Kevin Ellis. Code repair with Ims gives an exploration-exploiittion tradeoff. arXiv preprint arXiv:2405. 17503. 2024a.
Hao Tang, Darren Key, and Kevin Ellis. Worldcoder, a model-based Ilm agent: Building world models by writing code and interacting with the environment. arXiv preprint arXiv:2402. 12275v1, 2024b.
13

---

William R Thompson. On the likelihood that one unknown probability exceeds another in view of the evidence of two samples. Biometrika, 25(3-4):285-299) 1933.
Mark Towers, Ariel Kwiatkowski, Jordan Terry, John U Balis, Gianluca De Cola, Tristan Deleu, Manuel Goulao, Andreas Kallinteris, Markus Krimmel, Arjun KG, et al. Gymnasium: standard interface for reinforcement learning environments. arXiv preprint arXiv:240.117032 2024.
Dweep Trivedi, Jesse Zhang, Shao-Hua Sun, and Joseph J Lim. Learning to synthesize programs as interpretable and generalizable policies. In Advances in Neural Information Processing Systems, volume 34, pages 45114625111 2021.
Pedro A Tsividis, Joao Loula, Jake Burga, Nathan Foss, Andres Campero, Thomas Pouncy, Samuel J Gershman, and Joshua Tenenbaum. Human-level reinforcement learning through theory-based modeling, exploration, and planning. arXiv preprint arXiv:2107. 12544, 2021.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, volume 30, 2017.
Abhinav Verma, Vijayaraghavan Murali, Rishabh Singh, Pushmeet Kohli, and Swarat Chaudhuri. Programmatically interpretable reinforcement learning. In International Conference on Machine Learning, pages
PMLR, 2018.
Abhinav Verma, Hoang Le, Yisong Yue, and Swarat Chaudhuri. Imitation projjcten programmatic reinforcement learning. In Advances in Neural Information Processing Systems, volume 32, 2019.
Guanzhi Wang. Yuqi Xie, Yunfan Jiang, Ajay Mandlekar, Chaowei Xiao, Yuke Zhu, Linxi Fan, and Anima Anandkumar. Voyager: An open-ended embodied agent with large language models. arXiv preprint arXiv:2305. /6291, 2023.
Jason Wei, Xuezhi Wang. Dale Schuurmans, Maarten Bosma, brian ichter, Fei Xia, Ed H. Chi, Quoc V Le, and Denny Zhou. Chain of thought prompting elicits reasoning in large language models. In Advances in Neural Information Processing Systems, 2022.
Lionel Wong, Gabriel Grand, Alexander K Lew, Noah D Goodman, Vikash K Mansinghka, Jacob Andreas, and Joshua B Tenenbaum. From word models to world models: Translating from natural language to the probabilistic language of thought. arXiv preprint arXiv: 2eee222366.2227 12672, 2023.
Lionel Wong, Jiayuan Mao, Pratyusha Sharma, Zachary S Siegel, Jiahai Feng, Noa Korneev, Joshua B. Tenenbaum, and Jacob Andreas. Learning grounded action abstractions from language. In The Twelfth International Conference on Learning Representations, 2024.
Sherry Yang, Yilun Du, Seyed Kamyar Seyed Ghasemipour, Jonathan Tompson, Leslic Pack Kael- bling, Dale Schuurmans, and Pieter Abbeel. Learning interactive real-world simulators. In The Twelfth International Conference on Learning Representations, 2024.
Shunyu Yao, Dian Yu, Jeffrey Zhao, Izhak Shafran, Tom Griffiths, Yuan Cao, and Karthik Narasimhan. Tree of thoughts: Deliberate problem solving with large language models. In Advances in Neural Information Processing Systems, volume 36, pages 11809-1 11822, 2023.
Eric Zelikman, Qian Huang, Gabriel Poesia, Noah Goodman, and Nick Haber. Parsel: Algorithmic reasoning with language models by composing decompositions. In Advances in Neural Information Processing Systems, volume 36, pages 31 1444666-1523, 2023.
Alex Zhang, Khanh Nguyen, Jens Tuyls, Albert Lin, and Karthik Narasimhan. Language-guided world models: A model-based approach to ai control. arXiv preprint arXiv:2402.0 01695, 2024.
Shun Zhang, Zhenfang Chen, Yikang Shen, Mingyu Ding, Joshua B. Tenenbaum, and Chuang Gan. Planning with large language models for code generation. In The Eleventh International Conference on Learning Representations, 2023.
Victor Zhong, Tim Rocktäschel, and Edward Grefenstette. RTFM: Generalising to new environment dynamics via reading. International Conference on Learning Representations, 2020.
Andy Zhou, Kai Yan, Michal Shlapentokh-Rothmaal Haohan Wang, and Yu-Xiong Wang. Language agent tree search unifies reasoning acting and planning in language models. arXiv preprint arXiv:2310.04406 2023.
14

---

A Broader Impact
The CWM framework enables LLMs generate world models for model-based Reinforcement Learning, which could potentially be employed for planning with a real agent, As the code generated by the LLM is untrusted, it should always be checked by a human expert before it is used under any circumstances. Alternatively, as CWMs are represented with Python code, this also allows for interpretable world models, which could be safer for critical applications after being vetted by an expert.
B Additional GIF-MCTS implementation details
Choice of Actions If a node doesn't 't contain bug, new generate and improve actions should always be available (with the exception of the root node, which will only have new generale 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 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 [Koesis and Szepesvári, 2006] as follows:
In N₂ Raze, +
(no(e) U/ C.
where o; the value of the node, C is a constant parameter used to balance exploration (empirically set to 0. ). N; is the number of visits to the node's parent and Ra=a; 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 Ra=o, 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, employ 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 separate classifier, and further differentiates local and global values. We define the global value eig as the average of all values of the nodes with the same action type any level of the tree and the local value UL as the average of all expanded children with the same action type. The linear model then simply learns to predict the value U; of a given action as balanced sum of the two values, normalized between zero and one, with the following formula:
WG UG WL 7L
02
WG + wt
where the we and wt. parameters are learned during the search using gradient descent.
Initially, the global average ao 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 weight the prior when computing the average (effectively acting as if there were nodes already discovered with the prior value).
Value Estimation for Buggy Nodes. As mentioned in Sec. buggy nodes will get a reward of and would thus never be explored. To allow the fix action to be chosen, we assign 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 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 ʃ unsuccessful fixes, the temporary value is set to zero, which strongly discourages the buggy parent node from
15

---

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. also important to note that the temporary values are excluded from the backtracking step of the MCTS algorithm, 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 Huggingface documentation" for an accurate description of each LLM parameter.
Table 4: GIF-MCTS hyperparameters.
Parameter
Description
Value
L
Number of new lines extracted from a generate action.
2
c
Visit count offset.
1.0
C
Exploration constant.
0.1
7
Discount factor.
1.0
Ug
Initial prior for generate actions (with its initial count).
0.5 (2)
U
Initial prior for improve actions (with initial count). 0.55 0.55 Number of allowed fixes to node.
3
Table 5: Llama hyperparameters. Note that for GPT-4 Turbo, the only parameter used was the number of maximum new tokens, to the same value used for Llama.
Parameter
Value
max max_new_token.
1500
temperature
1.0
top_k
100
top_P
0.8
gum_retuurn_seuencee
num jemms
1
C Ablation Study on GIF-MCTS
We perform an ablation study to validate the individual contribution of each action type of GIF-MCTS. We run the MCTS procedure on CWMB with only two out of the three actions available and compare the accuracy with the full method in Table 6. Note that for the Fix and Improve MCTS variant, one generate action applied at the root node to obtain an initial program, which the algorithm expands from with the available budget. All ablations 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.
The performance the method drops after the removal each action, most significantly in the harder of continuous environments (while there more statistical uncertainty for the discrete environments). Fixing bugs appears to be the most important action: it is much more efficient to try fixing 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 generate fully functioning program from the start. On the other hand, improve seems be the least impactful: this makes sense, as intuitively improving code snippet that already works is has less room for improvement.
Results.
D Qualitative Study
To investigate the specific effectiveness of each individual type of action, analyze the trees produced by GIF-MCTS and report some statistics interest in Table 7. We specifically focus on the
16

---

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 emor.
Discrete Action Space
Continuous Action Space
Method
Budget
Accuracy (t)
RR(1)
Accuracy (t)
R(t)
GIF-MCTS (ours)
50
0.08
0.83=0.06
0.38
0.23
No Generate action
50
0.87
0.73=0.09
0.25=0.06
0.16
No Improve action
0.05
0.79=0.07
0.34=0.05
0.17#0.02
No Fix action
50
0.81
0.55
0.21
0.1010.01
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 find specific biases towards 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).
Discrete Action Space
Continuous Action Space
Quantity
Llama 3 70B
GPT-4 Turbo
Llama 3 70B
GPT-4 Turbo
% generates
50.0
88.3
18.5
33,4
% improves
44.7
8.3
35.3
34.8
% fixes
5.3
3.4
46.2
31.8
Path length
5.7
2.3
3.2
2.3
% path generates
73.2
100.0
47.0
59.0
% path improves
17.5
0.0
5.0
6.3
% path fixes
9.3
0.0
48.0
34.7
Tree depth
15.6
5.0
10.8
4.5
From the results, the method presents pretty clear bias towards the generate action 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 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 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 sample that is actually increasing the performance, while generate actions have the advantage of being chosen at the beginning, where it is possibly casier to find good programs.
Still, the fact that many improve actions are taken that result in either the same value 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 explaining code which provides feedback to another specialized in generating it).
17

---

There is also clear difference between the set of casier discrete action space problems, for which the percentage of fix actions very low (with GPT-4 Turbo only needing generates in order to synthesize perfect or near-perfect models, shown in Table 11) and the harder continuous action space problems, where fixing bugs becomes much more prominent.
E Data Contamination
With any experiment involving LLMs there is a concern about data contamination: the model's pre-training corpus could have included the original implementation for the various programs we 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 (ArCoder and Codeforces, which include 264 and 41 problems respectively). This means that for the third website, Kattis, which makes up 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 observe that some parts of the generated CWMB environments recall implementa- tions 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, the code is necessarily different. Examples of generated CWMs along with their ground-truth implementations can be found in Appendix O for more thorough comparison.
As we use modified version of the RTFM environment (with fixed manuals and no stochasticity), there is no reference solution for online, which provides evidence that our solution is not merely retrieving information from the LLM's 's training data.
Generally speaking, there is of course way to outright dismiss these concerns. However, method compared to baselines using the same underlying models, ensuring that the superior performance reported for GIF-MCTS is not biased by potential data contamination.
F
Data Quality
As part of the CWMB, for each environment the collected dataset 'D contains both low-scoring and high-scoring trajectories. As discussed in Section 3, this is fairly standard practice for offline RL, as the general assumption that 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 fair model even in hard environments when provided with only 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. other words, this a worse version of the same buffer used for the main experiment, which by construction carries less information. We synthesized CWM with GIF-MCTS and GPT-4 using 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
18

---

contains both positive and negative rewards), the model still scores 100% accuracy, on par with the model generated with the full range of data.
G Additional Related Work
We expand in the following the Related Work section, covering the works that try to build world models with language and those who explored using programs to express RL policies.
World Models with Language. Model-based RL methods are built around learning a predictive model of the environment to inform the agent's decisions [Sutton, 1991]. A recently growing body of research is focusing on building world models that can include information in natural language, as opposed to approaches using only vision or full state observations [Hafner et al., 2021]. Dynalang [Lin et al., 2024] predicts the future text and image representation of the environment with an encoder- decoder architecture with joint input of previous frames and text, while Zhang 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 meoder-decoder model for language-drivee 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 RI.. Verma al. [2018. 2019] first introduced Programmatically Interpretable RL (PIRL), which focuses on representing RL policies as interpretable and verifiable programs by first learning an oracle policy with deep RL and then distilling a program with a domain specific language that can model tree-like programs. Similarly, Bastani et al. [2018] focus on extracting decision trees from an oracle policy with imitation learning and Inala al. [2020] use itate 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 RI. 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 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 robotie policy, with similar approach concurrently introduced by Singh et al. [2023]. Voyager [Wang et al., 2023] generates an incrementally growing skill library using JavaScript code play Minecraft.
H Comparison of Inference Times
We further demonstrate the efficiency of CWMs compared to directly using an LLM as the world model in Table 8. On selection of three environments from the CWMB we ask GPT-4 Turbo to directly predict the next observation of the environment given its description and some in-context examples of the task, and compare the inference time with calling the step function of the CWM. Calling the Python program is four orders of magnitude quicker for the easiest environment and seven orders of magnitude quicker for the hardest environment. We additionally observe that none of the predictions made by GPT-4 Turbo were accurate.
I Code World Models Benchmark Details
We include a detailed list of statistics for each environment in the CWMB in Table 9. Notice that when creating the descriptions from the Gymnasium docstrings, we left out documentation sections that do not relate to the environment definition itself, such as versioning information, Gymnasiun-rrelated arguments, and external references, from these descriptions. For the reported number of tokens
19

---

Table 8: Comparison: inference times between GPT-4 and CWM. Results are calculated from sample of 10 transitions from the replay buffer used during GIF-MCTS.
Environment
GPT-4 Time (s)
CWM Time (8)
CartPole-vl
2.2
0,00005
HalfCheetah-v4
6.1
0,0001
Humanoid-v4
146.7
0.0001
choose OpenAl's open source tiktoken tokenizerS. The code lines and code tokens are reported from the corresponding CWM generated by GPT-4 Turbo using GIF-MCTS with budget of 10. This is meant to be a general indication of how long a typical implementation of the environment would be, but can of course vary. All environment descriptions were parsed from Gymnasium v.0.29.1.
Table CWMB details. Detailed statistics for each environment in the CWMB. An Action Space or Observation Space indicated between bars (I.A), [8] n) indicate discrete space with different choices. The value intervals for each space are omitted for visual clarity.
Description Lines
Description Tokens
Action Space
Ohservation Space
Code Lines
Code Tokens
Environment
Dimensionality
Dimensionality
Blackjack-v1
601
A
(32, 11,2)
94
826
ClifWWalkig-vOO
47
456
A
S 48
611
483
Taxi-v3
89
724
A 6
S 500
83
767
Acrobot-v1
66
859
A
S R
76
794
CariPole-VI
53
663
A =
S
62
639
MountainCar-v0
47
454
A =
8 R²
62
Ant-v4
148
2983
AER"
S R
33
267
HaliCheetah-v4
86
1674
ACR
8 R
58
554
Hopper-v4
87
1529 AERI
8 R
91
847
Humanuid-v4
204
4578 AERIT
S R
68
617
HumanoidStandup-vv
202
4551 A AER"
$ E R75
50
442
InverteDDoubbeeeenbbuuss
84
1364 A ACR'
8 R 8
54
465
InvertedPendulum-
683
AFE'
66
Pendulum-v1
ACR
8 $ R
Pushcr-v4
2035
AER"
76
587
Reacher-v4
87
1472
AER
S R1
78
699
Swimmer-v4
68
1168
AER"
S R
80
700
Walker2d-v4
92
1785
A CCR
S R
81
770
Indicative number sampled from single result, can vary.
J Results for Individual Environments
We report the individual accuracy and return for each environment in the CWM when using I.lama in Table 10 and when using GPT-4 Turbo in Table 11.
K Comparison with Offline RL
We compare the overall performance of a SOTA offline RL method, Conservative Q-Learning (CQL) [Kumar et al., 2020], against a planning agent using the synthesized CWM with our method. We report in Table 12 the average raw reward obtained over 10 episodes for a random policy, CQL, planning agents with the CWM obtained by GIF-MCTS (ours) respectively with Llama and GPT-4, and 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 balance between CQL and CWMs, with CWMs being more suited discrete tasks and CQL outperforming CWMs in complex physics tasks, where our method struggles. However,
'nttps://yypi. wwtg/project/tikookenl
20

---

Table 10: CWMB results. Individual results for each environment in the CWMB using Llama 3 (we report the results for the first seed only).
GIF-MCTS
WorldCoder
Environment
Action Space
Accuracy (1)
R(t)
Accuracy (t)
R(t)
CartPole-vl
Discrete
1.00
1.11
0.92
1.09
CliffWalking-vo
Discrete
1.00
1.01
1.00
0.97
MountainCar-v0
Discrete
1.00
N/A
0.83
N/A
Taxi-v3
Discrete
0.92
0.67
0.44
0.23
Blackjack-vl
Discrete
0.83
0.53
0.85
0,41
Acrobot-v1
Discrete
0,54
N/A
0.73
N/A
InvertedPendulum-vá
Continuous
0.66
0.14
0.66
0.01
Pusher-v4
Continuous
0.41
0.74
0.41
0.77
Pendulum-vi
Continuous
0.34
0.15
0.31
-0.15
Walker2d-v4
Continuous
0.34
0.07
0.34
0.08
Hopper-v4
Continuous
0.33
0.15
0.00
0.02
Swimmer-v4
Continuous
0.33
0.01
0.33
0.07
HalfCheetah-v4
Continuous
0.33
0.13
0.33
0.15
Ant-v4
Continuous
0.33
0.67
0.33
0.69
InveredDDubliPennduuu-wt
Continuous
0,25
0.06
0.34
0.05
Reacher-y4
Continuous
0.13
0.93
0,42
0.67
HumanoidStandup-ve
Continuous
N/A
0.00
N/A
0.00
Humanoid v4
Continuous
N/A
0.00
N/A
0.00
Table 11: CWMB results. Individual results for each environment in the CWMB using GPT-4 Turbo.
GIF-MCTS
WorldCoder
Environment
Action Space
Accuracy (t)
R(t)
Accuracy (t)
R(t))
CartPole-vl
Discrete
1.00
0.99
1.00
1.00
CliffWalking-vo
Discrete
1.00
0.98
1.00
0.89
MountainCar-v0
Discrete
1.00
N/A
1.00
N/A
Taxi-v3
Discrete
0.99
0.87
0.99
0.67
Blackjack-vl
Discrete
0.93
0.41
0.79
0.59
Acrobot-v]
Discrete
0.53
N/A
0.42
N/A
InvertedPendulum-
Continuous
0,66
0.08
0.66
0,00
Humanoid-vá
Continuous
0.43
0.01
0.00
0.00
HumanoidStandup-ve
Continuous
0.42
+0.04
0.00
0.00
Reacher-v4
Continuous
0.42
0.88
0.42
0.71
Pusher-v4
Continuous
0.41
0.72
0.41
0.70
LnvereedDoublePendduum
Continuous
0.41
0.02
0.00
0.00
Pendulum-vl
Continuous
0.38
0.51
0.38
0.50
Walker2d-v4
Continuous
0.34
0.03
0.01
0.03
Hopper-v4
Continuous
0.34
-0.04
0.33
-0.01
Swimmer-y4
Continuous
0.33
0.04
0.33
0.02
HalfCheetah-v4
Continuous
0.33
0.23.
0.33
0.24
Ant-v4
Continuous 0.33
0.69
0.00
0.20
CWMs also reach competitive results some of these harder environments, such as Pendulum-vzs Reacher- v4 and to lesser extent Ant-v4, Pusher-v4 and HalfCheetah-uue even without direct access to the original physics simulator. Particularly these tasks, but also in general, we observe severe overfitting happening in CQL almost immediately (for example, CQL performs worse than random in Pendulum-ve um-11, likely due to the small size of the provided dataset. mentioned previously, sample efficiency is one of the main promises of the CWM approach, as very few trajectories are needed validate the model, whereas traditional methods are typically designed to work best with large amounts of data.
21

---

Table 12: Comparison with CQL. We report the average raw reward obtained over 10 episodes for random policy, Conservative Q-Learning (CQL), planning agents with the CWM obtained by GIF-MCTS (ours) respectively with Llama 3 and GPT-4, and planning agent with oracle access to the true environment (Oracle). CQL was trained with 10 epochs for 100 steps per epoch (1000 total steps) using the same dataset used learn our CWMs.
GIF-MCTS (ours)
Environment
Random
CQL
Oracle
Llama 3
GPT-4
Blackjack-vl
0
-0.3
-0.6
-0.1
1
CliffWalking-v0
-1169.2
N/A* *
-90.2
-100
-100
Taxi-v3
-798.5
-740
-353.9
408.8
-124.5
CartPole-vl
24.4
317.6
277.4
310.4
494
MountainCar-vO
-200
-200
-200
-200
-200
Acrobot-v1
-500
-295
-500
-494.2
5500
Pendulum-v1
-1122.
1218.2
-1322.2
-739.8
-373.6
Reacher-v4
-43.7
-1..5
-9.2
-11.2
-6.8
Pusher-v4
-149.9
-2.44
-61.1
-63.3
-30.3
InvertedPenduluu--v
8.3
66.7
13.1
10.9
42.5
InvertedDDoublePeeduuluu-
49
164
60
53.4
241.6
HalfCheetah-v4
-304.5
-1.3
-50.33
-22.8
893.3
Hopper-v4
32.2
137.4
62.6
23.3
229.1
Swimmer-v4
-5.9
28.4
-2.7
8.1
317.8
Walker2d-v4
0
278
22.3
11.5
334.7
Ant-v4
-33.2
998
867.7
896.8
1304.7
Humanoid-v4
139.4
393.3
N/A
162.3
1860.7
HumanoidStandupp-v.
33240.2
51045.7
N/A
29405.9
138075.6
N/A for CQL indicates a failed run, while for GIF-MCTS it indicates a failure in synthesizing a syntactically correct CWM.
It is also worth noting that outperforming state-of-the-art methods for offline RL 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, 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, 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 call a physics simulator in the continuous environments) future work, now that the effectiveness of the method has been proven.
L Planning algorithms details
In this section we report all the parameters used in our implementations of Monte Carlo Tree Search (MCTS) [Kocsis and Szepesvári, 2006] and Cross Entropy Method (CEM) [Rubinstein, 1997]. together with brief explanation of the meaning of those parameters within the context of the two algorithms.
MCTS. At each time-step, run Tmats simulations with MCTS select the best action play. At every simulation, starting from the root node, we select one action via the Upper-Confidenc Bound formula for Trees (UCT)
In N n4 +
UCT(node;)
C.
(4)
where is the estimated value of node i, C the exploration constant, N, is the visit count of the parent of node i, n, is the visit count of node and a factor offsetting the visit count. Once select unexplored action at of the nodes, we expand the node that the action leads to
22

---

and perform a rollout with random policy for up to max actions to estimate its value. The value backpropagation is done in standard MCTS and we use discount factor of y. The values of all parameters are reported in Table 13.
Table 13: MCTS planner parameters.
Parameter
Description
Value
mets
Number of iterations.
25
max actions
Max actions per rollout.
100
C
Exploration constant.
1.0
Visit count offset. Discount factor. Softmax temperature. 0.01
0.99
mets
CEM. In this case, assuming deterministic environments, we plan directly for the next Teem time- steps, meaning that we choose the actions for up to T Teem steps ahead, using the CEM algorithm. At every iteration we sample Neem action plans from zero-mean Gaussian with dimensions T Teem. x 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 rollouts in the environment, starting from the current state. We then select the top Keem action plans (clites samples), fit the Gaussian parameters to them and repeat. At the last iteration, return the top scoring action plan. All parameters are reported in Table 14.
Table 14: CEM planner parameters.
Parameter
Description
Value
T Teem
Time horizon.
100
Ieem
Number of iterations.
20
Norm
Number of samples.
1000
Keem
Number of clites.
100
M Computational Resources
In the following section report as accurately 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 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 problems for which we ran our method, one baseline and 3 ablations, which should be roughly equivalent to single experiment with 100 APPS problems, or 10 jobs of 15 hours with GPUs, for 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 similar amount of CPU hours.
Furthermore, have paid approximately $62.3 in OpenAI calls GPT-3.5 Turbo (used only for prototyping) and GPT-4 Turbo (used with a budget of 10 calls on the CWMB experiments 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 single consumer CPU in a few hours.
23

---

N
Prompts
In this section we report the main prompts used for GIF-MCTS. These prompts are also shared by WorldCoder implementation, while 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 experienced Python developer. You will be provided with incomplete code snippet from a Python program. The task this program supposed to perform described 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 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 the standard output (stdout). Output your code solution with the following format: "'python [your code] "
</user>
<assistant>
"python
(CODE_SO_FAL SO_FAR)
</assistant>
Figure 3: Prompt the APPS benchmark for the generate action.
24

---

<system>
You are an experienced Python developer. You will be provided with an incorrect code snippet from a Python program. The task this program supposed to perform is described 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 the standard
output (stdout).
First, write an explanation of the difference between the ground-truth output and the program's ss 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 concrete, actionable fix for it. Finally fix the program in its entirety following the suggestion. The expected output is in the format:
## Error explanation
[your explanation of the error]
## Error location and wrong logic
[where the error comes from and why]
## Fix suggestion [how fix the error]
## Correct code
"python [your code]
## Incorrect code
You are provided with the following code snippet fix. "'yython (CODE)
The code additionally makes wrong prediction about this input.
## Input
(NPUT)
## Ground-truth output
(OUTPUT)
## Code incorrect outputs
[PREDICTION]
</user>
<assistant>
## Error explanation
</assistant>
Figure 4: Prompt on the APPS benchmark for the improve action.
25

---

<system>
You are an experienced Python developer. You will be provided with an incorrect Python program. The task this program 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_DE DESCRRITIIONN
Please read the inputs from the standard input (stdin) and print the outputs the standard output (stdout).
First, write an explanation the error and point out the part 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 its entirety following the suggestion. The expected output is in the format:
### Error explanation
[your explanation of the error]
## Fix suggestion
[how fix the error]
## Correct code ppython
[your code]
## Incorrect code
You are provided with the following code snippet fix.
"'yython
(CODE)
(ERROR)
</user>
<assistant>
## Error explanation
</assistant>
Figure 5: Prompt the APPS benchmark for the fix action.
26

---

N.2 CWMB Prompts
<system>
You are experienced Python developer. You will be provided with incomplete code snippet from a Python program. The task this program supposed to perform described 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 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 set up the Environment, which defines all the variables described the above documentation, plus any additional variables needed to maintain the environment state or implement its functionality.
stt_state function to set custom value for the environment and its internal representation (you can assume that when "set_ state" is used, the task not done and internal variables should be set consequence). set state takes single argument as input: a state observation from the observation space defined above. a step function to predict 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 any additional functions you may need to complete this task. Do not write an example of how the class anything else. Be careful about edge cases. Make 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 a score It is of VITAL importance that you do not leave undefined any function, but implement each of them completely.
state and step functions and
write all
</user>
<assistant>
"python
(CODE_SC SO SO_FAR/ </assistant>
Figure 6: Prompt the CWMB for the generate action.
27

---

<system> You are an experienced Python developer. You will be provided with an incorrect code snippet from Python program. The task this program is supposed perform is described in the following user prompt. Your task is to rewrite the program so that it performs the task expected without any errors. You will be rewarded based the number of test cases your code passes. </system>
<user> {ENV DESCRIPTION)
## Class Definition
The class should be called "Environment". should have least:
an init function to set up the Environment, which defines all the variables described in the above documentation, plus any additional variables needed maintain the environment state or to implement its functionality.
a set_state function to set custom value for the environment and its internal representation (you can assume that when "eet_sttte" 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 predict step in the environment. The input parameters for the step function
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, described in the reward definition above.
boolean variable indicating if the episode done.
## Important Notes
Only produce the environment class, containing the _iii__ set_state and step functions and any additional functions you may need to complete this task. Do write 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 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 explanation of the difference between the ground-truit truth transition and the step function's '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 concrete, actionable fix for it. Finally, fix the program in its entirety following the suggestion. The expected output is in the format:
## Error explanation [your explanation of the error]
## Error location and wrong logic [where the еrror comes from and why]
## Fix suggestion [how to fix the error] ## Correct code
"'yython [your code]
## Incorrect code
You are provided with the following code snippet to fix.
"pythonn (CODE)
The code additionally makes wrong prediction about this input.
## Input
(INUT))
## Ground-truth output
(OUTPUT)
## Code incorrect outputs
[PREDICTION] <luser>
<assistant> ## Error explanation </assistant>
Figure 7: Prompt on the CWMB for the improve action.
28

---

<system>
You are an experienced Python developer. You will be provided with an incorrect Python program. The task this program 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 the variables described the above documentation, plus any additional variables needed to maintain the environment state or to implement its functionality.
set state function to set custom value for the environment and its internal representation (you can assume that when "set_state" is used, the task not done and internal variables should be set consequence). set state takes single argument as input: state observation from the observation space defined above.
a step function predict step in the environment. The input parameters for the step function are:
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.
boolean variable indicating if the episode 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 example of how to use the class or anything else. 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 a score It is of VITAL importance that you do not leave undefined any function, but implement each of them completely.
First, write explanation of the error and point the part of the code responsible for the error and why its logic is erroneous. Second, suggest how you would fix the error, reasoning about the problem. Finally fix the program in its entirety following the suggestion. The expected output is in the format:
## Error explanation
[your explanation of the error]
## Fix suggestion
[how to fix the error]
## Correct code
"pythhnn
[your code]
## Incorrect code
You are provided with the following code snippet to fix.
"python
(CODE)
[ERROR]
</user>
<assistant>
## Error explanation
</assistant>
Figure Prompt on the CWMB for the fix action.
29

---

N.3 Sample Environment Descriptions
For the CWMB we extract the description for each environment directly from the Gymnasium source code". 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-vl environment 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 the cart and the goal is to balance the pole by applying forces in the left and right direction on the cart.
## Action Space
The action is a 'ndarray' with shape '(1.)' which can take values '0, 1 indicating the direction of the fixed force the cart is pushed with.
-O: Push cart to the left- I: Push cart to the right
** ** The velocity that is reduced 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
## Observation Space
The observation is 'ndarray' with shape '(4.)* with the values corresponding to the following
positions and velocities:
Num Observation Min Max
101 Cart Position 14.8 4.81
111 Cart Velocity -Inf Inf
12 Pole Angle 1 -0.418 rad (-249) 0.418 rad (249)
13 1 Pole Angular Velocity IInf Inf
**Note:** While the ranges above denote the possible values for observation space of each element, is not reflective of the allowed values of the state space in unterminated episode. Particularly: The cart x-position (index 0) can be take values between '(-4.8, 4.8),, " but the episode terminates if the cart leaves the '(-2.4, 2.4) range. The pole angle can be observed between (-418, .488)* radians (or **4240**), but the episode terminates if the pole angle is not in the range "(-22095, 20955' (or #*
## Rewards
Since the goal keep the pole upright for long as possible, reward of '+1" for every step taken, including the termination step, is allotted. The threshold for rewards is 500 for v1 and 200 for vO.
## Starting State
All observations are assigned a uniformly random value in '(-0.05, 0.05)
## Episode End
The episode ends if any one of the following occurs:
1. Termination: Pole Angle is greater than +12 2. Termination: Cart Position is greater than
+2.4 (center the cart reaches the edge of the display) 3. Truncation: Episode length is greater than 500 (200 for v0)
Figure Example description for the CartPole-vl environment.
Examples of Generated Programs
We report examples of Code World Models generated by both Llama 3 and GPT-4 Turbo. We include CartPole-v] as an example of a discrete environment, Ant-v4 as an example of a continuous environment and RTFM.
'nttps://github.co httpss///ggthub..cco/aaaaaa Foundation//ymmaattttoooooo http:///aaaaooooooao//o//aa///ts hww.aaawwwwnnnnnnnnan..aaam envs/ classic_ control/cartpoe..py
30

---

Figure 10: Code World Model of the CartPole-vl environment generated by GIF-MCTS with Llama 3. For reference, the official implementation of the environment can be found at the official Gymnasium GitHub repository of the Farama Foundation (at gymnasium. envs. cllasscc_cnntrolt cartpole)
import numpy as up
class Environment:
def (seli:
self colf.paott.ooiitonn np random. uniform (-0.05, 0.05)
self.cart_velocity = np random. uniform (erfform-00000 (50.05, 0.05)
self self.pole_angle up. random. uniform( (-0.0 05, 0.05)
halutpoole_angular.elllooe np. up.random. uniform(-O. 05, 0.05)
self opld-eppssdee_eenth self ggaviity 9.8
10
11
self malffmass_cart cart 1.0
self mas__polseeeeeeeeppeppppp 0.1
self.length__pnl folceeerc__mag 10.0
0.5
14
15
self self.time_step 0.02 time step for the simulation
16
def set stt_statessel,, state):
18
self. corttposiiiion state
19
self colf.cart_vvleee veloctyy state[i]
20
self pole_aole state self aagllaangular_eelocci self aapisod__eengee
state
23
24
def step(self, action):
25
force self force_macaeggee if action -- else -self www.efforee..e...eeeee..eee force_rreepee costheta = np sintheta mp. (olleellf.ooeeeaggeeee
28
temp (force self mass_ self self. wwl.thhnnoool pole
self
* sintheta) (self mass_cart
self mass_pols)
angular_acceless accellrraccoliitiioon (self gravity sintheta costheta temp)
(self (self.1ength.poll. (self soelf.mass_care self.mass.pp....eeeeeeeeeaeeeeee (affmass__poe))
(4/3 selfmmass_pole costheta**2
30
linear_acee acceeraacceleotioonn temp self mass_pole self length_pole
angular, aacitaarcccleettiion.. costheta / (self carrt self mass_pole)
31
32
Update the state variables using the calculated accelerations new new_cart_velocige = self cart_velocity + linear_acceloratt
33
self time_step
34
new_pole_ pir_pole_anguarr..mm
self vole_angular__elocit;
- angular_ paauuaaracceleratiinn self, time_step
35
new cor__art_ppll cart_posiition self. cart..osttioo molf.carttvelocity.
self time_step
new_pole_and anglle self pole_angle self welf.pole_eeeeee
self. time_step
37
38
self wwffccartvvellll verl_velocity nes_cart_velocie
self cart_position new_cart_position
40
self aaalle_angular__eleeee self pole_angle new_pole_angle
coo.polle.ngguaar.ello.c.
41
42
self weweeppe_eeenneeeeeeeeeee +=
done abs (self. pole_angle) > 0. 2095 abs( self. cartcctappposiion) 2.4 or self, episode_length
45
46
reward 1 observation array'
47
[self cart_position, self cart wtt.laaartttettteeeeee velocity, self polpppplaaaaaa angle, self self.pole, wawleaagnullllllleeeeeeee return observation, reward, done
49
31

---

Figure 12: Code World Model of the CartPole-vl environment generated by GIF-MCTS with GPT-4. For reference, the official implementation of the environment can be found at the official Gymnasium GitHub repository of the Farama Foundation (at gymnasium. envs classicc cartpole))
import numpy np
class Environment:
def
._init__
self. cor_tpttstooo 0.0
self wel.ccart_elloitty 0.0
self.pole_angle self haaa.pooee.angula..ellooo..
0.0
10
self gravitr = 9.8
self cart_ mass 1.0
12
self solf.pole_mass
13
self tatal_ttlaaaaaa mass self cart_mass self polppplllaamee m_mass
14
self pole_length 0.5 # actually half the pole's length
self worce_rrceeeeeee 10.0
16
self. tau 0.02 time step, 20
17
18
self resetr)
def reset(self):
self cart_position np random. uniform (-0.05, 0.05)
22
sulf.cart_velocie np random. uniform (-0. 05, 0.05)
23
self pole_pole_an up random. uniform (-10.rmm000000 05, 0.05)
self wtw.ppaaalllrreeeeeeeeeeeee velocity np random. imiform(-O.00 0.05)
return mp. array([sell) cart_posiiion self.cart, cart_ velocity, + self pale.oole_angee self pole pole_angular, (aanuaar_eloccty]]]
26
def (self,,, state):
self cart_ position, self velocty,, self pole_angle self. pole.pole_aae wat.:poolaangllll angular_ velocity state
29
30
def step,,elf, action):
force self force_mag mag if action -- else -self force_mag
33
Equations for the model costheta = np. caelf pole_angle) sintheta = np. sin(self pole_angle)
34
35
31
temp (force welf.pole_neeeeee
+ self half.pole_saee poole_angulrr_aee velocity== aeteeaattalrreleotttya self welf.pole__engee sintheta) self totattttallaaaae mass
38
theta_ace acc (self. gravity * sintheta costheta temp)
beelf.pole_lengee
(4.0/3.0 self welf.pole_ppeeeee polpppolssasseeeeeeeeeeee mass costheta==2
self tatatootammmass
39
x_acc temp self. pole_mass self cole_length theta_acc
coatheta / talf.total_maan masa
40
Update the four state variables, using Buler's method self cart_position self. tau self. vertcvl self, cart caar__velcctty self tau x_acc self polf.pole_angle += self, tau * self self wangular_velociigge += self tau theta_ace
42
43
44
47
Check if the episode done done = colf.cart_position -2.4 or self. cart.postttoo 2.4 or self self.pole, polpppolaaggll angle < 2092095 self olg_aagle > 2095
48
49
52
32

---

53
# Since the problem spectfies to always give a reward of +1 reward 1.0
54
up- array ([self wtr.trrrrattiiooc.. self cart_ velocity, self pole_angle, self pwwppangularrraaeeeaaaa (angular_veloctyy)] reward, done
Figure 14: Code World Model of the Ant-v4 environment generated by GIF-MCTS with Llama 3. For reference, the official implementation of the environment can be found at the official Gymnasium GitHub repository of the Farama Foundation (at gymnasium. envs mavee.mjjco..o Ant_v4).
import numpy as np
class Environment:
def ...init__ (self,
healthy
ctri_cost__eiight0.55
contact_
cottact_coot_eeightt-----
10
contact_force_range-10.0
11
healthy.2_raane- 1.0),
12
terminate_vian hww.aaateeeeeeeemmaaaaieeeeo
reset_noise__caleelee-2
15
self hel--healthyrrowardd healthy_reward
16
self wtw.owarrd_rwwar__weigghh forvard__rward_veiige
17
self.ctrl.cosft celff.cttl.cost_wweittt strl_cost_weight
18
self wet.coontctt_costueeiggtt self wtt.caccttfooreeeaaneeeeeee self healh healthy.z_range self. hollittemminate__bbu_uueeeee
contact_cost_weil
19
contact_force_reat
20
21
terminate_when. when_unhealthy
self www.rreetnoiie_salllll reset_noise_scaal sos_sssss_cceeeeeeeeeeeeee___eee
self wwweeeeeccurreeeeeeeeeeeeeeeeeee cwsrenttppssiiionssssessecssses sarree_fr______seeeeeeee cotrmmoobeeeviionn..
Betbuue_ccureet_poosiie frommmss_oooooooooooo
wwtas/_ffooooootiiioooo observation
24
state self self._reede rit.s_reeettssattee
25
def www.ee_ssattttee (sset_state(self):
implement state reset logic here state mp. ze.zerrs(27)
28
29
state[o] [o] = 0.75 initial z-coordinate
30
state += random. normal (0, self wwweeeestiiitt pesst__oiss__caaee size=27)
return state
32
33
def set set_ ssa__t state):
34
self state state
def step(self, action): 37
implement dynamics model or physics engine to predict neat state next_state self pelff..predctt.erssttteeeee state, action) observation next_state
38
39
reward self www.caaateeeeerrreedeeoooooaaede self htt___hhec_eeeeeeee eeweeeeeeeeedeeeeeeeeeeeeeeeeeee leroodeeenddde (neet_sttaee
next_state)
41
42
return observation, reward, done
43
def wwppdiccteeeee srcictt_ext__seeeee mext_stateeeelf, state, action): imp lement dynamics model physics engine to predict neat state for ezample, using simple integrator # assume velocities are the last 14 components of the state velocities = state [13:]
46
47
48
new velocities velocities copy()
50
update front left leg, front right leg, back left leg, and back right leg
33

---

51
new ww_velocittie(::]] +- 0.01 action(:4] update corresponding body parts
52
new, [ += 0.01 action[4:]
Fim: slice new, velocsties to match shape of state[:13]
55
new_positions state[:13] 0.01 new_velocities[11
56
next_state = concatennat (ner_positions, new_velocities) return next_state
39
def calluulate_rvvrd(eelll action, next_state):
60
implement reward calculation logic here healthy_rewae reward self healthy_reward forward_reward aelf forward
61
[next_state[13] self s13]) # a-coordinate velocity ctrl_cost self cost__ohst_wihht np.sum sum square (action))
64
np.
65
contact_cost 0 # smp lement contact cost calculation logéc here
reward healthy_rewaard + forward_reward ctrl_cost contact_cost return reward
69
def _cceckk __chck_episode episode_ (self, state): # implement episode end check logic here
70
done = not (self hwwlthhaa_raagggeee [0] <=
72
state <= self healthy_z_raanue range [1])
73
return done
Figure 16: Code World Model of the Ant-v4 environment generated by GIF-MCTS with GPT-4. For reference, the official implementation of the environment can be found the official Gymnasium GitHub repository of the Farama Foundation (at gymnasium. envs. mujoco Ant_v4).
import numpy np
class Environment:
def _nnii___sselt
ssee_comtaattfooocssfallssss terminate__ _eeettte_w________eeeeee____e_ee unhealthy-True,
health (0.2, 1.0), healthy maalthh_reward1.00
+ ctrl_cost_weigee -eroosstteetttt.555__11 5, contact_comt corslosssttwihihieeeeeeee weight-0.5
mantact_force_range (0, 10), ceset_nnise_scale-tee frame_skip-5):
self seccudd_cureent_poattiosss from
ttt.:iimes.rrooobeerttiiioo observation
+ exclude_curremn_.piiiioo... laaoanss_fo_oooeeeeeeeee httprrccco___rooooooerrrtrioo observatieennn
self.use_com www.__cootttttfooreeeea use_contact_ forces
self. saaniiaatt_when_ubballll
= tateiaattt_uheuuunhllaaaa
self healthy-z_raagge healthy_z_range
11
self heaad healthy_reward
12
solf.cctr_.cst_weeigee ctxl_cost_weight
13
calfccontact_coott
contact_cost_ veight
self self. www.tttctfooreeeeeetteeeee self reset__nosse_scale reset_noise_scale
contact_force_rran
16
self frame_skip frame_skip self dt self frame_skip 0.01
17
18
if self srccuudecurrett.ppttttosso._rr self obserrvatiou..scccssii. 27
20
21
else:
22
self observation.space,e 29
if self cootccctcc___o:es:: self. wbwerrvttioonsspce__.ii..
25
# 14 body parts contact
+ force values
self state np. (self observation_space, 76
28
self. self.last_x_post last_x_posiition 0.0
29
34

---

30
def set set_state@self,, state):
31
self.atata state state
self lwsllass__________ __ositixosiiion state [0] if not aroomoobeervation else
self.
33
34
def ssep, action):
if not (-1 <= action) al1() not (action <= 1). alaal1):
raise ValueError (aarroor""ctttt ("Action must be within the range [-1, 1] for each element.")
37
38
Simulate state transition (simplified) self state += na-random normal (0, self reset_ reset_noise_scale, self wtt..oberratiiosssceeesiit
40
41
x_position self state[] if not
self wew.eeeeeeeeeeeeeeee Bec_ludd_currettpeee_ wttp:pppttstttotooooeeeeeeeoeeee
else 0.0
z_position
if not
self beccludecurrenttppoittttt httrrriossttoommaerrtttittt else self state
43
44
Calculate reward components
forward_ reward (x_position self lwallaa___p__ s_ttt__posiiion) / self dt
cost self wttp_coosttteeeeeeeeee
mp sum (up..nnpuuna...ee square (action))
47
contact_cost
48
49
if self muse_contact_foreee
Last 84 entries contact forces
contact_forces self [-84]]
52
clipped_forces np app.cllpccontat_fooree
www.looottttttooccecceeo fartact_forceerageee range [o], self contact_force_ramt [1])
53
contact_cost self contact_cost_wweigt
np. sum (wp. square (cuarrecliiped__force)))
54
55
reward self healthy_reeee
forward_revard ctrl_cost.
contact_ cost self _ast____po_tt__ poosition x_position
58
Determine if the epssode done unhealthy = not (self. healthy <= z_position <=
59
self hww.hhelhht..... (_range[]]] [1])
done unhealthy and self
return self state, reward, done
61
Figure 18: RTFM environment generated by GIF-MCTS with Llama 3
import numpy np
class Environment:
def
_seltt (oatontttsdddeeeee
self ceinster_elemntte - ('demon': 'fire', 'dragon': 'fire', 'jimn'
'lightning', mmedusa:: 'poison', 'bandit'; 'cold', 'wolf': 'cold', 'goblin': 'lightning', 'mage': 'lightning' 'spider': 'poison') self, @elffitte_modiifeeee mod.miodiieerseeeeo ('axe': grandmastera'', 'bow': 'blessed',
'daggers': 'shimmering', 'hammer': 'gleaming', 'polearm': 'fanatical', 'shield': 'myaterious', 'staff': 'arcane', 'sword': 'soldiers'}
self monster-groups { ('demon'; 'rebel enclave', 'dragon'; 'order
4 the forest', 'jinn': 'star alliance', 'medusa': 'order of the forest', 'bandit': 'rebel enclave', 'volf': 'order of the forest', 'goblin': 'star alliance', 'mage': 'star alliance', 'spider': 'rebel enclave')
self woa..goooggouppp 'order of the forest'
35

---

self.
('fire: [' I'grandmasters' 'soldiers') ,
'lightning': ['arcane', 'blessed'] 1, 'poison': I'gleaming',
'mysterious 'cold': ['fanatical', "shhmmering '}>
self.board mp. full 6, 2), 'empty', dtype=object)
11
self bolddbo 0, :] 'wall'
12
self self.board[;, -1, :] = 'wall'
self board :] 'wall'
self Solf.board[-1 ;, :] 'wall'
15
ieve.tnventory 'empty'
16
self.player_positae (1, 1)
17
boad. player.position(0], self wtw.:pllllayst.aaee. 0] 'you'
19
def set_state(self, state):
20
self. board sta self self. inventory state[1] for in range
for range
24
self biardli, j, 'you':
25
huffppaaye__osittee (i, j)
break
28
def step(self, action):
29
action_dict (0: "Stay", 1: "Up", 2: "Doun", 3: "Left", 4: "Right")
30
new_player_position - List(sell player_ppsitionn if action_dict -- "Up":
elif action_dict([cttion]] -
new.player_poaitiiee
31
"Doun":
34
new www_player_positionn [0] += 1
elif action_dict([ctiool) == "Left":
dev.playeer_posiiion/]
31
elif action_dict([ctionn) "Right";
38
new her_player_positiion
39
if (0 <= hew_player_positionn00 and <= new.player.position[[]
and
41
self. board www..eee..lllleepppttiiolllooooa new.player.positiont]] 0]
+ I= 'wall')
42
self boardd Seaffppaaee_ppee playtr__ossition(0),
self player_position[[]] 0] 'mpty''
44
self player_positiioon
45
self. boarddserdd..fooooooo @ooardsseff.llayr:: player playye__ossitonn[0] [0],
46
self sutf.player.position
0] = 'you'
frame (self board bopyd.copy(), self. inventoren))r))
49
valid_actions
for in r5nge(5) self _is
(i)]
50
reward = self cefff-gge__eeardd_ done = self._is_donee return frame, valid_ actions, reward, done
51
53
54
def iss ia_waalid_action(see. action):
55
action_dict = {0: "Stay", 1: "Up", 2: "Down", 3: "Left", 4: "Right") new.player.position = list (self waw.lff.layee__pooo posttir_position) if action_ action_dicc((ctiionl
"Up": --
58
59
elif action_dict(action)
"Down":
60
new (icioo__dct((cctionn] == "Left":
+= 1
elif action_dict new, elif action_dict
"Right";
64
dew.player__position +=
return (0 new.player_positi.on [0] 6 and <=
new.playor.pooitionnl] and
36

---

67
self. board haarrdeew_paaer_ppsstto/// now.player_position[]] o]
I= 'wall')
def _get, reward
70
for in range(6);
71
for in r6nge(6);
if self board J. o] in self canster_slemente and molstmonster monster_groups [self board(i, J, 0]] ==
self goal.r
73
if self inventrry !- 'empty':
74
item_name = aelf inventory.spliti split(' [)[
stem_modifier = self wtwm_modifiersle (item [item_name]
monster_ee ceeater_element self monater elements [self board
4 J. 0]]
77
if item @eem_modifier in self Alement_weaknessee [monster_e eetteerreeemeetee element]]
return
79
-1
80
return O
81
def ss_done
for range
84
for in range
85
if self board [1, J, 0] in self monster. conster_elementee elements and self www.mmnnnterggrooo [self board[i, J. 0]]
self get.gglllggrupp return False
87
return True
88
Figure 20: RTFM environment generated by GIF-MCTS with GPT-4.
import numpy np
class Environment:
def _se1)t
self board full 6, 2),
'empty', dbjee=object)
self.board[: 0, :] 'wall'
self belf.board[: [:, 5, :] 'wall'
self bolrdboard[] [0, :] 'wall'
10
self bolrd [5, :] 'wall'
12
self Salf-player_por (3,
13
self.board. [3, 1, 0] 'you'
14
self self.invent inventryy 'empty'
self moooooonnnoouue. =
17
'star alliance': ['mage', 'goblin', 'jinn'], 'order of the forest': ['dragon', 'medusa', 'wolf'], 'rebel enclave': ['demon', 'spider', 'bandit']
18
19
22
self mnoosster_eemment {
23
'dragon': 'cold', 'medusa': 'poison', 'volf': 'poison', 'demon' 'free', 'spider': 'lightning', 'bandit': 'cold', 'mage': 'lightning', 'goblin': 'fire', 'jinn': 'cold'
26
27
78
self item hewemmimoiiereeee =
grandmasters': 'fire', 'blessed'; 'lightning', 'shimmering'
cold',
30
gleaming': 'poison', 'fanatical' 'cold', 'mysterious': 'poison',
37

---

31
'ooldier': 'fire', 'arcane: 'lightning'
32
}
self wet.aictionnidce :o: "Stay", "Up",
35
"Down", 3: "Left", 4: "Right")
36
self goal_group order of the forest'
def set sst_stateseelf, state): board, inventory state self. board board
39
40
41
self inventoyy inventory
pos np. where (self. board,.; :, 0] == 'you')
if poa[0] size 0:
44
self playerrayr_pos (pos[o] [0], pos [1] [o])
45
46
def step(self, action):
direction = self aataaacttiiiaeeeeaeee lottm_dict[atttt [action)
x, y self player_pos
49
dx, dy x, if direction == "Up": dx = -1 elif direction "Down": dx +1
50
53
54
elif direction "Left": dy - y elif direction -- "Right": dy y
55
58
59
if self board[dx, dy, 0] == 'wall':
dx, dy x,
new_pos (dx, dy)
63
cell content self. board[neu_poo nev_pos [1], 0] game_done False reward
64
cell_content 'empt'' and cell_content 'wall': content_parta = cell_content .plit()
68
if len wecontett_ppartteeeee -- and content.parts(1] in
self wwnstensteere.. mamattereeeeents::
70
monster content_parts
71
self monstor monster_elementa elements [monster] monster_group = next ((group for group, monsters in self self.monster www.monsterrrruup.... proupa.items:
monster_elemene
72
if monster in monsters), None)
74
self inventory != 'empty':
75
item_modifier, iten self.inventory. haventory.sslitte
76
if tal.ffitem_modifficm [ites_moddiiii ceteemmodiiier == monster_element
and monster_group == self self.goal_group reward game_done True else:
7M
79
reward -1
game_dome done True
else:
83
reward -1
84
game_done = True
elif content_parts [o] in self.item iittttttmmmma modemiediieer:
self inventory cell.content
88
if not game_done:
self board y, 0] = 'empty'
self. board (nev_pos [0], new_pos [1], 0] ='ou''
self player_pos new_pos
38

---

92
91
valid actions [a for in aelf action_dict if self board[self player_ [po[[]] +
0, -1, 1, 0, 0) [a], self playpr_paaalll (0, o, -1, 1) [a], o] 'wall']
95
return (self. board. copy O), self inventory), up array (valid_ artiidd)) reward, game_ done
39

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: All claims in Abstract and Introduction are backed by experiments found in the Experiments Section 5 or in the Appendix, and clearly make reference to them.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
Limitations discussed in the Limitations Section 6.1.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model vell-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
40

---

Justification:
No theoretical results are presented in this work.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We comprehensively report, either in the main text of in the appendix, all parameters, prompts and experimental details to reproduce results.
Guidelines:
The answer NA means that the paper does include experiments.
If the paper includes experiments, No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
41

---

Answer: [Yes]
Justification: Anonymized code will be provided in zip file together with the submission and released with URL referenced in the paper upon acceptance of this work.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (uttps://nips.cm cc/
Policy) for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to reproduce the results. See the NeurlPS code and data submission guidelines (https: //niss "nips.ccppppi wwtpcccppppppppeeeeeeeeeec wttpc//ggddeeeeeeee /eodddeS/ooommmiii wttpsuummissoooooololoooooo ssionPolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: All details of our experiments are specified. We will also release, together with the code, the data collected in RI. environments, as they present stochasticity in terms of initial conditions.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: Errors are reported on all the main experimental tables. All errors are errors of the mean value, not standard deviations.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
42

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: All computational resources have been accounted for Section M of the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics attps://neurips: wwtppcpwpppcceeeeeeeeeeeeeeeeeee ctt.e//Ethhiiiieeeeeeeeee ines?
Answer: [Yes]
Justification: The authors confirm that the research conducted in the paper conforms, in every respect, with the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No. they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We report discussion on potential positive and negative impacts in Appendix
A.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
43

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess 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 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 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 system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
To the best of our knowledge, risk is posed by this work.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have 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 usage guidelines 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 best faith effort.
12. Licenses for existing assets
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?
Question:
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 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 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.
44

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. wwwccoo////oo/oo datasatssss has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Question:
Answer: [Yes]
We do introduce new benchmark derived from existing openly available resources, detailing how to re-create and reporting its details appendix. The open- sourcing of the code for this work will contain instructions on how to run the benchmark.
Justification:
Guidelines:
The answer NA means that the paper does not release assets.
Researchers should communicate the details the dataset/code/mooell as part 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 used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsoureing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well 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 the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
45

---

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


Reviews:
[{"review_id": "VQTdUhrRRS", "reviewer": "Reviewer_WBDc", "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.", "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)?\n\nOffline pre-training seems to be expensive. For domains like code generation, is pre-training necessary? For new domains that the agent may not have seen in the pre-training data, does it help to use in-context learning to show the agent some example trajectories?", "limitations": "The limitations are clearly discussed in the paper.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper is well-written and effectively demonstrates that MCTS can facilitate agentic behaviors and problem-solving across multiple domains. The motivation and the rationale are clear, as we expect the idea of agentic behavior of searching with diverse actions can help solve more complex problems.\n\nThe appendix includes extensive ablative studies that offer a detailed analysis of the method's components.", "weaknesses": "**Reliance on pre-training.** I didn’t fully understand why learning the world model is necessary. MCTS is an inference-time algorithm. The reliance on pre-training for the large language models might be a limitation.\n\nSpecifically, does GIF-MCTS need to predict all transitions well in the offline dataset to perform well on a task? For coding, the agent doesn’t need to accurately predict how a human solves a problem. It’s sufficient if the agent can find a correct solution, possibly using MCTS.\n\n**Novelty.** There seems to be limited novelty in this framework. Related works have explored agentic behaviors of using “implementing a function”, “fixing bug” as actions, like SWE-agent [1]. Using tree search and sequential decision-making are also explored in the literature [2]. The novelty seems to be a combination of both.\n\nReferences\n\n[1] Yang, John, et al. \"Swe-agent: Agent-computer interfaces enable automated software engineering.\" arXiv preprint arXiv:2405.15793 (2024).\n\n[2] Zhang, Shun, et al. \"Planning with large language models for code generation.\" arXiv preprint arXiv:2303.05510 (2023).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review our work and provide your detailed feedback. We are grateful for your comments and we hope that the following will clarify the details and novelty of our work.\n\n**1. Pre-training:**\n> Reliance on pre-training. I didn’t fully understand why learning the world model is necessary. MCTS is an inference-time algorithm. The reliance on pre-training for the large language models might be a limitation.\n\nand \n\n> 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?\n\nWe would like to clarify a potential misunderstanding regarding our approach. We do not train nor fine-tune any LLM, but only use off-the-shelf LLMs (e.g., Llama-3, GPT-4). \nSpecifically, we learn a Code World Model (CWM) that can be used in model-based RL as a substitute for the dynamics and reward functions of the environment, allowing an agent to search over action strategies using methods such as CEM or standard MCTS. \nSeparately, we propose GIF-MCTS, an iterative generation method tailored for generating CWMs with LLMs. Within it, we make use of in-context learning to show trajectories to the model as you suggested, specifically with the improve action. We developed GIF-MCTS specifically because we found prompting the LLM directly to not be effective.\n\n**2. Predicting transitions:**\n>Specifically, does GIF-MCTS need to predict all transitions well in the offline dataset to perform well on a task? \n\nIn our work, we assume this is the case. More precisely, GIF-MCTS produces the CWM, which in turn should predict correctly all transitions in the offline dataset to perform well as a world model. However, there may be cases where the model does not need to be uniformly accurate everywhere, particularly in suboptimal regions rarely visited by the planning policy. \n\n> For coding, the agent doesn’t need to accurately predict how a human solves a problem. It’s sufficient if the agent can find a correct solution, possibly using MCTS.\n\nIt is not fully clear to us what the reviewer means by this. Our work involves two agents: the GIF-MCTS one, whose objective is to write a correct CWM, and the model-based RL agent, which uses vanilla MCTS (or CEM) with the CWM to solve the RL task (e.g., CartPole). Neither of the two ever tries to predict how a human solves a problem and no human data is used to learn the world model. For coding, the reward function used by GIF-MCTS (which guides the generation process) is the fraction of passed unit tests-- no full solutions are compared against.\n\n**3. Novelty:**\n> Novelty. There seems to be limited novelty in this framework. Related works have explored agentic behaviors of using “implementing a function”, “fixing bug” as actions, like SWE-agent [1]. Using tree search and sequential decision-making are also explored in the literature [2]. The novelty seems to be a combination of both.\n\nWe respectfully disagree. Our main novelty lies in formulating the concept of Code World Models, a framework to learn RL world models written in code. Additionally, we introduce a benchmark to evaluate this approach, and finally propose GIF-MCTS as a code-generation algorithm with an LLM, specifically tailored for the CWM framework. \nHence, while GIF-MCTS can be seen as novel combination of agentic behaviour and tree search, our novelty goes beyond that.\n\n**4. Comparison with WorldCoder:**\n> 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)?\n\nGood point, thank you for raising it. We believe that GIF-MCTS outperforms WorldCoder because it considers a more diverse set of programs. \nThe main difference is that WorldCoder initially generates a single complete program, which becomes the ancestor for all future programs. 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, which better explores the solution space. \nA further ablation study (\"No Generate action\" in Table 6 of the Appendix) supports this finding: using a tree search (like GIF-MCTS) but only with refinement actions (like WorldCoder) results in lower performance compared to our method.\nWe will add this explanation to the Discussion section in the main text.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CitLr2AVzm", "reviewer": "Reviewer_u7M2", "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.\n\nTo me, the contributions of this paper are:\n- The GIF-MCTS method for effective iterative program synthesis wrt available test cases, which is separately evaluated on the APPS code generation benchmark.\n- The concept of code world models (as outlined above) which shows promise in efficiently obtaining world models from observational data\n- A benchmark which covers common environments used in Deep RL that enables the community to propose further advances", "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": "I found the paper's discussion of its limitations to be upfront and comprehensive and have nothing to add.", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "I found the main idea of the paper to be quite original and interesting, i.e., generating the code for a world model rather than utilizing LLMs directly. The authors show that this works quite well in environments with discrete action space: they often obtain perfect or near-perfect prediction accuracy as well as rewards approaching those achieved by planning with an oracle world model. In continuous environments, where the original implementations rely on a physics simulation engine, the results are less convincing -- this is understandable to me as I would expect LLMs (or humans, for that matter) having trouble reimplementing those without access to the original or a substitute simulation engine. Whether code world models are a promising direction to ultimately tackle planning in complex or real-life environments remains to be seen, of course; however, I would expect paper at hand and the proposed benchmark to inspire further works in this relevant direction.\n\nI also liked the separate evaluation of their MCTS code synthesis method on APPS where it produces strong results, and it's worth noting that there are full papers proposing similar iterative prompting for code generation. I hope the authors make this part of the paper available for easy re-use.\n\nThe paper itself is well-written and, while it is pretty packed, I found that I could easily locate and understand the method and details on problem setup and evaluations.", "weaknesses": "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\n\nFor the planning results it would be good to list the true rewards obtained along with the rewards of a random policy and SOTA model-free/model-based results from the RL literature. While beating SOTA results here is not the focus of the paper, it would provide a helpful perspective for readers familiar with those environments.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the kind words about our work. We are glad that you found it interesting and believe it will inspire future works! We agree that continuous physics simulations are particularly challenging and we also were not particularly surprised by that, but we will be looking to address this in future work, for example by integrating the method with external physics simulators and APIs.\n\n**1. Examples of LLM-generated environments:**\n>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\n\nIn the limited space of the rebuttal, we include a Code World Model (CWM) example for the Ant environment (continuous action and state spaces) in the common response.\nIn the final manuscript, we will include two more (one gym environment with discrete action space and RTFM), together with how they compare to the ground truth source code. \nWe performed an additional qualitative study on the generated programs and found that while certain aspects of the gym environments resemble available online implementations (such as the values of constants in the CartPole environment), the logic of the step function is distinct from that of the reference model for all generations.\n\n**2. Comparison with RL:**\n> 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.\n\nWe performed a direct comparison to a random baseline, the oracle baseline (using the same planner as the CWMs but with the ground truth environment) and with Conservative Q-Learning [1], a popular offline RL baseline that can work for both discrete and continuous action spaces. We train CQL on the same small dataset of trajectories used for the CWMs. The results are reported in the markdown tables of the common response. \nOverall, there is a balance between CQL and CWMs, with CWMs being more suited to discrete tasks and CQL outperforming CWMs in complex physics tasks. We also observe severe overfitting happening in CQL almost immediately, likely due to the small size of the provided dataset. As we point out in the paper, sample efficiency is one of the main promises of the CWM approach, as very few trajectories are needed to validate the model. There is also a considerable gap between CQL and the oracle planner, which represents an upper bound for the CWM approach given our choice of planners.\n\n**3. LLM sample budget:**\n> 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.\n\nThe mention of 50 calls in the Appendix was a typo, thank you for catching it! Indeed, the correct budget used for APPS is 20 LLM calls corresponding to 20 programs generated and evaluated on the unit tests. As for the number of problems, we work with the test set of the APPS dataset which is composed of 5000 total problems, of which 1000 make up the competition split, so 1000 is the correct number of problems we evaluated on. Could you point us to where you got the impression that the competition split is smaller?\n\n**References:**\n\n[1]: Kumar, Aviral, et al. \"Conservative q-learning for offline reinforcement learning.\" Advances in Neural Information Processing Systems 33 (2020): 1179-1191.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9Pug8GoMRV", "reviewer": "Reviewer_3vW3", "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. \n\nThe system is evaluated on the synthesis of programs for solving competition-level problems and programs encoding world models. As baselines, the experiments include recent works such as CodeRL, Parsel, and WorldCoder.", "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.\n2. How about running the fix bug action a number of times whenever a state representing a buggy program is generated? Instead of allowing MCTS to decide when to debug, why not debug right away?\n3. Please comment on the data contamination issue mentioned above and clarify what is the research question behind this work.\n4. Please comment on the statistical metrics and empirical design of the experiments.", "limitations": "In weaknesses, please see the parts related to overclaiming and data contamination.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "After reading the paper once, I believe its main contribution is to show the power of tree search in the context of synthesis with LLMs. A simple MCTS-based algorithm already produces programs able to attain better quality metrics than other methods not using search.\n\nThe experiments are nicely done from an ablation perspective: I quite enjoyed Table 6 in the appendix, where different versions of GIF-MCTS are evaluated.", "weaknesses": "The paper has 4 weaknesses that are worth mentioning, so the authors can react to them in the rebuttal.\n\n**1. Clarity of the paper could be much improved.**\n\nWhen reading Figure 2 and the description around it, I got confused with the use of the word \"rollout.\" I think it is being used to mean two different things. The first is a portion of the code that a node represents. The second is the LLM call that generates that portion.\n\nI did not find in the paper the specification in terms of how many lines are generated in each \"generate new lines\" action and how many are generated in the rollouts.\n\n**2. The paper overclaims**\n\n> we improve the trade-off between exploration and exploitation in the UCT formula used for action selection.\n\nWhat was probably meant was that the paper uses a heuristic that achieves better empirical results. There are no theoretical improvements in the trade-off.\n\n> Calling code instead of LLMs for planning has the advantages of being precise, reliable, interpretable, and extremely efficient.\n\nNone of these properties are actually properly evaluated in the paper. This might be a bit of nitpicking, but code can also be imprecise, unreliable, uninterpretable, and extremely inefficient. Anyone working as a programmer has already experienced all these negative properties of code.\n\n**3. Lack of discussion on data contamination**\n\nWith the exception of the RTFM benchmark, nothing is said or discussed about possible data contamination. Is the search algorithm helping the LLM remember what is in its training data or is the overall solution doing something else? Are the models for RL tasks available online as code? Are the solutions to the programming problems available online?\n\nThe lack of discussion about data contamination brings me to the question of what is the research question behind this work. Is it a retrieval question (i.e., can the system recover information used in the LLM's training data) or a reasoning question (i.e., can the system write programs that solve problems)?\n\n**4. Statistical significance is unclear**\n\nIt is not clear whether the search was run only once or many times. The averages presented seem to be over a number of problems and episodes, but not on the number of times the entire system was executed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the thorough review and the valuable suggestions to improve the clarity and soundness of our work. We appreciate your recognition of the power of GIF-MCTS and of the usefulness of the ablation studies on our method. We detail in the following how we are going to address your comments in the final version of our work.\n\n**1. Clarity.**\n\n**Rollout.** We consider each node as a full program, comprised of a \"state\" part and a \"rollout\" part. The **state** is the **first $L \\times d$ lines**, where $L$ is the number of generated new lines per node and $d$ is the depth of a node in the tree. The **remaining part of the program** is denoted as the **rollout**, and we consistently use the term to mean only that. **LLM calls** correspond to actions (edges) in the tree, and each of them **produces a full program** (node).\nPlease let us know if there is a specific part of the manuscript that creates this confusion and we will clarify it in the final version.\n\n**Generated lines.** The number of generated new lines $L$ is equal to 2 (Table 4 in the Appendix); we will report this information also in the main text.\nWe also realised that in the main text we used $l$ rather than $L$ and we will fix this. \n\n**2. Claims.** We never meant to present imprecise claims and we thank the reviewer for pointing out the need for better wording of them.\n> we improve the trade-off between exploration and exploitation in the UCT formula used for action selection.\n\nWe will change to \"we *propose a heuristic that empirically improves* the trade-off between exploration and exploitation in the UCT formula used for action selection\" (justified in Appendix A). \n\n> Calling code instead of LLMs for planning has the advantages of being precise, reliable, interpretable, and extremely efficient.\n\nWe will change to \"Calling code instead of LLMs for planning has *potential* to be *more* precise, reliable, interpretable, and extremely efficient.\"\n\nAddressing the mentioned properties one by one:\n- Precise: code can natively access a calculator and as such is numerically precise. \n- Reliable: a program passing all unit tests can be deemed more reliable than a black-box prediction from an LLM.\n- Interpretable: as mentioned in our response with reviewer u7M2, we will add concrete examples of Code World Models (CWMs) generated by our method, as mentioned to Reviewer u7M2. We include one example on Ant-v4 in the common response. In general, we find that all generated programs are clearly interpretable by a proficient human programmer.\n- Extremely efficient: we report in Table 7 of the appendix (referred in the Introduction section) the inference time of our generated CWMs vs inference time of an LLM. **CWMs are 4 to 6 orders of magnitude faster.**\n\n**3. Data contamination.**\n\n**RTFM Benchmark:** The RTFM benchmark is not available online, and our method's success on it provides evidence that our solution is not merely retrieving information from the LLM's training data.\n\n**Programming Problems:** 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. 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, 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.\n\n**Gym Environments:** While we observe that some parts of the gym 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. We include one model-generated environment in the common response, and we will add more in the final manuscript, along with links to the most popular online implementations.\n\n**Fair Comparison:** Data contamination is a known issue in almost all studies involving LLMs, not just ours. However, our method is compared to baselines using the same underlying models, ensuring that its superior performance is not biased by potential data contamination.\n\nWe will add this discussion to the Appendix of the paper.\n\n**4. Statistical significance.** The search was run once per each problem, due to the computational cost. For APPS, considering that we have 1000 independent problems, this is more than enough to tell apart with statistical significance the results of GIF-MCTS and the baselines. For CWMB, the uncertainity is higher. We ran 2 extra seeds for each environment (resulting in 18 environments * 3 seeds total executions for our method) for GIF-MCTS and WorldCoder using Llama 3 70B and obtain the following results:\n\n| Model                | Method          | Budget | Discrete Accuracy     | Discrete Return        | Continuous Accuracy      | Continuous Return      |\n|----------------------|-----------------|---------|---------------|---------------|---------------|---------------|\n| Llama 3 70B| GIF-MCTS (ours) | 50      | **0.84±0.03** | **0.76±0.03** | **0.35±0.03** | **0.22±0.01** |\n| Llama 3 70B | WorldCoder    | 50      | 0.79±0.04     | 0.60±0.04     | 0.32±0.03     | 0.19±0.01     |\n\nCWM still outperforms WorldCoder with smaller error margins, confirming the statistical significance of our method, especially in the discrete case. We unfortunately could not repeat this for GPT-4 due to budget concerns.\n\n**5. When to fix bugs.** Great question. Sometimes, generating a bug-free program from scratch will be more promising than trying to fix a buggy program multiple times, whereas other times the opposite will be true. Rather than deciding a priori how many times to debug right away, we believe that letting MCTS decide will lead to lower regret. A further relevant discussion is present in Appendix A.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kH2xomkeeu", "reviewer": "Reviewer_49z7", "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.", "questions": "See above", "limitations": "Yes", "rating": 5, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "**Motivation and intuition**\n- The motivation for modeling environments using codes is interesting and convincing, and this work presents an effective framework to achieve it.\n\n**Technical contribution**\n- Given that WorldCode has explored the online setup, this work explores the offline setup, which offers an unique perspective to the code world model idea.\n\n**Clarity**\n- The overall writing is clear. The authors utilize figures well to illustrate the ideas.", "weaknesses": "**Dataset collection**\nThe authors stated that the dataset \"includes at least some low-scoring and some relatively high-scoring behavior.\" I am wondering how we can do this without accessing true reward functions. Given a new environment that we have not seen before, how could we collect such a dataset to learn a code world model? \n\n**The effect of dataset quality**\nIt is unclear how the quality of the dataset would affect the code world model performance. Can we learn a reasonable code world model with just trajectories collected by a random policy? \n\n**Online setting**\nThis paper focuses on an offline setup, where the code world model learns from a pre-collected offline dataset. I am curious how we can extend the framework so that it can improve the code world model and the policy constructed from the offline dataset, just like how the research studies how to fine-tune policies learned using offline RL algorithms.\n\n**Access to physics engines as a tool**\nThe proposed framework writes a code to represent a world model for each environment from scratch. It could be difficult to precisely simulate some interaction that requires modeling physics. It would be interesting to see if granting access to physics engines, e.g., calling MuJoCo, could generalize to physics-intensive environments.\n\n**Comparison to offline RL methods**\nThis work mainly compares the proposed framework with WorldCoder, which also builds a world model. I am curious about how the proposed framework compares to offline RL methods, which simply learn a policy from the dataset to maximize the return.\n\n**Related work: programmatic RL**\nThe related work section is not comprehensive. Highly relevant programmatic RL works, which represent policies as programs and interact with environments by executing the programs, are missed from the section, such as:\n- Programmatically Interpretable Reinforcement Learning (ICML 2018)\n- Imitation-projected programmatic reinforcement learning (NeurIPS 2019)\n- Programmatic Reinforcement Learning without Oracles (ICLR 2022)\n- Learning to synthesize programs as interpretable and generalizable policies (NeurIPS 2021)\n- Hierarchical programmatic reinforcement learning via learning to compose programs (ICML 2023)\n- Show me the way! Bilevel search for synthesizing programmatic strategies (AAAI 2023)\n- Reclaiming the source of programmatic policies: Programmatic versus latent spaces (ICLR 2024)\n\n**Clarity: offline setup**\nWhen I read the introduction and the abstract, I thought it focused on an online model-based RL setting. Then, I realized this is an offline setup while reading Section 3. I suggest the authors revise the introduction and the abstract to align the expectations.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful review. We appreciate that you found our motivation compelling and our framework effective. We also appreciate your recognition of our unique perspective on the offline setup of the code world model idea. In the following we address your concerns and we are confident that by including the changes described below the manuscript's quality will be significantly improved.\n \n**1. Dataset collection.** To collect a dataset for a novel environment, one could in practice use a specialised exploration policy, such as in [1], [2] and [3] to mention a few. Including some relatively high-scoring behavior trajectories is common-practice in many offline RL datasets (see e.g., [4]).\nIn case no reward is accessible from the environment (if we interpret correctly your question), one could still collect transitions without rewards and learn a partial Code World Model (CWM) that predicts only the next state. In general, dataset quality and exploration for offline RL are active areas of research, for which orthogonal advancements could benefit CWMs.\n\n**2. The effect of dataset quality.** Yes, in principle we can learn a reasonable CWM only with trajectories collected by a random policy. For example in our experiments on RTFM, we employed exclusively a random policy to collect the dataset, which covered the state space sufficiently well and in the best case we obtained a perfect model (Table 3 of the paper). \nIf the environment is hard to explore and we are provided only a few random trajectories, it is still be possible to learn a CWM if the language description of the environment is accurate enough.\nWe ran an additional experiment on RTFM: we collected 10 trajectories all resulting in failures, so that a reward of +1 is never observed. We synthesized a CWM with GIF-MCTS and GPT-4 using 50 calls. The resulting CWM is 100% accurate on the 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.\n\n**3. Online setting.** This is certainly possible. The simplest procedure is to iteratively **collect data** by interacting with the environment and **update the world model** using the new data, similarly to what is done in the model-based RL literature (see e.g. [5]). To collect data one could use either an exploration policy (point 1) or a variant of the planning algorithms used in our paper with extra noise added for exploration. Our method can be easily modified to update the CWM, for example by re-using the previous tree search and evaluating each program (node) on the updated dataset, before continuing to expand the tree.\n\nIn general, we chose to work within the offline setup to focus on studying the overall applicability of the CWM approach without relying on specific exploration approaches that would affect the quality of the training dataset and possibly skew the results. We will add the discussion about dataset collection, dataset quality and extension to online setting to the paper.\n\n**4. Access to physics engines as a tool.** This is most definitely a valid direction for future research and we will mention that. We believe CWMs could be greatly empowered by access to simulation tools. In the specific case of MuJoCo, the simulations are almost fully taken care of by the library, and the resulting CWM would reduce to a trivial function call. Therefore, we feel that to be generalizable to other environments or embodiments, the simulation tool should operate at a more fundamental level.\n\n**5. Comparison to offline RL methods.** We performed a direct comparison with Conservative Q-Learning (CQL), a popular offline RL baseline. We train CQL on the same small dataset of trajectories used for the CWMs. The results are reported in the markdown tables of the common response.\nOverall, there is a balance between CQL and CWMs, with CWMs being more suited to discrete tasks and CQL outperforming CWMs in complex physics tasks. We also observe severe overfitting happening in CQL almost immediately, likely due to the small size of the provided dataset. As we point out in the paper, sample efficiency is one of the main promises of the CWM approach, as very few trajectories are needed to validate the model.\n\n**6. Related work: programmatic RL.** We thank the reviewer for pointing us to this interesting line of research and we will include an additional paragraph about programmatic RL in the Related Work Section. We also added some other works (with [6] and [7] specifically focusing on model-based programmatic RL) and some references for modern works using LLMs to generate programmatic policies.\n\n**7. Clarity: offline setup.** Thanks for the suggestion, we will clarify our setup.\n\n**References:**\n\n[1]: Pathak, Deepak, et al. \"Curiosity-driven exploration by self-supervised prediction.\" International conference on machine learning. PMLR, 2017.\n\n[2]: Savinov, Nikolay, et al. \"Episodic Curiosity through Reachability.\" International Conference on Learning Representations (2019).\n\n[3]: Ecoffet, Adrien, et al. \"First return, then explore.\" Nature 590.7847 (2021): 580-586.\n\n[4]: Fu, Justin, et al. \"D4rl: Datasets for deep data-driven reinforcement learning.\" arXiv preprint arXiv:2004.07219 (2020).\n\n[5]: Hafner, Danijar, et al. \"Dream to Control: Learning Behaviors by Latent Imagination.\" International Conference on Learning Representations.\n\n[6]: Azad, Abdus Salam, et al. \"Scenic4rl: Programmatic modeling and generation of reinforcement learning environments.\" arXiv preprint arXiv:2106.10365 (2021).\n\n[7]: Tsividis, Pedro A., et al. \"Human-level reinforcement learning through theory-based modeling, exploration, and planning.\" arXiv preprint arXiv:2107.12544 (2021).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank all the reviewers for engaging with our work and providing their precious feedback. In addition to our individual answers, we include the following material:

**Example of generated Code World Model.** Reviewers u7M2 and 3vW3 asked for examples of Code World Models (CWMs), being interested in the generated world models for continuous environments and in possible overlaps with the online implementations. In Figure 1 of the attached PDF we report the CWM generated by our method for the Ant-v4 environment. For reference, the official implementation of the environment can be found at the official Gymnasium GitHub repository of the Farama Foundation (at gymnasium.envs.mujoco.Ant_v4).

**Comparison with RL baselines.** Reviewer u7M2 asked for reference a comparison with SOTA RL and Oracle baselines, and reporting the raw returns, while reviewer 49z7 was specifically interested in a comparison with offline RL. We report in the following tables the average reward obtained over 10 episodes for a random policy, a SOTA method Conservative Q-Learning (CQL) [1], 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 (Oracle). CQL is a SOTA offline RL method, trained with 10 epochs for 100 steps per epoch (1000 total) using the *same* dataset 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, we started to observe severe overfitting for CQL with more training steps.

| **Environment (Discrete)**              | **Random** | **CQL**  | **GIF-MCTS (ours) Llama 3** | **GIF-MCTS (ours) GPT-4** | **Oracle** |
|------------------------------|------------|----------|-----------------------------|---------------------------|------------|
| Blackjack-v1              |      0   |    -0.3 |      -0.6 |    **-0.1** |      1   |
| CliffWalking-v0           |  -1169.2 |   N/A   |     **-90.2** |  -100   |   -100   |
| Taxi-v3                   |   -798.5 |  -740   |    **-353.9** |  -408.8 |   -124.5 |
| CartPole-v1               |     24.4 |   **317.6** |     277.4 |   310.4 |    494   |
| MountainCar-v0            |   **-200**   |  **-200**   |    **-200**   |  **-200**   |   -200   |
| Acrobot-v1                |   -500   |  **-295**   |    -500   |  -494.2 |   -500   |


---

| **Environment (Continuous)**              | **Random** | **CQL**  | **GIF-MCTS (ours) Llama 3** | **GIF-MCTS (ours) GPT-4** | **Oracle** |
|------------------------------|------------|----------|-----------------------------|---------------------------|------------|
| Pendulum-v1               |  -1122.8 | -1218.2 |   -1232.2 |  **-739.8** |   -373.6 |
| Reacher-v4                |    -43.7 |   -11.5 |      **-9.2** |   -11.2 |     -6.8 |
| Pusher-v4                 |   -149.9 |   **-52.4** |     -61.1 |   -63.3 |    -30.3 |
| InvertedPendulum-v4       |      8.3 |    **66.7** |      13.1 |    10.9 |     42.5 |
| InvertedDoublePendulum-v4 |     49   |   **164**   |      60   |    53.4 |    241.6 |
| HalfCheetah-v4            |   -304.5 |    **-1.3** |    -150.3 |   -22.8 |    893.3 |
| Hopper-v4                 |     32.2 |   **137.4** |      62.6 |    23.3 |    229.1 |
| Swimmer-v4                |     -5.9 |    **28.4** |      -2.7 |     8.1 |    317.8 |
| Walker2d-v4               |     -0   |   **278**   |      22.3 |    11.5 |    334.7 |
| Ant-v4                    |    -33.2 |   **998**   |     867.7 |   896.8 |   1304.7 |
| Humanoid-v4               |    139.4 |   **393.3** |     N/A   |   162.3 |   1860.7 |
| HumanoidStandup-v4        |  33240.2 | **51045.7** |     N/A   | 29405.9 | 138076   |

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

In general, planning with the generated CWMs works best on discrete environments, where our original results also indicate that the CWMs are of a higher quality. However, CWMs also reach competitive results in some complex physics tasks, such as Pendulum-v1, Reacher-v4 and to a lesser extent Ant-v4, Pusher-v4 and HalfCheetah-v4, without direct access to the original physics simulator.

**References:**

[1]: Kumar, Aviral, et al. "Conservative q-learning for offline reinforcement learning." Advances in Neural Information Processing Systems 33 (2020): 1179-1191.

Meta-review: 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.
Decision: Accept (poster)