WorldCoder, a Model-Based LLM Agent: Building World Models by Writing Code and Interacting with the Environment

Thank you for the thoughtful review. Please see below for our responses.

What if integrating the proposed approach, in particular, \phi_1 & \phi_2 to deep RL, if feasible, maybe together with an LLM for language understanding?

Interesting idea! One way we've been thinking about this is to frame the integration as hierarchical RL: have conventional deep RL handle low level control on small timescales, and our method act as a high level controller.

rely on an imperfect LLM and an imperfect method for program synthesis, and there is likely no way to achieve perfection Is there a way to measure the quality of the learned model, e.g., comparing with the true model, rather than relying on performance study in the paper? There should be such a perfect model for Sokoban and Minigrid.

We learn a perfect model on Sokoban and sometimes learn imperfect models on Minigrid which nonetheless achieve high reward, because perfectly modeling the transition function is not necessary to solve all tasks, especially when not all dynamics need be used to achieve the goal.

But the real world is not as clean and tidy as videogames, so applying WorldCoder to something like a robot would require probabilistic programs, a direction we mention in the paper, and which we are excited to explore. We hope that by getting this paper out we can pursue these next steps. This relates to another comment:

Under Line 97 Conditions \phi_1 fit data and \phi_2 optimization What if there are noisy data points? What if data points are contradicts with each other?

We currently consider deterministic environments without noise, but as line 311 of the paper says, we are pursuing future work that overcomes this limitation.

it is not clear the effect of exponential search space on the polynomial sample complexity

We prove good sample complexity but not good computational complexity: in the worst case both planning and program synthesis are exponential time. We rely on the LLM to act as a good heuristic guide to make program synthesis not exponential in practice.

it is not clear if the evaluation is fair: 1) are the compared algorithms the best? 2) how about comparing problems deep RL is good at?

We're happy to discuss any specific baselines you'd like performed, but we think any deep RL approaches are going to require many thousands or even millions of episodes, although they likely will asymptote to 100% success.

The purpose of our work is not to show that Deep RL is always inferior, and in fact there are many problems that it can handle that our method would flop on. In the limitations section we mention that we require low-dimensional symbolic states and a deterministic environment, and comparing on any such problems would make deep RL look better. But given that our method is new and relatively unusual, we wanted to study interesting problems where it has advantages.

How about MuZero?

We think it would perform very similarly to Dreamer_V3, which we compare against.

“preferring an optimistic [world model] guarantees that the agent can at least start making progress toward the goal"... Is this correct? Consider the shortest path problem. At some node, if an agent selects a neighbour / node, the agent may never achieve the optimal solution.

We assume an episodic MDP, so even if the agent gets stuck on the current episode, it always has another try on the next episode---and if it was too optimistic, then it's guaranteed to not make the same mistake again, under our algorithm.

"Before ever receiving reward, \phi_2 forces the agent to invent a reward function." Why and how to "invent a reward function"? Why should it work?

With no reward function, or a reward function that always returns zero, it is impossible to satisfy $\phi_2$, because that optimism constraint says that it should be possible to achieve positive reward from the initial state. In order to satisfy that constraint, there must exist a reward function such that a planner predicts reward can actually be achieved, using the current learned transition function. So the LLM is free to invent a reward function, subject to constraints in $\phi_2$

How about the diameter and the logical dimensionality for problems like chess and go? Will they grow exponentially as the search space grows exponentially?

The diameter is at most twice the maximum length of a game. The logical dimensionality is approximately the log of the number of semantically distinct transition functions the language model can produce consistent with the replay buffer.

"assumes similar goals have similar reward structures" This may not be true. And how to define similarity?

Indeed, this assumption is merely a heuristic. For example, if the agent has learned that the goal "pick up the ball" gives a reward of 10 for picking up the ball, then this assumption would say "pick up the balloon" should also give a reward of 10 for picking up a balloon. Similarity is implicitly defined by the LLM's in-context learning.

Last, about the perceived contradiction between these statements:

L138: "After learning how the world works, our agent should be able to receive natural language instructions and begin following them immediately without further learning from environmental interactions"

L149: "If the LLM makes a mistake in predicting the reward function, then the agent can recover by subsequent rounds of program synthesis, which update $R(c)$ based on interactions with the environment."

These statements say that the agent should be able to receive a new command and conjecture a plausible reward function, but also modify that reward function if it turns out to be incorrect, based on experience it collects later on down the line.

what is the meaning of $\forall i \in [l]$

It means all integers $i$ where $1\leq i\leq l$. We'll revise to clarify this and other notation.

Thanks for your input and please let us know if you have further questions.

The proposed method relies on an imperfect LLM and an imperfect method for program synthesis... the learned world model is not guaranteed to be correct... Formal verification is required... The proposed method may... achieve good experimental performance. However, without the guarantee of correctness of code / world model, we may not know what may happen.

This is a fundamental issue for model-based RL and machine learning broadly, especially when neural networks are involved. However, representing knowledge as a program opens up unique opportunities for formal methods (see [1,2] for some of our favorite examples in RL+Programs), which we'd be excited to explore in the future. We'll revise to mention these factors, but our opinion is this criticism is not unique to our work and orthogonal to our contribution: An identical critique can be levied against neural network world models. Our method has the advantage that, in principle, there is hope of deploying formal methods in the future, owing to the symbolic nature of its learned knowledge.

(If you're curious, we have ideas about how to formally verify safety properties of world models, which is an open problem, because past work [1-2] instead verifies policies. We can continue the discussion in that direction if you'd like. Might be interesting to get your feedback.)

Is the proposed world model learned with a symbolic approach equivalent to a transition model?

The world model consists of a transition function $T:S\times A\to S$ and reward model $R:S\times A\times S\to \mathbb{R}$. The former is a transition model.

Moreover, is the world model tabular? It seems so

The world model is not tabular. It does not memorize every possible case in a large table. The world model is a program which can generalize over states and actions.

[1] An Inductive Synthesis Framework for Verifiable Reinforcement Learning. Zhu, Xiaog, Magill, Jagannathan. PLDI 2019.

[2] Verifiable Reinforcement Learning via Policy Extraction. Bastani, Pu, Solar-Lezama. NeurIPS 2018.

Thanks for continuing to engage with us. We think there's a misunderstanding we'd like to clarify.

In the large-compute limit, our approach does converge to a program that perfectly fits the training examples. This is because the LLM always puts nonzero probability on at least one such program (at nonzero temperature), and so if code repair runs arbitrarily many times, it will perfectly fit the training examples. In practice we limit the compute budget, so might terminate before finding an optimal program, but the same is true for the other approaches you mention (FunSearch; AlphaGeometry). Our main theorem we prove only holds in the asymptotic compute regime, because it assumes we can perfectly fit the training examples. The theorem is nontrivial because it is not fundamentally about fitting the training examples, but instead about RL sample-complexity, i.e. achieving high reward by successfully exploring the environment in relatively few steps.

We hope that this large-compute limit clarifies the relationship between imperfect LLMs and our method, as well as FunSearch and AlphaGeometry.

I understand that improving LLMs may be out of scope for you. But your work is based on imperfect LLMs. This is an issue not only to your submission, but to the LLM-based method.

Indeed, we agree that it's not unique to our submission and likely out of scope. We just want to emphasize that our method, like others such as FunSearch and AlphaGeometry, are only asymptotically complete in the large-compute limit---although they are always sound because they can decline to output a program if it fails to verify/satisfy the examples: So we always have soundness but only get completeness asymptotically. Distinguishing soundness and completeness is important here.

Thank you for the thoughtful review. We really think we can improve the paper to address your main points. Please see below.

GPT-4, the model used in the experiments, certainly was trained on the code implementing the models of the domains used in the experiments. How much of what we see is from the general knowledge of the LLM versus memorizing the implementation of the models used in the experiments?

This is an excellent point and it deserves careful consideration:

1. In the attached PDF (general response) we show what happens when the agent is trained on a new variation of Sokoban we designed this week, and which therefore could not be in pretraining data. The new variant has warp gates/teleportation portals allowing the agent to instantly move between a pair of locations. The logic of the teleport gate is subtle, because it is deactivated whenever one of the portals is blocked by a box. The agent nonetheless learns how to use the warp gates, and solves problems faster by exploiting them to take shortcuts, as shown in the PDF.

2. When solving AlfWorld, it makes many edits to its world knowledge to correct problems with its prior understanding of this environment (Fig 5 bottom), for example learning that only one thing can be picked up at a time, or that objects cannot be removed from closed receptacles. When solving Sokoban, it has to make analogous fixes, such as understanding that boxes can block the movement of other boxes, or that there is a per-step reward of -0.1. If it were simply retrieving the answer from training data, we would not expect these mistakes. Instead, it seems to be retrieving general knowledge about how videogames work, and making lots of small changes to the transition function based on the details of the data in the replay buffer.

How much of the code of the models for the three domains is available online?

1. AlfWorld is publicly available as a zip file containing PDDL configs. It is therefore unlikely that GPT4 was trained on its implementation, both because the raw data is zipped, and because it is not in a common source code file extension (PDDL looks like a weird lisp with an unusual file extension).
2. Minigrid (released 2022) and Sokoban are probably in the pretraining data. Our new Sokoban result is not in the pretraining data.

We will revise the paper to put more emphasis on the AlfWorld results and the new Sokoban results, because those are the ones that we are the most sure are not affected by data contamination.

The paragraph about React (Section 3) attempts to deal with this issue, but it misses the point

Upon reading your argument, we agree with you that the ReAct baseline does not address these concerns. The revision will swap that discussion for the discussion above.

How would WorldCoder perform on a domain that is different from everything that is out there?

We've designed a new Sokoban domain to answer that question. The best way of answering that question would be to design completely new games for the agent that have never been on the internet, or barring that, to use domains like AlfWorld whose environment implementation is very unlikely to have been trained on.

Data contamination is the only reason why I am giving a "poor" score to soundness

We hope that you can reconsider your score in light of the new experiment and the above discussion.

Thank you for the unusually helpful review. Please let us know if you have any further questions.

Thank you for the thoughtful review, and for your support. Please see below for our responses.

Task generalization achieved by prior knowledge: Env generalization with uncertain/incomplete prior knowledge of LLMs is not properly addressed. E.g., the dynamics behavior of grid worlds is obviously known by GPT-4

The system leverages its prior knowledge, but is not bound by it. Here are two examples, including new results in the attached PDF:

1. In the attached PDF we show what happens when the agent is trained on a new variation of Sokoban we designed that has warp gates/teleportation portals allowing the agent to instantly move between a pair of locations, as long as one of them is not blocked by a box. The agent learns how to use the warp gates and can solve problems faster by exploiting them, as shown.
2. When solving AlfWorld, it makes many edits to its world knowledge to correct problems with its prior (Fig 5 bottom), for example learning that only one thing can be picked up at a time, or that objects cannot be removed from closed receptacles.

the paper's contribution boils down to prompt-engineering with output constraints

The contribution includes a new logical/formal framing of optimism under uncertainty, its theoretical analysis, and empirically showing this framing works in practice and helps accelerate learning. We further contribute an LLM agent with better big-$\mathcal{O}$ complexity in terms of amortized # API calls per action, a consequence of being model-based instead of model-free.

Scalability is not properly addressed. E.g. looking at the prompts in Appendix for MiniGrid: the states are encoded in a rigerous manner with a lot of designed structure and tokens. How does the performance behave for large state space encodings. Is the LLM still able to create adequate transition functions?

We scale to complicated symbolic environments: AlfWorld has a very large state space with typically around 60 objects and around 20 predicates, for a states space size of $2^{60\times 20}$. Appendix E lists an example learned world model for this setting. We do not consider nonsymbolic environments, although it is a next step that we are actively working on.

Experiments on grids or highly abstracted env states: Is this an analogy to Kahnemann's System 2 or a current limitation?

It is indeed analogous to System-2 abstract reasoning, or task planning in robotics more specifically

The authors do not address input data integrity: The method relies on perfect(?) input sample quality… The generated Python programs can only model highly abstracted MDPs, since the input samples are a text-based, tokenized abstractions… Experiments on grids or highly abstracted env states => How to ensure hiqh quality input data in the general case (e.g. pixels)?

We are working on extensions that learn probabilistic programs (to handle noise in the data), and which use multimodal models to define the basic symbols/tokens (so we can learn from e.g. pixels), which would allow handling, for example, robotics environments. We consider these extensions sufficiently different that they should go into a different paper, while WorldCoder_v1 is more concerned with the fundamentals of the problem and its conceptual/mathematical solution.

Thank you for the comprehensive review and for being supportive of the paper. Please below for answers to your questions.

The method lacks generality, needing a hand-defined curriculum in minigrid and a handcrafted language-based reward bonus on AlfWorld;

We think this might be a misunderstanding:

1. The method doesn't need a curriculum (see Fig 4C), but can take advantage of good curriculum if one is provided (Fig 4B)
2. The reward bonus does not contain any AlfWorld-specific features: It relies on generic text features. It can apply to other domains where the objects and actions have names in natural language.

the environments considered here… are pretty simple

We started with simple environments because our pilot experiments showed that new basic principles needed to be developed (optimism under uncertainty applied to programs). We are excited to explore richer application domains such as those you mention, and we hope that by getting this work published we can soon begin that follow-up work, which requires addressing the limitations outlined in the submission.

Please note that in the general response and in the attached PDF we have designed a more complicated Sokoban environment---with warp gates/teleportation portals--and tested our approach on it as part of responding to a different reviewer.

a more principled baseline for tasks like sokoban would be Tree of Thoughts (ToT)

We will work on this and hopefully have something to share during the discussion period next week. However the asymptotic analysis would be the same as ReAct, with cost growing linearly as a function of the number of episodes.

The method remains slightly complex and nontrivial to implement... easily available code an important part

Code and data will be released upon publication.

if policies are parametrized by the goal, which programming languages allow you to do: go_to_object(goal) would certainly transfer to many different objects

Good point! We will mention that.

The paper lacks baselines in minigrid and alfworld

Please see Appendix Fig 6 for minigrid Deep RL baselines. We agree though that more baselines on those domains would help the paper, but we also believe strongly that deep RL baselines would always (1) eventually solve every task but (2) take millions of episodes. On the other hand LLM agent baselines would (1) rapidly achieve sub-100% performance but (2) never asymptote to 100%, due to hallucinations and (3) have API cost linear in the number of episodes. We observed these characteristics for Sokoban (deep RL/ReAct) and Minigrid (deep RL), and given the high cost of running experiments, felt it was low priority to run further simulations.

Could a version of WorldCoder work for hard linguistic/reasoning tasks

Thanks for suggesting these directions. WorldCoder could be used for the tasks you mentioned if the world model can delegate parts of its computation to a "fuzzy" reasoning module like an LLM, similar to ViperGPT. This would be a fascinating direction to explore.

How do you ensure the self-repair program synthesis step does not break previously satisfied conditions

We keep all the old experience data in a replay buffer to ensure that we do not break the world model when updating it for new data.

Thanks again for your support and please let us know if you have further questions.