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

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.

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]). In 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.

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). If 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. We 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.

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.

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

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.

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. Overall, there is a balance between CQL and CWMs, with CWMs being more suited to discrete tasks and CQL outperforming CWMs in complex physics tasks. 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.

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.

7. Clarity: offline setup. Thanks for the suggestion, we will clarify our setup.

References:

[1]: Pathak, Deepak, et al. "Curiosity-driven exploration by self-supervised prediction." International conference on machine learning. PMLR, 2017.

[2]: Savinov, Nikolay, et al. "Episodic Curiosity through Reachability." International Conference on Learning Representations (2019).

[3]: Ecoffet, Adrien, et al. "First return, then explore." Nature 590.7847 (2021): 580-586.

[4]: Fu, Justin, et al. "D4rl: Datasets for deep data-driven reinforcement learning." arXiv preprint arXiv:2004.07219 (2020).

[5]: Hafner, Danijar, et al. "Dream to Control: Learning Behaviors by Latent Imagination." International Conference on Learning Representations.

[6]: Azad, Abdus Salam, et al. "Scenic4rl: Programmatic modeling and generation of reinforcement learning environments." arXiv preprint arXiv:2106.10365 (2021).

[7]: Tsividis, Pedro A., et al. "Human-level reinforcement learning through theory-based modeling, exploration, and planning." arXiv preprint arXiv:2107.12544 (2021).

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.

1. Clarity.

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). Please 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.

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. We also realised that in the main text we used $l$ rather than $L$ and we will fix this.

2. Claims. We never meant to present imprecise claims and we thank the reviewer for pointing out the need for better wording of them.

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

We 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).

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

We will change to "Calling code instead of LLMs for planning has potential to be more precise, reliable, interpretable, and extremely efficient."

Addressing the mentioned properties one by one:

• Precise: code can natively access a calculator and as such is numerically precise.
• Reliable: a program passing all unit tests can be deemed more reliable than a black-box prediction from an LLM.
• 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.
• 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.

3. Data contamination.

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.

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.

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.

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.

We will add this discussion to the Appendix of the paper.

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:

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

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.

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.

1. Examples of LLM-generated environments:

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

In 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. In 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. We 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.

2. Comparison with RL:

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.

We 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. Overall, there is a balance between CQL and CWMs, with CWMs being more suited to discrete tasks and CQL outperforming CWMs in complex physics tasks. 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.

3. LLM sample budget:

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.

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

References:

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

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.

1. Pre-training:

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.

and

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?

We 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). Specifically, 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. Separately, 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.

2. Predicting transitions:

Specifically, does GIF-MCTS need to predict all transitions well in the offline dataset to perform well on a task?

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

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.

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

3. Novelty:

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.

We 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. Hence, while GIF-MCTS can be seen as novel combination of agentic behaviour and tree search, our novelty goes beyond that.

4. Comparison with WorldCoder:

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

Good point, thank you for raising it. We believe that GIF-MCTS outperforms WorldCoder because it considers a more diverse set of programs. The 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. A 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. We will add this explanation to the Discussion section in the main text.

Thank you for the responses and clarifications. The clarification on pre-training does clear up my confusion. I thought the training data set $D$ (line 148) is used to train a language model to generate world models. That was a misunderstanding.

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

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

If these points are correct, it would be helpful to clarify them in the paper. Initially, I thought we needed to learn a world model for APPS, which confused me.

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

I have raised the score to 6 (Weak Accept).