WorldCoder, a Model-Based LLM Agent: Building World Models by Writing Code and Interacting with the Environment
Authors: Hao Tang
Darren Yan Key
Kevin Ellis
Venue: NeurIPS 2024 poster (2024)
Keywords: program synthesis
LLM
reinforcement learning
Primary area: reinforcement_learning

Abstract: We give a model-based agent that builds a Python program representing its knowledge of the world based on its interactions with the environment. The world model tries to explain its interactions, while also being optimistic about what reward it can achieve. We define this optimism as a logical constraint between a program and a planner. We study our agent on gridworlds, and on task planning, finding our approach is more sample-efficient compared to deep RL, more compute-efficient compared to ReAct-style agents, and that it can transfer its knowledge across environments by editing its code.

WorldCoder, a Model-Based LLM Agent: Building World Models by Writing Code and Interacting with the Environment
Hao Tang
Darren Key
Kevin Ellis Cornell University
Cornell University haotang@cs.cornell.cco
Cornell University dyk34Gcornell.edu
kellis@cornell.edu
Abstract
We give model-based agent that builds Python program representing its knowl- edge of the world based on its interactions with the environment. The world model tries to explain its interactions, while also being optimistic about what reward it can achieve. We define this optimism as a logical constraint between a program and planner. We study agent gridworlds, and on task planning, finding our approach is more sample-efficient compared to deep RL, more compute-efficient compared ReAct-style agents, and that it transfer its knowledge across environments by editing its code.
Introduction
new device or play new game. Given the right prior knowledge, Consider yourself learning interactions, most people can acquire basic knowledge how many together with relatively few devices games work. This knowledge help achieve novel goals, be transferred to similar devices or games, and can be communicated symbolically to other humans. How could an Al system similarly acquire, transfer, and communicate its knowledge of how things work? We cast this as learning a world model: mapping, called the transition function, that predicts the next state of affairs, given current state and action 63 55 26 Our proposed solution is architecture that synthesizes a Python program to model its past experiences with the world, effectively learning the transition function, and which takes actions by planning using that world model.
In theory, world models have many advantages: they allow reasoning in radically new situations by spending more time planning different actions, and accomplishing novel goals by just changing the reward function. Representing world knowledge as code, and generating from LLMs, brings other ad adnatages. It allows prior world knowledge embedded in the LLM to inform code generation, allows sample-efficiee transfer across tasks by reusing pieces old programs, and allows auditing the system's knowledge, because programming languages designed to be cuman-interpretable.
But there are also steep engineering challenges. Learning the world model now requires a combina- torial search over programs to find transition function that explains the agent's 's past experiences. Obtaining those experiences the first place requires efficient exploration, which is difficult long horizon tasks with sparse reward. To address the challenge of efficient exploration, we introduce new learning objective that prefers world models which planner thinks lead to rewarding states, particularly when the agent is uncertain to where the rewards problem, we show how curriculum learning can allow transfer of knowledge across environments, making combinatorial program synthesis more tractable by reusing successful pieces of old programs.
To address the program search
Fig. T diagrams the resulting architecture, which we cast as model-based reinforcement learning (MB RL). In Fig. 2 we position this work relative to deep RL as well as LLM agents. In contrast to deep
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

state, reward.goal
world
world nodel.py
planner
def update(state, action, goal):
LLM
replay buffer
return new_star statt reward
action
Figure 1: Overall agent architecture. The agent also inputs goal in natural language
Sample
World model representation
Priors
Inputs
LLM calls/task
complexity
Deep MB RL
Low
High
Neural, learned
High-dim
LLM Agents
High
Zero"
Neural, fixed
Symbolic
(T) O(1)
Ours
High
Low
Symbolie, learned
Symbolic
Qualitative comparison of our method against deep model- basel RL and LLM agents (ReAct, RAP, Yao al. (77). Hao al. [29], Zhao al. [78], Liu [4T).. Sample complexity refers to the number of environment interactions needed to learn a world model ("LLM agents do not update their world model). LLM calls/task the number of LLM calls needed solve new task in fixed environment, amortized over many such tasks, as function of the maximum episode length T. Asymptotically, after learning world model, our method can accomplish new tasks by only at most one LLM query to update the reward function.
Figure 2:
RL 55 i.a.]. we view the world model something that should be rapidly learnable and transferable. Central to our work is particular claim about how an LLM should relate to world models. In our setup, the LLM does not simulate the world, but instead builds simulation of the world. This should be contrasted with LLM agents such as ReAct (77] where the LLM plays the role a world model by reasoning about different actions and their consequences. We also do not expect the LLM to perform planning, which they are known to struggle with (68]. Instead, we require the LLM to possess fuzzy prior knowledge of how the world might be, together with the programming skill code and debug transition function. Our use of LLMs is closest to [2 74 79,, which generate planning productions, which be closer to 13 67], which learn world models in domain-specific programming languages. We further show how use Turing-comples languages like Python-whit we believe important for general- purpose learners-an and we also study efficient transfer learning and exploration strategies.
particular kind of world model. Overall though, our problem statement is
We make the following three contributions:
1. An architecture, which we call WorldCoder, for learning world models as code. The architecture supports learning that is more transferable, interpretable, and dramatically more sample-efficient compared to deep RL, and also more compute-effiiiett than prior LLM agents.
2. A new learning objective for program-structure world models that favors optimism in the face of (world model) uncertainty (Section 2.2). We show theoretically and empirically that this learning objective generates goal-driven exploratory behavior, which can reduce by orders of magnitude the number of environment interactions needed obtain reward
An analysis of different forms transfer world models, finding that code can be quickly adapted new world dynamics within grid-world and robot task planning domains.
Methods
2.1 Problem statement and core representations
We start with the standard MDP formalism but modify it in three important ways. First, we assume a goal is given in natural language. The goal could be something specific such as "pickup the ball", or underspecified, such "maximize reward." Second, we restrict ourselves to deterministic environments, and assume that environment dynamies are fixed across goals. Third, we assume an episodic MDP with the current episode terminating upon reaching the goal.

---

We formalize this as Contextual Markov Decision Process (CMDP: Hallak et al. 128)), which tuple (C, S, A, M) where C' is a set of contexts (i.e. goals), S is a set of states, actions, and M is a function mapping context e. E C' to Markov Decision Process (MDP). The context-conditioned MDP, M(c), is a tuple (S, A, R 7) with transition function T S A + S, discount factor y, and reward function R S x A A S + R. The transition function does not depend on the context. The objective is to select actions to maximize cumulative discounted future reward, which is given by Lt T'T,, where is the reward received timesteps into the future. Termination is modeled by assuming there is a special absorbing state.
set of
State representation. Motivated by robot task planning 20 55 [57 we represent states as sets of objects, each with a string-valued field name, fields and y for its position, and additional fields depending on the object type. For example, if name="door then there are two Boolean fields for if the door is open/closed and locked/unlocked. This can be seen as an Object Oriented MDP 15]
Representing world models as code. The agent uses Python code to model the transition and reward functions, Mathematically we think of this Python code as tuple (T, It) of transition function : x A - S and reward model R : C - (S x A x S - R x 1)) Note again that the reward depends on the context, and returns an extra Boolean indicating whether the goal has been reached, in which case the current episode terminates. Both functions are implemented separate Python subroutines, which encourages disentangling the dynamics from the reward.
2.2 The world model learning problem
What objectives and constraints should the world model satisfy? Clearly, the learned world model should explain the observed data by correctly predicting the observed state transitions and rewards. This is standard training objective within model-based RL 62]
One less obvious learning objective is that the world model should suffice to plan to the goal. Given two world models, both consistent with the observed data, the agent should prefer model which implies it is possible to get positive reward, effectively being optimistic in the face uncertainty about world dynamics. Assuming the low-data regime, there will be multiple world models consistent with the data: preferring an optimistic one guarantees that the agent can at least start making progress toward the goal, even if it later has to update its beliefs because they turned out to be too optimistic.
Concretely, the agent collects dataset D of past environment interactions, each formatted as a tuple (s, r',' c, d) of current state next state s', action and reward in context with d indicating if the episode ended upon reaching s'. (The variable d should be read as "done.") The agent also stores the initial state so and context of each episode, that it prefer world models that can reach the goal from an initial state. The learning problem is to construct Python programs implementing a transition function π and reward model satisfying constraints Φ1 and Φ2, defined below:
Φ (D.T.R) = V(s.a,r, E D: (f..)
fit data
(1)
where (I, (T, + (,, ir)) =' Aice) = (r.d)
2 Sp. c, T. S..RR. (so.c.T.
optimism
Ea,, $1, 02 82, ags St
Vi E [A] f(s a.) = Ar > 0: ((c)((s-1. az, sz) = (r, 1)
where so. e is the initial state/context of the episode, and the turnstile (I-) should be read as meaning that a given program entails a given logical constraint or predicts given replayed experience.
Constructing a world model satisfying Φ1 A d2 is a program synthesis problem. We solve it by prompting an LLM with random subset of DD and asking it to propose a candidate program. If the resulting program does not fit all the data, the LLM is backprompted with members of ' inconsistent with its latest program, and asked to debug its code, following 77 47
Sec. 2.5 details this process.
2.3 Understanding optimism under uncertainty
Optimism under uncertainty is pervasive principle throughout learned decision-making 63 in-
cluding model-based RI. [75 40) We contribute new instantiation of that principle as a logical constraint between program and a planner, and which we show can be made compatible with
3

---

LLM-guided program generation. Below we give formal and intuitive guides to how this logical constraint plays out in our context.
Exploration guided by goal-driven behavior. Before ever receiving reward, Φ forces the agent to invent reward function. This reward function will be something that the agent believes it could achieve, but has not already achieved: roughly, something within its 'zone of proximal development" [71]. After probing zone of proximal development, the agent updates its model based on the experience of actually trying to achieve a new goal (these model updates come from enforcing @1). This exploration continues until finishing an episode with positive reward.
For example, suppose the context (goal) is "open the door with the key", and the agent has not yet achieved any reward, nor has ever even found the key. Then Φ2 would (for example) prefer R's which reward touching the key. Even if touching the key does not actually give reward, the agent can make good progress by pretending it does.
Formal analysis of exploratory behavior. to be discovered in time that polynomial several key quantities, defined below. contrast to prior work model-basee optimism under uncertainty 775
We prove that good-enough world model is guaranteed
31 40) our analysis is not probabilistic, and instead leverages combinatorial and logical properties of the solution space, because our models represented as discrete programs.
3 60
Definition 2.1. A dataset D is mutually independent w.r.t a solution space M, denoted, D M, iff for all data points in d E 'D there exists a solution in M which explains all the data except for d:
Vd 3m E M, m V d A (Vd' E D {d), m d').
Definition 2.2. The logical dimensionality of model space M, written KILM, is the size of the biggest dataset that is mutually independent of M:
KLL = max [P] VD where D TL M}
Within this paper the model class corresponds to the cross product of possible transition functions and possible reward functions, denoted and R. respectively, so M = T x R.
Definition 2.3. The diameter of deterministic episodic MDP, written Ds.A.T, is the maximum shortest path length between all pairs of states:
DS.A.T
max
min
len(p)
s,''ES ES pa path from to
Putting these together, Appendix A shows that we can learn in time polynomial w.r.t. the diameter (a property of the true MDP) and the logical dimensionality (a property of the model space):
Theorem 2.4. Assume an episodic MDP (S, A, T, R, 7). Assume an agent acting according to an optimal planner operating over world model (T, It) E T x R satisfying Φi A 42, and that the true MDP is in the agenf's 's model class: T E T and R E R. Then the maximum number of actions needed to achieve the goal (exit with positive reward) is Ds.A,T X (KTXR + +1).
Intuitively, the polynomial sample complexity is possible because the agent systematically finds independent data points, which serve counter-examples the agent's current model. At most КıTTæR such counterexamples are needed, each requiring at most DS.A.T actions.
Following natural-language instructions. 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. Mathematically, given learned program that implements reward functions for previously seen goals, together with a new goal e where c &f domain R), the agent should update its reward model in order to cover
This zert generalization to new goals occurs as a consequence of enforcing optimism under uncertainty ( in Eq. 2 Upon observing a new context the constraint & is violated. This triggers debugging R so that it covers c, subject to the constraint that f((c) allows reaching a goal state.
Given the importance of cnatruction-follwinng abilities, we use differrntt prompt for trying to enforce 62 upon encountering a new goal (Appendix Sec. F.5). This prompt is based on retrieving previously learned reward functions. This retrieval strategy assumes similar goals have similar reward structures, so that the LLM can generalize from old reward functions in â when generating the new R(c). 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 interactions with the environment.

---

Algorithm WorldCoder Agent Architecture
Hyperparam:
random exploration probability (default to 5%)
Hyperparam: MINDATASIZE, min actions before learning begins (default to 10)
D, D,D.c +- 2,
& replay buffer. Dac holds initial states and contexts needed for de init empty world model
T null,, null
loop forever through episodes:
EPISODEGOALL)
get context (goal) record initial state Replay buffer of initial conditions for d2
80 t CURRENTSTATE()
Da < Dac {(50, c))
loop until episode ends:
X-CURRENTSTATEE
not Φ1 A @ and [D] 2 MINDATASIZE then
T, R SYNTHESIZE(T,, D.
>> Section 2.5
with probability do --RANDOMACTIONO
> e-greedy explore
else
PLAN(S (((() $r., d + ENV.STEP()) + DU (,
>> Value Iteration take action in state b record experience
2.4 Overall architecture
Ultimately our world models exist to serve downstream decision-mann maaking: and in turn, taking actions serves provide data for learning better world models. There many architectures for combining acting with world-model learning, such as via planning (55), training policies in simulation [26], or hybrid approaches [62]. We use the architecture shown in Algorithm performs random exploration to initialize a dataset of environmental interactions; it updates its world model using the program synthesis algorithm Sec. 2.5] and, past its initial exploration phase, performs planning. Different planners are possible, and we use depth- value iteration (in simple domains) and MCTS (for more complex domains).
[ At high level, it initially
2.5 Program Synthesis via Refinement
One approach to program synthesis is to have an LLM iteratively improve and debug its initial outputs 7 [47), which call refinement. For world models as programs, this means prompting the LLM with an erroncous program it previously generated, together with state-action transitions that program cannot explain, and prompting to fix its code. This process repeats until the program satisfies Φ1 A Φ2. Refinement can be very successful when the target program has many corner-cased each of which can be inferred from few examples, because the LLM can incrementally grow the program driven by each failed test case. This is exactly the case for world models, which might need to handle wide range of objects and their interactions, but typically don't 't demand intricate algorithms. Refinement also allows computationally efficient transfer between environments, because new world model can be built by refining an old one, instead of programming it from scratch.
We use a concurrently developed algorithm called REx to determine which program to refine next (65 REx prioritizes refining programs that appear more promising, which in the context of world model learning, we instantiate by defining by a heuristic h, which measures progress towards satisfying Φ1 62
x] 1/ (D.T.R) (D) + 1Decl
E.o.cED.e 1
[022000 c, T, RR)
bT, R.
The above heuristic computes the fraction of the replay buffer which is consistent with Φ1, and the fraction which is consistent with Φ2. (We incentivize satisfying Φ1 first by multiplying the average accuracy on the optimism objective by indicator [Ø) (D, T, RR)].) REx also prioritizes refining programs that have not been refined very many times, and uses a bandit formulation to balance (1) exploring programs that have not been refined very many times against (2) exploiting by refining programs that are better according to h. See (65) for more details.
5

---

We implement this refinement process using GPT-4 because recent work 47 finds it is the strongest model for repairing and improving code (as opposed just generating code from scratch), Using this technique and this LLM we can generate world models with 250+ lines of code.
Experimental Results
We study our system in three environments, Sokoban, Minigid, and AlfWorld, with the goal of understanding the sample efficiency and computational efficiency of the learner, especially when transferring knowledge across environments, as well as the impact of optimism under uncertainty.
Sokoban is puzzle-solving task where the agent pushes boxes around 2d world, with the goal of
pushing every box onto a target (Fig. (A). Solving hard Sokoban levels challenging planning
[8] (4). Unlike these other works, our emphasis not on solving the hardest Sokoban levels. Instead, we wish show that our agent can rapidly achieve basic competence. Master-level play could then be achieved via any of the cited works that focus on sophisticated planning and search.
task that has received recent attention from the planning and RL communities [10] TI 37 19
Starting with only the natural- language goal of "win the game", our agent builds world model over the first 50 actions. The resulting code human- numan-understandable (Appendix(D) DD, and generalizes to solving levels with more boxes (Fig. 3). While the system cannot solve very hard Sokoban levels (eg, boxes), that an artifact of the difficulty of planning, and could be addressed by plugging the world model into any of the techniques cited above. In contrast to this work, both model- and model-free ffeee deep RL require millions of experiences to solve basic levels (Fig.3) 3D).
(A) Sokoban example gameplay, 3 boxes
c
G
20 actions
reward
reward
reward
(B) Learning curves. Solve rate test
token
(D) deep RL comparison
boxes
boxes,
led
PPO
drcamerVa
solve rate
1.0
1.0
tokens
0.5
aptimism yes optimism
yes optimism
0.0
BaAut
0.0
10a
10a
10a
10a
10a
6
100 10a 10*
steps
steps
episodes
steps
(E) Teleportation Portals: Example gameplay by learned agent
sale
posh
teleporter
Figure 3: (A) Sokoban domain (per-step reward of -0.1 ellided from figure). (B) Learning curves. ReAct has the same pretrained knowledge of Sokoban but cannot effectively play the game. (C) Our method has different asymptotic I.I. cost compared to prior L.M agents, which consume LLM calls/tokens at every action. (D) Deep RI. takes >1 million steps to learn 2-box Sokoban. (E)
Nonstandard Sokoban with teleport gates
Almost surely, an important reason why our system learns quickly is because the underlying LLM already knows about Sokoban from its pretraining data, and can quickly infer that it is playing similar game. However, simply knowing about Sokoban does not suffice for the LLM to play the game, as demonstrated by the poor performance of ReAct (Fig. BB). ReAct is a baseline which prompts the LLM with the state-action history, then asks it to think step-by-step (Reason) and before predicting an action (Act). Quantitatively, ReAct succeeds on only 15% + 8% of basic levels, showing that pretrained knowledge of Sokoban does not, by itself, allow strong play. ReAct-style i.a.). also require expensive LLM calls every action, and so asymptotically their cost grows linearly with the number of actions taken. Our approach has different asymptotics:
architectures (29)

---

after front-loading 400k LLM tokens (about S15), it can issue as many actions as needed without subsequent LLM queries (Fig. BPC).
To further demonstrate that pretrained knowledge of Sokoban cannot explain the success of our system, we modify the game include extra rules by adding pair of warp gates/teleportation portals to each level which the agent can use to instantly transport itself (Fig. BE). The teleportation portals have subtle dynamics, because they become deactivated whenever they are blocked by Sokoban box. Our agent continues to be able to learn in this environment, including modeling the blocking behavior of the teleportation gates, and learns to use the teleporters to more rapidly solve levels.
On Sokoban, the optimism under uncertainty objective (02, orange curves in Fig. 3B) has little effect: Sokoban has dense reward structure that allows easy learning through random exploration. Next we consider problems with sparse rewards and natural language instructions, which the optimism objective is designed to handle.
Minigrid. To better understand the transfer learning and instruction following aspects of our approach, we next study Minigrid [9 [], suite of grid games designed for language-conditionee RL. Minigrid environments include objects such as keys, doors, walls, balls, and boxes.
Fig. illustrates results for our agent playing a sequence of minigrid environments, while Ap- pendix A.1 gives walk-through of an example learning trajectory. The agent interacts with each environment episodically through different randomly-generated generated levels. We order the environments into a curriculum designed to illustrate different forms of transfer learning. For example, when transferring from the first to the second environment, the agent needs to extend its knowledge to incorporate new objects and their associated dynamics (keys and doors). Learning about these new objects requires extra environment interactions, during which the agent experiments with the objects to update its transition function. In contrast, the third environment presents no new dynamics, but instead introduces new natural-languanl goals. Our agent can follow these natural- -atural-language instructions by enforcing optimism under uncertainty (Φ2)..
(A) minigrid environments and curriculum ordering: testing different forms of transfer enw door key
eriv 1:empty
ervv 3: unlock
ertv letch
cnv S: unlsccu pickup
new thems tiore few shot pentrolization
neo nevand: Bolltt tentraltzation
new transtions: loosiall gensralzation
transhtion: shot-sott gen.
get to the green goal squaaee
"use the key open the door get goal"
"open the door"
"fetch a green ball"
"pick the red box
(B) curriculum learning results
empty
doorkey
unlock
fetch
unlockpickup
1.0
2.0
0.5
0.5
0.5
0.0
10
10
10
10P
30 steps
10
0.0
10a
10
10'
10a
0.0
10F
10a
10a
10f
10
10a
stops
steps
stegs
(C) no curriculum results
empty
doorkey
unlock
fetch
unlockpickup
1.0
1.0
1.0
0.5
0.5
0.5
0.0
0.0
0.0
10
10-
10a
100
10a
10
10a
10a
10
10a
10
10
10P
10f
10'
10
10*
steps
steps
steps
steps
steps
Figure 4: (A) Minigrid environments ordered into a curriculum that tests different kinds of transfer learning. (B) Transfer learning performance, compared with (C) performance when solving each environment independently. Appendix Fig. deep RL comparison.
Transfer is especially helpful in quickly solving new environments (Fig. 4B). Without transfer, more episodes of exploration are needed to collect experience to build a good world model. However, optimism under uncertainty (0²) helps in non-transfer setting by promoting goal-directed exploration, and in fact, absent transfer, it is only with de that WorldCoder can solve the harder environments. Optimism is also necessary for zero-shot generalization to new goals (Fig. 4B, env 3).
To better understand the importance of which both encourages exploration and enables natural- language instruction following-we contrast against a version of our approach which simply prompts GPT4 to generate a new reward function upon receiving a new goal (green Fig. 43, labelled 'prompt4reward"). Theoretically this suffices to follow new natural-languag instructions, provided
7

---

the LLM is reliable. Surprisingly, this ablation struggles to follow new instructions (e.g.. transfer from env 2-+3 in Fig. [A)). showing the importance @2 in correcting mistakes made by the LLM when generating reward functions.
AlfWorld. To test the scalability of our method we work with the AlrWorld robot task planning domain. This is a household robot planning environment with variety of different kinds of objects, such as microwaves, cabinets, utensils, food, and different rooms the robot can navigate, We convert AlfWorld into a MDP by representing the state as a set of fluents (in the style of PDDL [2 ). and study our agent as progresses through representative AlfWorld tasks.
In Fig. ( we find that our method learns models of picking up and placing objects, and then improves its model by learning how use cooking appliances such refrigerators and stoves. This stresses the scalability of the exploration and program synthesis methods: We find that we can synthesize a world model containing 250+ lines of code, which serves as an adequate model of how AlfWorld works (Appendix E,, and that the optimism under uncertainty objective is necessary for nonzero performance on all of these tasks. Typically the agent solves the task the first episode, but requires around 20 exploratory steps to get reward, with subsequent episodes serving to make small corrections the world model. Fig. I iiniicates with arrows the episodes where these small corrections occur.
These tasks have relatively long horizons, so we use a more sophisticated planner based on MCTS. Because these tasks also have sparse rewards, engineer -onain- partial reward function that incentivizes actions and states which have more textual overlap with the natural-language goal, based on established document retrieval metrics (BM25 (50), Appendix B framework can interoperate with different kinds of planners.
This shows that our
Pick-Place
Heat-Place
Cool-Place
Pick-Two-Place
gota, put pickup
no no-pickup pickup--trom from- closed-receptacle
ND- put-in/on heat receedi-e
no wo-tool-withh. is.cool closed-fridge
reward
no- pickup- when-holding
not_ishhot
correction
Reward
20
sdo15
20
20
20
SdRIS
10
episode
episode
episode
episode
yes optimism reward
yes optimism steps
no optimism reward
no optimiem steps
Figure 5: AlfWorld environments and tasks. Each learning curve shows average reward at each episode (solid line) and how many steps the episode took (dashed), averaged over seeds. Curves annotated with arrows and text explaining what was learned at each episode. Optimism objective is necessary for any non-zero performance.
4 Related Work
World models. World modeling is classic paradigm for decision-making: It is the basis of model-
based RL 22 (66), lok-aad-motioon motion-planning robotics [32], and it is corroborated by classic
behavioral experiments suggesting biological decision- making features a similar architecture [66].
In practice, world models are hard-won, requiring either large volumes of training data [26 25 or careful hand-engineering
32 (cf. Mao et al. 44 Konidaris et al. (36)).
Neurosymbolic world models, such as Cosmos and NPS [56] 22], learn a factored, object- neural world model. This factoring helps compositional generalization like in our work-but importantly they can learn from raw perception, but at the expense of transfer and sample complexity. Combining our work with these others might be able to get the best of both.
8

---

LLMs as a world model. Whether LLMs can model the world is an open question, but there is evidence that, given the right training data in large quantities, transformers can act as decent world models, at least within certain situations 38 76 45 These works aim to learn a rich but frozen world model from a relatively large volume of examples. We tackle different problem: building a simple world model on-the-fly from modest amount of data.
LLMs for building world models. Recent works 79 74 24] consider using LLMs to generate plan- ning operators: kind of world model that as abstract, symbolic, and expressed in domain-speciffc programming language for planning (cf. DECKARD [46], another LLM system which generates state- Maatemmahhine world models). In these works, the primary driver of world-model generation-wil the LLM first inputs- is natural language describing affordances and goals. Our work considers a different problem: building world models from cost-and-foremod from interacting with the environ- ment. In practice, agents have knowledge both from language and from acting in the world, and so these families of works should be complementary.
LLMs for (ecision-making is an emerging paradigm that includes ReAct [77] and many others [29 78 41 I i.a.], which directly use LLMs to issue actions and reason about their consequences in the world. For instance, ReAct works by prompting the LLM think step-by-step and then predict an action. To the extent that these methods use a world model, it is implicitly encoded within the weights of neural network. We instead build an explicit world model.
Programs as Policies. Instead of learning a world model, one can learn a policy as a program. The first wave of these works 69 70 considered domain-specifee languages, while recent LLM work 73 39 61 uses more flexible general-purposs languages like Python. An advantage of learning a policy is that it does not need to model all the details of the world, many of which may be irrelevant to decision making. A disadvantage is that policies cannot readily generalize to new goals-unlike world models, which can be used by a planner to achieve a variety of objectives. Relatedly, other recent work considers synthesizing programs that implement reward functions [42) and then generating policy with conventional deep RL.
Programs as world models. We are strongly inspired by existing program synthesis algorithms for constructing world models from state-action trajectories [13] 67]. We believe that this family of methods will not be generally applicable until they can support general-purpon programming languages: So far these works have used restricted domain-specifee languages, but we show that a general-purpose computational language, like Python, can be used to learn world models, which hope expands the scope of this paradigm. We also show how to bias learning toward goal-directed behaviors, and how to support transfer across environments and goals. Last, we simplify the core program synthesis algorithm: the cited prior works required relatively intricate synthesis algorithms, which we can avoid by using LLMs as general-purposs synthesizers. We hope our work can help make this paradigm simpler and more general.
Turing-complete
Other works have also explored how humans can manually provide knowledge to RL agents via source code: e.g., RLLang [5T] uses programs to specify parts of policies and world models, which could be combined with our system to integrate prior knowledge.
Exploration & Optimism in the face of (model) uncertainty. uncertain consequences is common in decision- making, including in model-based RL 60 31. We mathematically instantiate that principle in new way-as logical constraint between program and a planner-and show how it can work with LLM-guided discrete program search.
Being optimistic about actions with
(63 40 75 3
5 Limitations and Open Directions
Our work has important limitations, and naturally suggests next steps. Currently we assume deter- ministic dynamics, which could be addressed by synthesizing probabilistic programs [14] 21]. Given recent advances in synthesizing probabilistic programs [53], together with advances in using LLMs for deterministic code, this limitation seems nontrivial but surmountable.
By representing knowledge as code, our approach delivers better sample efficiency and transferability, but at high cost: Our states must be symbolic and discrete, whereas the real world is messy and continuous. While the obvious response is that the agent can be equipped with pretrained object detectors-a common assumption in robotic task planning [36] i.a. alteemaive routes include

---

multimodal models [3]] and using neurosymbolic programs [6 56 perception and symbol processing, which might be more robust to missing or misspecified symbols.
64 to bridge the gap between
Last, our method uses only a very basic mechanism for growing and transferring its knowledge. Instead of prompting debug its code, could have built library of reusable subroutines and classes shared across different environments and goals, reminiscent of library learning systems [16] 72 23 4. which refactor their code to expose sharable components. Further developing that and other ways of managing and growing symbolic knowledge about the world remains a prime target for future work.
Acknowledgements.
This work supported by NSF grant #2310350 and gift from Cisco.
References
[1] Michael Ahn, Anthony Brohan, Noah Brown, Yevgen Chebotar, Omar Cortes, Byron David, Chelsea Finn, Chuyuan Fu, Keerthana Gopalakrishnan, Karol Hausman, Alex Herzog, Daniel Ho, Jasmine Hsu, Julian Ibarz, Brian Ichter, Alex Irpan, Eric Jang, Rosario Jauregui Ruano, Kyle Jeffrey, Sally Jesmonth, Nikhil Joshi, Ryan Julian, Dmitry Kalashnikov, Yuheng Kuang, Kuang-Huce Lee. Sergey Levine, Yao Lu, Linda Luu, Carolina Parada, Peter Pastor, Jornell Quiambao, Kanishka Rao, Jarek Rettinghouse, Diego Reyes, Pierre Sermanet, Nicolas Sievers, Clayton Tan, Alexander Toshev, Vincent Vanhoucke, Fei Xia, Ted Xiao, Peng Xu, Sichun Xu, Mengyuan Yan, and Andy Zeng. Do as i can and not as say: Grounding language in robotic affordances. In arXiv preprint (iXi::2204.01999 2022.
[2] Peter Auer. Using confidence bounds for exploitation Learning Research, 3NNov):397-422, 2002.
trade-offs. Journal of Machine
[3] Mohammad Gheshlaghi Azar, Ian Osband, and Rémi Munos. Minimax regret bounds for reinforcement learning. In International conference on muchine learning, pages 263-272.
PMLR, 2017.
[4] Matthew Bowers, Theo X. Olausson, Lionel Wong, Gabriel Grand, Joshua B. Tenenbaum, Kevin Ellis, and Armando Solar-Lezama Top-down synthesis for library learning. Proc. ACM Program. Lang., 7(POPL., jan 2023. doi: 10. 20.1145/3571234, URL httpt:ttt//ooiii http///oooogg/..///
1145/3571234
[5] Cameron B Browne, Edward Powley, Daniel Whitehouse, Simon M Lucas, Peter Cowling, Philipp Rohlfshagen, Stephen Tavener, Diego Perez, Spyridon Samothrakis, and Simon Colton. A survey of monte carlo tree search methods. IEEE Transactions on Computational Intelligence and Al in games, 4(1):1-43, 2012.
[6] Swarat Chaudhuri, Kevin Ellis, Oleksandr Polozov, Rishabh Singh, Armando Solar-Lezama, Yisong Yue, et al. Neurosymbolic programming, Foundations and Trends@ in Programming Languages, 7(3), 2021.
[7] Xinyun Chen, Maxwell Lin, Nathanael Schärli, and Denny Zhou. Teaching large language models to self-debug. arXiv preprint arXiv:2304.05128, 2023.
[8] Maxime Chevalier-Boiswar Lucas Willems, and Suman Pal. Minimalistic gridworld environ- ment for openai gym. http:///gggtuuub.coomm/iiieeeeee
2018.
[9] Maxime Chevalier-Boisverr, Bolun Dai, Mark Towers, Rodrigo de Lazcano, Lucas Willems, Salem Lahlou, Suman Pal, Pablo Samuel Castro, and Jordan Terry. Minigrid & miniworld: Modular & customizable reinforcement learning environments for goal-oriented tasks. CoRR, abs/2306. 13831, 2023.
[10] Leah Chrestien, Stefan Edelkamp, Antonin Komenda, and Tomá' Pevný. Optimize planning heuristics to rank, not to estimate cost-to-goal. In seirth- Conference Neurul In- formation Processing Systems, 2023. URL. https: http:::/openreview.nettt www.wevwwwwwwe/oruuueuee.
Pgy6sgsLPY
10

---

[11] Stephen Chung, Ivan Anokhin, and David Krueger. Thinker: Learning to plan and act. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL h //openreviee ew wwtpe/iiweeeeeetuueteemm
[12] Rémi Coulom. Efficient selectivity and backup operators in monte-carlo tree search. In International conference on computers and games, pages 72-83. Springer, 2006.
[13] Ria Das, Joshua B. Tenenbaum, Armando Solar-Lezama, and Zenna Tavares. Combining functional and automata synthesis to discover causal reactive programs. Proc. ACM Program. Lang. 7(POLL) jan 2023. doi: 10/1149/3572249 URL https: http://doo.oorgg/////// wwg. 1145/ 3571249
[14] Luc De Raedt, Angelika Kimmig, and Hannu Toivonen. Problog: probabilistic prolog and its application in link discovery. In Proceedings of the 20th International Joint Conference on Artifical Intelligence, IJCAI'07, page -468-2473, San Francisco, CA, USA, 2007. Morgan Kaufmann Publishers Inc.
[15] Carlos Diuk, Andre Cohen, and Michael L Littman. An object-oriented representation for efficient reinforcement learning. In Proceedings of the 25th international conference on Machine learning, pages 240-247, 2008.
[16] Kevin Ellis, Lionel Wong, Maxwell Nye, Mathias Sable-Meyer, Luc Cary, Lore Anaya Pozo, Luke Hewitt, Armando Solar-Lezama, and Joshua B Tenenbaum. Dreamcoder: growing gener- alizable, interpretable knowledge with wake-sleep bayesian program learning. Philosophical Transactions of the Royal Society A, 381(2251):20220555 2023.
[17] Richard Evans, Matko Bošnjak, Lars Buesing, Kevin Ellis, David Pfau, Pushmeet Kohli, and Marek Sergot. Making sense of input. Artificial Intelligence, 299:103521, 2021.
[18] Dieqiao Feng, Carla P Gomes, and Bart Selman. A novel automated currieu- lum strategy to solve hard sokoban planning instances. In H. Larochelle, M. Ran- zato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Informa- tion Processing Systems, volume 33, pages 3141-3152. Curran Associates, Inc., 2020. files//aperr/220000
URL
[19] Dieqiao Feng. Carla Gomes, and Bart Selman. Left heavy tails and the effectiveness of the policy and value networks in dnn-based best-first search for sokoban planning, 2022.
[20] Caelan Reed Garrett, Rohan Chitnis, Rachel Holladay, Beomjoon Kim, Tom Silver, Leslie Pack Kaelbling, and Tomás Lozano-Pérez. Integrated task and motion planning, 2020.
[21] Noah D. Goodman, Vikash K. Mansinghka, Daniel Roy, Keith Bonawitz, and Joshua B. Tenen- baum. Church: language for generative models. In Proceedings of the FheentyFourrte Confer- ence on Uncertainty in Artificial Intelligence, UAI'08, page 220-229, Arlington, Virginia, USA, 2008. AUAI Press. ISBN 0974903949.
[22] Anirudh Goyal, Aniket Didolkar, Nan Rosemary Ke, Charles Blundell, Philippe Beaudoin, Nicolas Heess, Michael Mozer, and Yoshua Bengio. Neural production systems: Learning rule-governed visual dynamics. arXiv preprint arXiv:2103.01937. 2021.
[23] Gabriel Grand, Lionel Wong, Matthew Bowers, Theo X. Olausson, Muxin Liu, Joshua B. Tenenbaum, and Jacob Andreas. Lilo: Learning interpretable libraries by compressing and documenting code, 2023.
[24] 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. NeurIPS, 2023.
[25] David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In Advances Neural Information Processing Systems 31, pages 2451-2463. Curran Associates, Inc., 2018. https: /wor Ldmode Ls. thub.

---

[26] Danijar Hafner, Timothy Lillicrap. Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. ICLR, 2021.
[27] Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models, 2023.
[28] Assaf Hallak, Dotan Di Castro, and Shie Mannor. Contextual markov decision processes. arXiv preprint arXiv: 02259, 2015.
[29] Shibo Hao, Yi Gu, Haodi Ma, Joshua Jiahua Hong, Zhen Wang, Daisy Zhe Wang, and Zhit- ing Hu. Reasoning with language model is planning with world model. arXiv preprint 230.:299...199.., 2023.
[30] Yingdong Hu, Fanqi Lin, Tong Zhang, Li Yi, and Yang Gao. Look before you leap: Unveiling the power of gpt-4v in robotic vision- lssionlanguuae planning. arXiv preprint 2rXiv:2311.178..
2023.
[31] Chi Jin, Zhuoran Yang, Zhaoran Wang, and Michael Jordan. Provably efficient reinforcement learning with linear function approximation. In Conference on learning theory, pages 2137- 2143. PMLR, 2020.
[32] Leslie Pack Kaelbling and Tomás Jozano-Pérez Hierarchical task and motion planning in the now. In 2011 1 IEEE International Conference on Robotics and Automation, pages 1470-1477, 2011. doi: N0.1109/CCRA22111.99900
[33] Vladimir Karpukhin, Barlas Oguz, Sewon Min, Patrick Lewis, Ledell Wu, Sergey Edunov, Danqi Chen, and Wen-tau Yih. Dense passage retrieval for open-domain question answering. arXiv preprint arXivv2004.04906, 2020.
[34] Mei Kobayashi and Koichi Takeda. Information retrieval on the web. ACM computing surveys (CSUR), 12(2):144-173 2000.
[35] Levente Kocsis and Csaba Szepesvári. Bandit based monte-carlo planning. In European conference on machine learning, pages 282-293. Springer, 2006.
[36] George Konidaris, Leslie Pack Kaelbling, and Tomas Lozano-Pered From skills to sym- bols: Learning symbolic representations for abstract high-level planning. Journal of Artificial Intelligence Research, 61:215-289, 2291 2018.
[37] Kalle Kujanpää, Joni Pajarinen, and Alexander Ilin. Hybrid search for efficient planning with completeness guarantees. In Teirth-seventh Conference on Neural Information Processing Systems, 2023. URL httL:htts:/// //tp:://////eeeeeeeeeee www.fwwr/orrrmmmm 1tte///rrrmmdmdo4625 d=bYoc46
[38] Kenneth Li, Aspen K Hopkins, David Bau, Fernanda Viégas, Hanspeter Pfister, and Martin Wattenberg. Emergent world representations: Exploring sequence model trained on synthetic task. ICLR, 2023.
[39] Jacky I Liang, Wenlong Huang, Fei Xia, Peng Xu, Karol Hausman, Brian Ichter, Pete Florence, and Andy Zeng. Code policies: Language model programs for embodied control. In arXiv preprins arXiv:2209.07755. 2022.
[40] Qinghua Liu, Praneeth Netrapalli, Csaba Szepesvari, and Chi Jin. Optimistic mle: A generic model-based algorithm for partially observable sequential decision making. In Proceedings of the 55th Annual ACM Symposium on Theory of Computing, pages 363-376, 2023.
[41] Zhihan Liu, Hao Hu, Shenao Zhang, Hongyi Guo, Shuqi Ke, Boyi Liu, and Zhaoran Wang. Reason for future, act for now: A principled framework for autonomous I1m agents with provable sample efficiency. arXiv preprint arXiv:2309.1 17322, 2023.
[42] Yecheng Jason Ma, William Liang, Guanzhi Wang, De-An Huang, Osbert Bastani, Dinesh Jayaraman, Yuke Zhu, Linxi Fan, and Anima Anandkumar. Eureka: Human-level reward design via coding large language models. arXiv preprint arXiv: Arsiv-2310.12931, 2023.
[43] Christopher D Manning, Prabhakar Raghavan, and Hinrich Schüitze Introduction to information retrieval. Cambridge university press, 2008
12

---

[44] Jiayuan Mao, Tomás Lozano-Pérez, Josh Tenenbaum, and Leslie Kaelbling, Pdsketch: Inte- grated domain programminng learning, and planning. Advances in Neural Information Process- ing Systems, 35:36972-36984, 2022.
[45] Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample-efficient world models. In The Eleventh International Conference on Learning Representations, 2023. URL
http::://ppeeeeveee net/ fow.mmid=vnuuuAAaCC..
[46] Kolby Nottingham, Prithviraj Ammanabrolu, Alane Suhr, Yejin Choi, Hannaneh Hajishirzi, Sameer Singh, and Roy Fox. Do Embodied Agents Dream of Pixelated Sheep: Embodied Decision Making using Language Guided World Modelling. In International Conference on Machine Learning (ICML), 2023.
[47] Theo X. Olausson, Jeevana Priya Inala, Chenglong Wang, Jianfeng Gao, and Armando Solar- Lezama. Is self-repair silver bullet for code generation?, 2023.
[48]
Antonin Raffin.
RI
baselines3
ZOO.
11-baselii biselines3- ZoO, 2020.
[49] Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dormann. Stable-baselines3: Reliable reinforcement learning implementations. Journal of Machine Learning Research, 22(268):1 2021. URL http:// org/ papers/v22 20-1364 html
/
[50] Stephen Robertson, Hugo Zaragoza, et al. The probabilistic relevance framework: Bm25 and beyond. Foundations and Trends@ Information Retrieval, 3(4):333-389 2009.
[51] Rafael Rodriguez- Sanchez, Benjamin A. Spiegel, Jennifer Wang, Roma Patel, Stefanie Tellex, and George Konidaris. Rlang: A declarative language for describing partial world knowledge to reinforcement learning agents, 2023.
[52] Stuart Russell and Peter Norvig. Artificial intelligence: modern approach. Pearson, 2016.
[53] Feras Saad. Scalable Structure Learning, Inference, and Analysis with Probabilistic Programs. PhD thesis, Massachusetts Institute of Technology, 2022.
[54] Max-Philipp B. Schrader. gym-sokoban. zym-sokkoan, 2018
httpss//ggiiuubb..commm/hhrrere
[55] Julian Schritwieser,, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Si- mon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7839):604-699 2020.
[56] Atharva Sehgal, Arya Grayeli, Jennifer J. Sun, and Swarat Chaudhuri. Neurosymbolic grounding for compositional world models, 2023.
[57] Mohit Shridhar, Xingdi Yuan, Marc-Alexandre Côté, Yonatan Bisk, Adam Trischler, and Matthew Hausknecht. ALFWorld: Aligning Text and Embodied Environments for Interactive Learning. In Proceedings of the International Conference on Learning Representations (ICIR), 2021. URw/https: // /ttpss/////xvv.o/o org abs 2010. 03768
[58] David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driess- che, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Mare Lanctot, et al. Mas- tering the game of go with deep neural networks and tree search. nature, 529(7587):484449) 2016.
[59] Tom Silver and Rohan Chitnis. Pddlgym: Gym environments from pddl problems. In Interna- lional Conference Automated Plamning and Scheduling (ICAPS) PRI. Workshop, 2020. URL /
com tecoo//omm toms pdd
[60] Sean R Sinclair, Siddhartha Banerjee, and Christina Lee Yu. Adaptive discretization for episodic reinforcement learning in metric spaces. Proceedings of the ACM on Measurement and Analysis of Computing Systems, 3(3):1-44, 2019.
13

---

[61] Haotian Sun, Yuchen Zhuang, Lingkai Kong, Bo Dai, and Chao Zhang. Adaplanner: Adaptive planning from feedback with language models. NeurIPS, 2023.
[62] Richard S Sutton. Integrated architectures for learning, planning, and reacting based on approximating dynamic programming. In Machine learning proceedings 1990, pages Elsevier, 1990.
[63] Richard Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press. 2018.
[64] Hao Tang and Kevin Ellis. From perception to programs: Regularize, overparameterize, and amortize. International Conference Machine Learning (ICML., 2023.
[65] Hao Tang, Keya Hu, Jin Peng Zhou, Sicheng Zhong, Wei-Long Zheng, Xujie Si, and Kevin Ellis. Code repair with Ilms gives an exploration-exploiaatio tradeoff. NeurlPS, 2024.
[66] Edward Chace Tolman and Charles H Honzik. Introduction and removal of reward, and maze performance in rats. University of California publications in psychology, 1930.
[67] Pedro A Tsividis, Joao Loula, Jake Burga, Nathan Foss, Andres Campero, Thomas Pouncy, Samuel Gershman, and Joshua B Tenenbaum. Human-level reinforcement learning through theory-based modeling, exploration, and planning. arXiv preprint arXiv:2107. 2021.
[68] Karthik Valmeekam, Matthew Marquez, Sarath Sreedharan, and Subbarao Kambhampati. On the planning abilities of large language models a critical investigation. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. UR//nttss:/oppeeevveee SettForrum?1-xGGEXXISEE
[69] Abhinav Verma, Vijayaraghavan Murali, Rishabh Singh, Pushmeet Kohli, and Swarat Chaudhuri. Programmatically interpretable reinforcement learning. In International Conference on Machine Learning, pages 5045-5054. PMLR, 2018
[70] Abhinav Verma, Hoang Le, Yisong Yue, and Swarat Chaudhuri. Imitation-projected pro- grammatic reinforcement learning. Advances in Neural Information Processing Systems, 32, 2019.
[71] Lev Semenovich Vygotsky. Mind in society: The development of higher psychological processs.. Harvard university press, 1980.
[72] 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 preprini arXiv: Arxiv-2305.1629) 2023.
[73] Huaxiaoyue Wang, Gonzalo Ponzaiez-Pumaricga, Yash Sharma, and Sanjiban Choudhury. Demo2code: From summarizing demonstrations to synthesizing code via extended chain-of- thought. NewrlPS, 2023.
[74] Lionel Wong, Jiayuan Mao, Pratyusha Sharma, Zachary S. Siegel, Jiahai Feng, Noa Korneev, Joshua B. Tenenbaum, and Jacob Andreas. Learning adaptive planning representations with natural language guidance, 2023.
[75] Chenyang Wu, Tianci Li, Zongzhang Zhang, and Yang Yu. Bayesian optimistic optimization: Optimistic exploration for model-based reinforcement learning. Advances in neural information processing systems, 35: 35:11420-14223 2022.
[76] Jiannan Xiang, Tianhua Tao, Yi Gu, Tianmin Shu, Zirui Wang, Zichao Yang, and Zhiting Hu. Language models meet world models: Embodied experiences enhance language models. NeurlPS, 2023.
[77] Shunyu Yao, Jeffrey Zhao, Dian Yu, Nan Du, Izhak Shafran, Karthik Narasimhan, and Yuan Cao. React: Synergizing reasoning and acting in language models. arXiv preprint arXiv:2210.03629, 2022.
[78] Zirui Zhao, Wee Sun Lee, and David Hsu. Large language models as commonsense knowledge for large-scale task planning. NewrIPS, 2023.
14

---

[79] Feiyu Zhu and Reid Simmons. Bootstrapping cognitive agents with a large language model. In AAAI, 2024.
[80] Yutao Zhu, Huaying Yuan, Shuting Wang, Jiongnan Liu, Wenhan Liu, Chenlong Deng, Zhicheng Dou, and Ji-Rong Wen. Large language models for information retrieval: A survey. arXiv preprint 2rXiv:3388.07007 2023.
15

---

A Theoretical analysis of Sample Efficiency using optimism under uncertainty (02)
The optimism under uncertainty objective (61 A 22 in Sec. is much more sample efficient than the traditional data-consistency objective (σ)) as shown in Figure 4 in our experiments. We will later show a learning trajectory of in the MintCrri--UnlackPPccpp---0
environment in Appendix |
This UnlockPickup environment is difficult for exploration due its large search space. PPO failed to gain enough positive-reward supervision signals after 3 x 10 actions as shown in Figure 6 Our method failed without the optimism objective well. Nevertheless, our method with the optimism objective learned the correct world model from scratch with no more than 100 actions due to its better sample efficiency in exploration.
Here we provide a simple theorem stating that significantly better sample efficiency is theoretically guaranteed when using this objective (polynomial w.r.t the diameter of the state space and the solution space) following an intuitive observation:
Observation. When planning with world model that satisfies the Φ1 A Ø2 there are only two possible outcomes:
Either the model correct: the agent achieves the goal successfully;
Or the model is incorrect: the agent then must find a counter-example to its world model's prediction and gain more knowledge of the environment.
shows that the world model that satisfies é1 A Φ2 is either correct or is efficient in guiding the agent to find a counter-exampee when in search of the correct world model, each interaction data (s, a, s', r, d) in the replay buffer implies constraint to satisfy for the potentially correct world models in the space of world models. Collecting data that can be explained by all potentially correct world models is useless because implies no stricter constraints than the current ones. Only through the counter-examples which cannot be explained by the current world model, the size of the set of the potentially correct world models can become smaller and smaller, which eventually leads set that only contains the good-enough world models that guide the agent towards the goal.
of the current world model. Counter-example are precious because,
The optimism under uncertainty objective is much more sample efficient to get those valuable counter- examples than random exploration. For deterministic environments, the number of actions to find counter-ews exaame with world model that satisfies Φ1 A 62 is guaranteed to be smaller than the diameter of the state space, Ds. when equipped with an optimal planner (by definition). Random exploration needs exponentially more number of actions find this counter-example in the worst case. (Note that we do not assume free reset to all possible states, so the agent needs to go through the trajectories from the initial state each target state.)
More formally, we show that the maximum number of actions to learn/find a good-enough world model is polynomial w.r.t. intrinsic properties of the state space and model space:
Definition A.1. A dataset D is mutually independent w.r.t. a solution space M, denoted, D JL M. if and only if for all data points in d E 'D there exists solution in M which explains all the data except for d:
Ad E D, 3m E .,, V d A (Vd' E D {d), m d').
Definition A.2. The logical dimensionality of M, denoted as K M, is
KIL = max[[D]: vD where D IL M]
is straightforward to prove that this is at most [M]:
Lemma A.3. K i M S [M], i.e., the maximum size of the mutually independent data sets w.r.l to solution space is smaller than the size of the solution space.
Proof: each data point d E 'D at least exclude one more unique solution in the solution space than the other data points. Otherwise, the data set is not mutually independent w.r.t. the solution space.
Note that this is the loosest upper bound of K M. We assume nothing about the model space and the learning algorithm. In practice, K M should be much smaller than the size of the solution space. For example, for linear models, we only need n independent data points to characterize the correct solution for d E I.".
16

---

Theorem A.A. Assume an episodic MDP (S,A,T, R, 7). Assume an agent acting according to an optimal planner operating over world model (T,, Rt) E T x R. satisfying A Q2 and that the true MDP is in the agent's 's model class: T T and R E R. Then the maximum number of actions needed achieve the goal (exit with positive reward) is Ds,A,T x (KTxR + 1).
Proof Outline:
Lemma a: Each replay buffer data set, D, can be represented by its mutually independent data subset, DUTRR, from the perspective of finding the correct world model, i.e..
DUTER S A 1T T X R E E (T, ( ( DITTR.
For the replay buffer, D(O), at any step if world model (710, R(0)) satisfies the optimism under uncertainty objective as well as the traditional data-consistency objective, i.c., (T', Rt) D(t), A (T, ft) }- dz2, then either the word model is correct, which means it can guide the agent to the goal successfully, the agent finds counter-example, d', the current world model in Dg steps when given an optimal planner.
Lemma 3:
Lemma y: This counter example, d', is mutually independent the replay buffer, D(t), because there a model (T()), R()) such that (T19, d' (T(), R(9) ) D(t).
Given these lemmas, we have
D (D ITXRR 2
D0 LLXR +1
and therefore
D ( 2 KILTXR XR +1+ D 2 KITXRR + 1 KLTXR- LRXR
Assume the world model after Ds x (KITTR + 1) steps is still incorrect, we then build mutually independent data set, (Dax (KLTR()) with size larger than Ki TxXR, which is contradictory to its definition.
The proofs of Lemma a and are straightforward by definitions. For example, we can build (()) by dropping data points in ( that are not mutually independent to the left subset one by one. We can build D (t++1) LR by adding the counter-example d' to D DITTR
I x
We prove Lemma ,8 by construction. Given the definition the optimism under uncertainty objective, Φ2 in Sec. 2
2 s0, (()))
7or 81, 02 82, ..., Ot, 8t
Vi E [f] 7(0) (si-1, 8sA
3r 0: ( (5c-1, ag, 88) (r,
there exists trajectory 01 $1, 02 82, ....... ag St such that either the model correctly leads the agent to the goal, i.e., 3r > R(c)(5t-1. ar, 31) = (r, or there exists a counter-example, i.e., Bi E [4] (s-11,)) * (()- v R( (c) a;, sz) * (( si)... which means (7(9, R(0) / (st_1, 04 86, ri, d.). We also have (710, P! D()) due to (T(1), Re(e)) + d1. We thus prove Lemma 3.
A.1 Example learning trajectory using (6) Φ2) MiniGridUnnlockPicuu
To demonstrate the effectiveness of the optimism objective (62 in Sec. (c in improving the sample efficiency through guided exploration, we show here an example of the learning trajectories in the MiniGrrddUnlockPickii
environment. In comparison with agents that only use the traditional data-consistency objective (Φ1), which merely relies on random exploration to gain new knowledge of the world through new data, agents that use the extra optimism under uncertainty objective can
imagine the functionality actions and their interactions with the necessary tools, order to achieve the goal, without any real interactions with the environment:
17

---

explore much more efficiently guided by the imagined world model;
and correct the imagined world model given the newly collected data after the efficient exploration.
The learning trajectory is as follows:
Original reward function
Original transit function
NULL
NULL
Random exploration No feasible move yet except for rotation.
Data-consistent transit-fume fune (figured out how to rotate)
def turn_ recttion)
if direction == (-1, 0):
return (0, -1)
elif direction (1, 8):
def ternn
if direction == (-1, 0):
return (0,
elif direction (1, 0):
def transition(state, action):
agents get get_eeettteeee by (sttess_b_nammmtttte, "Agent") assert len(agents) == 1, "There should be exactly one agent the state"
agent = agents[0] if action == "turn right": agent.direction
elif action "turn left": Add a condition for the "turn left" action Surn_leef(agent.direeee
agent.direction
The other actions does not to affect the state return state
18

---

Pata-consistent reward-func
(random guess given the textual mission)
def reward_func (state, action, next_state):
# If the state does not change after action, the reward is 0.0 and not done
if state next_state:
return 0.0, False
agent_n, agent_o = net_ennities_bre bymme
next_state, "Agent") [0], get_eenittees_by_nee
state, "Agent") box_n, box_o [B)x")[o], get_entittie__y_naat (nextt_s_ssaa statt, benniitteee.b__mmetttttee
"Box") [0]
# If the agent's position gets closer to the box give bigger reward
if abs(agent_n.x box_n.x) abs(agent_n. box box_n.y) abs (agentt_oo box_o.x) abs( agent_o.y box_o.y):
return 1.0, False
If the agent changes direction but position, the reward is and done elif agent_o.x agent_n.x and agent_o.y. agent_n.y and agent_o.direction != agent_n. direction:
return 8.0, False
the agent picks up the green box, give big reward and finish
elif agent_n.carrying and agent_n.carryingg.oooo
== "green":
return 100.0, True
the agent moves but does not make significant progress, give small penalty to make it efficient
else:
return -0.01, False
Goal-driven transit-func (imagine how move forward without interactions) (imperfect though as not considering the locked door)
def transition(state, action):
elif action 'move forward':
dx, dy agent.direction new_x, new_y = dx, Y dy
if not get_entitiee_by_natt
new_x,
new_y), 'Wall'):
agent.x, agent.y new_x, new_y
19

---

Goal-driven reward-furg fune (more efficient representation <- already correct)
def ffrrrdttetssttte,, action, next_state):
# Get all agent entities in the next state next_agents get_entities__y_nnet (next_sstaee Agent?)
Check if any agent is carrying a Box for next_ aagen__carryinng in lagent.carrying for agent in next_
if [ext_agents];
(next_agent_caml eext_agent_carrying.co return 1.0, True
Box) and 'green';
return 0.0, False
Guided by imagined transit-
Pata-consistent transit-fure (figured out the agent cannot step on the wall and the locked door. but still do not know how to go to the goal.)
func Trying to cross locked door
ddef transition(state, action):
elif action == 'move forward';
dx, dy agent.direction new_x, new_y = dx, dy door_at_new_lue
sr__entities_by_naet
senettitieesbyooosttionnnnnoo new_y), 'Door')
new_x,
if door door_at_new_loog
door_at_new_log = if return next_state
[oc
== "locked":
if get_entities_by_name
oet_entitiies.by_ssiiiin state, new_x, new_y), 'Wall'):
agent.x, agent.y new_x, new_y
20

---

Goal-driven transit-func (imagine how to pickup, toggle, and drop. without interactions) (imperfect though trying to pickup at the same position the agent)
def transition(state, action):
elif action == 'pick up':
Ltemm_attagent_Locattii logantt_oocaion
crnrttie_bbyposstionnneetttt agent.x, agent.y)
pickable_items [item for item in
grams__atagent__ocatte if naem.name not in ['Door', 'Wall', 'Agent']]
if pickable_iteme 1itblee_tttms__
agent.carrying pickable_items([0
n (pickable_eettmmmscss
elif action 'drop':
if agent.carrying:
agent.carrying.x, agent.carrying.y =
agent.x, agent.y
agent.carryiig None
elif action == 'toggle':
if issistaacceageet.caariiin, Key):
door_at_right [door for door in
right_x, right_y) if door name 'Door'
]
if door_at_ aight and door door_at_ _a__tttrrrhht [t_ color == cgrnt
door_at_right(0],
"unlocked'
return next_state
21

---

Guided by imagined transit- fune Trying step the key pick it up
Data-consistent transit-fum
(figured out the agent cannot step on the key, but still do not know how to go to the goal)
def transition(state, action):
elif action == 'move forward':
antities_at_front [entity for entity in front_x front_y) if entity.name != 'Agent']
if crriitess_atfrontee
if antities__t_ffrott/ name == "Door" and entities_at__root state == unlocked":
agent.x, agent.y front_x, front_y # other entities should prevent agent from moving forward
else:
agent.x, agent.y front_x, front_y elif action == 'pick up": = (inenttttesbbb_oootttioonnttttt
agent.x, agent.y)
pickable_items [item for item in items_at_agens aggen_location if name not in ['Door', 'Wall', 'Agent']]
if pickable_itee _itle::tttmt__
agent. ggent.carrying pickable_ Bickable_items__e next_state. remove (pickable [0])
Goal-driven transit-func
(imagine the agent can also pickup the object in front it, which is correct, without interactions.
def transitionnstate, action):
elif action 'pick up":
items_at_ litis_at_aeen__l_c_eee conniiifes_by_ossttiorrrressee agent.y) www.tiiiee.bb.pssstttoeeeeeeee
agent.x
ittms
front_x, front_y)
pickable_items [item for item in
ttems_st__geent_locttton
fttmm_att_front_locctttt if item.name not in ['Door', 'Wall', 'Agent']]
if pickable_itene items:
agent.carrying pickable [ next_state. remove (pickable tems [0])
22

---

Guided by imagined transit-
Data-consistent transit-fune (correct some detail about opening the door)
fune
Trying to toggle the door
def toggle_door (agent, next_state, next_x, next_y); doors_in_ sin__exx_positiinn = [door for door in next_x, next_y) if door name == 'Door']
if phors_in_nexx_ next_position and
doors_in itx_iieeettposiiio@i...o...... color == agent. carrying. color
state "open*
23

---

Guided by imagined transit- fune Trying pickup the box while carrying the key
Goal transit-func
(figured out cannot pickup objects while carrying others; im- prove the drop action that was imagined before) (SUCCEED!!)
def drop_ (oogiteemaggenn, next_state, next_x, next_y): entities, entitiis_tt_next_poee wtt._eee_e___stttonn
next_x,
next_y)
if not entities_in iifiiss_i_neex__oostiioo and agent.
carrying:
# Drop can only drap object if there's no obstacle and agent carries something.
agent.carrying.x, agent.carrytng.y = next_x, next_y
hww.sstate.appendaggntt.cm agent.carrying None
def wwwpoootcceeeeeeeeeeeeeeeeeeeeee y):
next_state, x,
dx, dy = x agent.x, agent.y
for i in range (dx), abs(dy))):
entities__t__nex_postte
agent.x if any iisiistanceeenttit Wall) or ( bsinstancceentiity Door) and entity.state != 'open') for entity in
dx, agent.y i dy)
return False
return True def Backup_item(agentt next_state): items_in mumns_iiincrrretttoatiioo
agent.x,
agent.y)
pickable_items = [item for item in location if item. Door', 'Wall', 'Agent']]
in agent.carrying is None: Agent can only pick an item when it is not carrying an item if not pickable_items: iimole
dx, dy agent.direction
facing_x, facing_y agent.x dx, agent .y + dy if next_state, facing_x, facing_y): Make there
wall door between
the agent and the item
items_in_fee conffaciig_lccattio = get_entities_ sentttiie_by_posittiieeee tny_ossttionnnext_sttt faciin__x, facing_y)
pickable_items [item item
locaaiiinlloltaioonnnonnoo.. name not in ['Door', 'Wall']]
location if item.
if pickable_iteme limolleiitemsse
agent.carrying pickable sitabbl_itemss tems next_state remove (pickable, items
The final synthesized world model code
class Entity:
24

---

def -_innit__self x, y, x*kwargs):
self. name self. self.__class_.. _amee.________ self.x =
self.y =
for key, value in wwargs.itemsi setattr(self, key, value)
def repr...
attr =,, (-join(ftkey)==vvalee if key not in ('name', 'x', 'y'))
for key, value in self.__dict__ ttems
attr: return f"[self. namel..namm)((ccc ([selfxx; {self.yh, (attr))"
else: return ([ssef.nname ((self.x), (self.y))
def
(self, other):
return (ge (self, key) == getattr key, None) for key in self___dfct_..kkes/s//
def (self=
return [ectrrmmenaotttttteessseeeeeeeer
())))
class Agent(Entity): pass
class Key(Entity): pass
class Door(Entity): pass
class Goal(Entity): pass
class Wall(Entity): pass
class Box(Entity): pass
class Ball(Entity): pass
class Lava(Entity): pass
def update__directioonaggee
action):
all_directions = [(0, -1), (-1, 0), (0, 1), (1, 0)]
current_dir_idt = all ditecirecti index
action 'turn right':
agent.direction all [It_dirrectioossuurrreettddii
1) % 4]
else: turn left
agent. dinettireetion (nl.irrectionssuuurtetttt.
1) % 4]
def pappittemaagent. next_state, next_x, next_y):
get_entities__y_pog tion (next_state, next_x, intittee_in_next_positte next_y)
if not
and agent.carrying
# Drop can only drop object if there's no obstacle and agent carries something.
agent.carrying.x, carrtncaryy next_x, next_y stt.stttee.pppeddd (ageettccanttnr
agent.carrying = None
def toggle_door (agent,, next_state, next_x, next_y): doors_in dbors__n__net__osittoe
[door for door in ret_entitiee__by_osittee
next_state, next_x, next_y) door. 'Door']
if nhors_in_next_positee and carrying.color
color == agent.
def beors__i__nex__psstttonne byntteeiiiteeeetiee name):
'open'
return entity for entity in entities if entity.name == name def centiiiee.by..osiiiinoeee.
x, y):
return entity for entity in entities if entity.x == and entity.y == y def haaccchoo_osstaleeeeeeeeeeeeed
next_state, x, y):
dx, dy for i in (ann abs(dy)))
agent.x, agent.y
entities__t entities_at_next staee__at_text_psstttn = bi i dx, agent.y
agent.
dy)
imy(ssinstanceeentii Wall) or (fsinstance(entitty, Door) entity .state 'open") for entity
return False
return True def pickup_iteemageet next_state): jiims_in_curreet_tocc agent.y)
get_entities.by.pot swn.iiiiee.y..sssttorrrettttttt. agent.X,
pickable_iteen itmss = [item for item in items_in_current_toosee if item. name not in ['Door', 'wall', ']
25

---

agent. carrying None: Agent can only pick up an when is not carrying an item if not pickable_itent items:
dx, dy agent.dfrectton
facing_x, facing_y agent.x dx, agent.y dy
if check_no_ brrcc__n_obsstcleeeceo wbwtaclle_betwee(((gg.... between(agent, next_ next_state, facing_x, x, facing_y)
: Make sure there is no wall or door between the agent and the item
items_ii_faccing__ location ge facing_x, facing_y)
pickable items [item for item items_in in fariiin_ffnngncaaacace location ntme not in ['Door', 'Wall']]
if pickable_items:
agent.carrying pickable_itemss [0]
next_state remove
(ee. (pickable_itens) [0])
def transition(state, actioo))
next_state = list(state)
agent dx, dy agent.direction
'Agent')
front_x, front_y agent.x dx, agent.y + dy if action == 'turn right' or action == 'turn left':
action)
elif action == 'move forward':
next_state, front_x, front_y) elif action == 'pick up':
update_posittoon
tackkp_itemaagenee next_state)
elif action == 'drop':
drop_ trop_ittemagget,, next_state, front_x, front_y)
elif action 'toggle':
toggle_door (agent, next_state, front_x, front_y)
return next_state
def update_ coanutteposittonnageenn next_state, next_x, next_y):
antities_at_next_posiiii next_y) if not (
Bet.entiiee.by.posiiiiooooo...e next_x,
isinstance Wall) (ennttyy Box) or (ennsttanceeettty... Ball) or Isiistaanceentitty Lava) or
Ilsinsttanceenntity Door) and entity,state 'open") Esinstance(entityy Key)
for entity in antities__t_nex__posiiii
):
agent.x, agent.y next_x, next_y
else:
agent.x, agent.y agent.x, agent.y # Agent stays in place
B Planning with Monte Carlo Tree Search and the BM25-based Heuristic
We propose a BM25-based heuristic to guide planning in text-based environments, and trade off the exploration and exploitation using variant of Monte Carlo Tree Search (MCTS) for deterministic environments with sparse rewards. The heuristic encourages the planner to focus on trajectories that are "closer" the goal as specified by users, while MCTS enables the planner explore other options after failing on the seemingly promising ones for too long.
B.1 Monte Carlo Tree Search for Deterministic Environments with Sparse Rewards
is a classic heuristic search algorithm and has achieved speculative success in difficult problems such Go (58). It uses tree policies such UCT (Upper
Monte Carlo Tree Search (MCTS) 12
26

---

Algorithm 2 Monte Carlo Tree Search for Deterministic Sparse-Reward Environment
function UCTSEARCH($))
> MCTS Search with the UCT tree policy
create the root node Uo with state 80 while within computation budget do 1i TREEPOLICY(e))
EXPAND(01)
h((())
more rollouts, BM25-based heuristic
BACKUP(4))
return a (BESTCHHILD(rD. 0))
function TREEPOLICY())
> Select node to expand using the UCT criterion
while is nonterminal do
e not fully expanded then return EXPAND(r)
else
V' BESTCCCHILD C)
return
function EXPAND(U)
>> Select a child to expand
choose a E untried actions from A(s(v)
add new child ef to e with s(v)),r(v)) WorldModel(s(1).. a) and a(u') =
return of
function BESTCHILD(E, c)
The UCT criterion for selecting nodes
Q(v")
21n (c) N(v))
> Exploit: Q(')) N(w Exploration: cV 2 (())
/ Noe")
return arg / o Mo)
+ CV
function BACKUP(r, A)
> Back propagation of the heuristic values
while is not nu null 11 do
((w) - N(v) +1
(o)) Qv)) + V parent of
to trade exploration and exploitation while deciding on which part of the search space to search more. It then evaluates the potential benefits of searching in this selected part of space, i.e., trajectories with the same prefix as the selected partial trajectory, by rolling them out in simulations. The UCT tree policy makes a balance between exploiting the search space with higher potential benefits and exploring the lesser-considered regions.
Confidence bounds for Trees) 35
However, the MCTS algorithm is designed for general planning problems and is not efficient enough for deterministic environments with sparse rewards, such as Alfworld in our experiments. We thus modify the algorithm to make it faster by making use of the additional assumptions of the environments: deterministic and sparse rewards. Specifically,, we remove the rollout procedure and instead substitute it with an optional heuristic to estimate the promise of the partial trajectories. Noticing that there is no randomness and no partial rewards in the environments, roll-outs are not cost-effective in the usage of simulations. We do not need to run a trajectory multiple times to estimate its expectations of returns. All returns of partial trajectories are zero and, once trajectory gets positive return, we find the goal state and should immediately return the trajectory from the initial state to this goal state, Computation spent on simulating the roll-outs should instead be spent on expanding the search tree with the smarter online tree policy rather than the fixed one in roll-outs. We thus remove the rollout procedure, or equivalently, set the rollout depth to zero with a hard-coded Q-value estimator, i.c., the heuristic. We adopt the classic UCT-based MCTS algorithm and replace the rollouts with a BM25-based hcuristic in our experiments, as shown in Algorithm 2 We refer the readers to the other surveys 5 for details about MCTS.
B.2 BM25 as a Heuristic for Planning in Text-based Environments
We propose to estimate the promise/goodness of a selected state with information on the similarity between the goal (specified by the user in texts) and the partial trajectory (texts that describe the path the agent has traversed from the initial state to the selected state). An example of the goal and the partial trajectory in the Alfworld environment is follows:
27

---

Mission example:
put alarmclock in desk
Partial trajectory example Goto(dest=siddetbblel)
The agentl at loc5 is now at loc20.
Pickup (ocku((obj-alarmcc (obj-alarmcloockee receptaclessfiettbblaa) The agent1 is holding alarmclock1. The alarmclock1 loc20 is (destdee)))))
The agentl at loc20 is now at locl.
Partial trajectory example
Goto(destt-microwvvel)
The agent1 at loc41 is now at loc35. Heat(obj=appleel, receptacle=microwaan The apple1 at None is now hot. (oto(dest=gaabbggcann)
The agent1 at loc36 is now loc41.
Intuitively, the assumption is that the more similar the partial trajectory is the goal, the more relevant the current state is to the goal, and thus the closer the current state to the goal states. For example, if the goal is "put hot alarmclock in desk", trajectories that involve "alarmclock", "desk". and "hot" should be more promising to exploit than other irrelevant ones.
There are various metrics and methods from the text retrieval field to estimate the text similarities 34 33 80]. We focus on the symbolic ones for efficiency issues and find BM25 (50), a state-of- the-art algorithm for web-scale information retrieval, works significantly well in our experiments. Our algorithm different from the classic BM25 algorithm in the sense that we do not have a static previously collected ttajectory/document corpus. Instead, maintain online trajectory corpus during the search, as shown in Algorithm 3 and Algorithm
43
One interesting finding in our preliminary experiments is that the saturating formulation of the term (1)+
frequency, -,-+h-(1---bb11 /.. in BM25 significantly outperforms its linear variant in TF-IDE. Without the saturating formulation, the heuristic will assign too high scores seemingly promising too long trajectories such as those that repetitively pick up "alarmclock" and go to "desk" (the concrete states are different in minor irrelevant details throughout the trajectory, so they are the same), which results in the planner to be stuck in meaningless local minimum until exhausting the budget.
C More Experimental Results
C.1 PPO for MiniGrid
1.0
empty
doorkey
0.8
uninck
felch
unlockpickup
mit
0.8
0.4
0.2
10a
10a
10c
10a
Figure 6: Performance of PPO in MiniGrid environments.
28

---

Algorithm 3 Online BM25-based Heuristic
Hyperparam: kn (default to 1.5)
Hyperparam: b (default 0.75)
N Nt(-) +0 l +0
Number of Trajectories Number of trajectories for each token Average length of trajectories
function BMI2SHEURISTII((nde mission m)
global N,Nt,
+ TRAJECTORY(()) Trajectory from the root to the node o, in text as list of tokens
m < PARSE(m)
Parse mission in text into list of tokens Update the number of trajectories for each token
loop L ESET(T)
NI(I) & NU(1) +1
L + N+1
IoxN+/r
Update the average length of trajectories b Update the total number of trajectories Initialize the heuristic value
N < N +1
loop & E SET(M)
n4 CCUNT((T,
Number of token in trajectory > Number of token t in mission m
< -COUNT(m, t)
idf < No + Nr
b. Inverse document frequency
h h + Rm idf n. +RI 61-
(K
BM25 Score
ml
return /
We evaluate PPO, as Deep RL baseline, in minigrid experiments. We use the tuned hwper-yarameeee from RL Baselines3 Zoo 488 for each environment and use the same symbolic memory state As shown in Figure 6, PPO is much more sample inefficient than ours. needs 10 10 steps to learn valid policies (not perfect though) in most environments. It cannot solve the UnlockPickup
parameters
environment even 3 x 10 steps.
C.2 PPO for Sokoban
The PPO baseline was implemented using Stable Baselines3 49 and the gym-sokoban library with 256 parallel environments, batch_size of 2048, horizon of 50 steps, and the rest default hyperparameters according to the Stable Baselines3 library (n_steps 2048, learning rate 0.0003, gamma 0.99, n_epochs 10, gae_lambda 0.95, clip_range 0.2, normalized advantage, ent_coef O, vf_coef 0.5).
The input to the network is (3, 7, 7) corresponding to the 7 by 7 Sokoban grid, with each item in the grid represent as an RGB pixel.
The policy network used was a convolutional neural network with convolutional layers with 16, 32, and 64 filters and all with kernel of size (2, 2) and stride of 1, followed by linear layer with outputs, corresponding to the 9 possible actions. All layers are separated by Rectified Linear Units (ReLu).
C.3 Dreamer V3 for Sokoban
The DreamerV3 baseline was implemented using the simple training script provided by the open- source DreamerV3 library created by Danijar Hafner (2], and the gym-sokoban library (54) with default hyperparameters and environment.
29

---

D Example of synthesized world models for Sokoban
Synthesized transition function for Sokoban
def transition(state, action):
Args:
state: set entities representing the state the environment action: the action can be "move right", "move left", "move up", move down"
Returns:
next_state: the state of the environment
unn
here we define how the player coordinates change for each action action_to_delta =
"move right": (1, 0), "move left": (-1, 0),
"move up": (0, -
"move down": (8, 1)
Here get the player and the boxes the current state 'Player')
player
cearities_by_naeesttee
boxes = get_entities__y_nne (state, 'Box')
walls get_entittee__by_naeee (state, 'Wall')
Then, calculate the player position according to the action delta_x, delta_y aattion_to_deltt(aceeee
new_player_x player.x delta_x new_player_y player.y delta_y
We check if the new player position is Wall
if get paaritie_bby_osittiie
new_player_x, new_player_y)
so, the player does not move pass
else:
If not, the player moves to the position pushed_boo get.ettiiiies.bb.ooottttoooooooe new_player_y)
new_player_x,
if pushed_box. box:
pushed_box. = pushed_ Laieeddbbxx[[ [0].x delta_x
bashed_box__y
pushed_box [0].y delta_y
Check if there is wall or other box at the pushed box destination
if get_entin cett_enttiie__b__pet by.pssbyyyysttt (boxes walls, bushed_box_5. pushed_box_y)
# If so, player and the box not move pass
else:
# If not, the box moves to the new position box[9] pushed_box
delta_x
pushed_box [0] delta_y
player.x += delta_x
player.y += delta_y
else:
player.x delta_x
player.y += delta_y
return state
Synthesized reward function for Sokoban
def reward_ deward_funccee (state, action, next_state) reward -0.1
30

---

done = False
boxes_prev prev get_ get_entities_be msaaeesstt 'Box')
targets_prev get_ erteeniiiessbby name (state, 'Target') boxes_next get_ syse_entttie__b__aaee (next_ state, 'Box') targets_next get_entities__y_nameee (next_state, 'Target') for box in boxes_next:
if any (box.xx == target.x and box.y t for target in targets_next):
not (box.x prev_box. and box.y prev_box.y for prev_box in boxes_ if any(prev_box.x == prev_target.x and prev_box.y == prav_targe for prev_target in t: for box boxes_prev:
reward +=
if any (box.x == target.x and box.y == target.y for target in targets_prev):
if not any(box.x next_box.x and box.y next_box.y for next_box boxes_next
if (ny box.x == tar__target.x and next_ box.y =
next_target.y for next_target in targets_next)):
reward - all (box.x target.x and box.y target.y for target in targets, next) for box boxes_next):
reward += 10
done True
return reward, done
E
Example of synthesized world model for Alfworld
Synthesized transition function for Alfworld
class Entity:
def __innt__sseit (self, name, loc=None, in_on=None, kwargs))
self name name
self. loc loc
self.in_on in_on self.ishot, self.iscools self. isclean None, None, None self. self.ison, self.istoggle = None, None, None self. pilf.pick self. oelfablp,,,le, self. toggleable None, None, None self helfhheat self.coolable, self.cleanable None, None, None self.isobject, self. serf.isree,,tacle, self. tr.eiiiseccptalleooeettt None, None, None
for key, value in kwargs. pesssiten itmms sstattrsself, key, value) def __eeer__seee))
attr attr_list for key in self. decf
if getattr(self, key) is not None:
value gsel,, key)
(svas,, str):
value f'"[value)")
attr_list. (err:llsttt.ppmm((((eeeettttttt
attr ', '.join(attr_lisz)
class_name wewf.entity_.typee).c.... .tatriy_ttye()ccciitaaee)(
return [ffcllas_namee((((e (lattr)))
def
(self,, other):
return all (getattrrrreff key) = gotht,, key, None) for key in self.__de deyf
def
(self)
31

---

return (itsssuuppsstteeesssff.......... tyyety)
def
return entity- type (self. name)
def ss_eentity_ttyeeseff,, etype):
return is ss__entity_tyyeesllf..oo Agent(Entity);
etype)
class
__entt__sself,, name, loc=None, holding=None, **kwargs): def
_u_ (name, loc, **kwargs)
self.holding holding def _..rer___selff
(neee=(se....ammm)) "nldddng=(selff.hodddnnn"""
/a loc=(self. loc),
class Action:
def (senf, name, **wwargs,):
for key, value in kwargs.items()
setattr(self, key, value)
def (sllf)
attr ;; (foy for key, value self._..dcct.... items()))
return _ame
def (self, other)
return allggettttrsttssss key) getattr(other, key, None) key keys)
__ def
return Srarrmmmnhssscuueeessfsssersssss (tendit)t_et__(())
class Goto(Action): def __init___seeff, entity_name,)
(nnterr))__niit___""ttttt dest=entity_nane
class Open(Action):
def
(senf entity_name,
(opper obj=ennity_name
class Close(Action): def __initt_(self, seerity_name,)
Superr)).__intt__((cclssss obj=entity_name
class Examine(Acttion)
__iniit_sselff, entity_name,) def
inttrr obj
abj=entity_name
class Use(Action):
def (seli entity_name,)
supper()._init___
obj=entity_name)
class Pickup(Action).
def __innit___eelff entityl_name, entity2_name):
intt obj=entityl_nane receptacle=entitv/.
class Put(Action):
def (slff,, entityl_name, entity2_name, cancontain,):
(nnper()__init____( obj=entityl_name name, receptaclernntityy_naas
cancontaiin=cancontein))
class Clean(Action):
__.init__sseff antityl_name, entity2_name): def
(niter obj=entityl_name,
receptaclesentttyyy_nneg
class Heat(Action):
def (self entityl_name, entity2_name))
supprr()__init__(hhettt obj=entityl_name,
receptaclesentityy__nne)
class Cool(Action):
def (_el entityl_nama entity2_name):
iunttrr).__iitt___"oooll obj=entityl_name, receptaclerentityyL.naa
def entity_tyyeenamm))
return name.stript (tt.ssrrpppe
def tenrity_type(namm. etype):
return etype.lower lower lower()..tring ().strip() == entity_.typennmm) .lower ()
def fs_receptacle, prcceptacle_ooen(stttee obj):
if obj in_on is None:
32

---

return True
receptacle get_ eeet_entit___oe sy_nameesttae, obj in_on)
if receptacle is None not receptacle.openablee return True return receptacle.
def
action_obj_name,
agent byyttiiee_b__tyyeesttee 'agent') [0]
heating_entity
bonttttyyy__ameetttt action_obj. objio____jaame
microwave_entitg get_ get_entity_by. by_naeyeenaattteeeeeeee action_ cewiooonrceptcllee.cmm name)
Ensure agent is in the correct location where the microwave is if agent.loc microwave_enti loc:
agent. loc
wewrowave_enttty...
Checking whether the agent has the object if agent.holding == heating_entity.nng name and haating_entityhheaableee Heat the object
heating_entity ishot True
heating_entity. iscool False Resetting cool
return state
pcreate_agently function creates an agent entity instance
def ureate_agent aname, loc=None, holding=None, **kwargs):
return Agent (name, loc=loc, holding=holdiing, **kwargs) Correcting and optimizing the Python code to better simulate world logic # Importing necessary Libraries
from copy import deepcopy
def get_entities_ contities_by__ypeeeeeeeeeeeeeeee (entitiee,,,,,,s,, etype):
return [entity for entity entities if entity.entity..tyee. lower etype. ()]
def rs_reccetacce (entity))
return (entity,, 'isreceptacle', False)
def move_and_heet ant_heettsttae,, obj_name, receptacle_name
Similar implementation as before is assumed pass
def get_entity_br (entiiiie,,e type_name):
return next( for entity in entities if entity.entity_.tyec)) lower () type_name. lower ()), None)
def get_entity_by_nant (entities, name):
return next( for entity in entities if _caoonicallize_nameeerrr
name) == www.cnnicalizeenaaaea (name)), None)
(nam)))
def _cwnonicalizeeeeccccccc (name):
return "t.join(nam..split.. .lower
def open_state obj obj_name)
obj get_entity_by. centitty_b_naaeeste (statte,,, (state, abj_name)
if obj and openable and not obj isopen: obj.isopen True return state def close_state(state, obj_name):
obj get_entity_by. seatiityyby_mmeesttttt obj_name)
if obj and openabee and obj isopen: obj.isopen False
return state
def move_and_put(statee obj_name, receptacle_name)) agent obj get_entity..b__nmeee (state, obj_name)
'Agent')
receptacle if obj and receptacle and obj. == agent. loc and agent.holding == abj_n
receptacle_namet
if receptacle openabbee and not ceopptaceeissopent
return state Can't t put into closed receptacle
obj loc = receptacle. loc obj.in in_on = receptacle. name hoeni.holdigg None return state
33

---

def set sreent_holdingstaateee agent_ name, obj_name):
Sets the agent's holding attribute to the object's name if possible.
agent get_ by agent_name)
obj byy_eennttyybb__ameetttttttt obj_name) if obj hijpppkuupabe::
agent.holding obj_ obj_name
obj.in_on None The object is no longer in or receptacle obj. None
def move_and__a _al_aaddccuusstattttteee_ obj_name))
Moves the agent to the object's location and then picks it up, if and pickupable "*.
agent hrt.eentttyybyyyyyestttttte obj get_ byyatttyyyynammmstaatt obj_name)
'Agent")
if obj and obj piccupbceee and obj.loc == agent. loc:
set_agent_holding(sten
agent. name, obj name)
return state
Assumed helper functions (simplified for clarity) def move_and and_use(statee obj obj_name):
Implement action logic here if needed
pass
Example function signatures, defining other required actions similarly def move aaee_gggentttttttt destination_name -amee:
agent byy.entityybby_yy((sttttttt 'Agent')
destination = get_entity_by. s_ynnemmmtttte destination_name)
if destination:
agent. loc destination. loc
return state
def open_receptacle(state,, receptacle_name):
receptacle get_entity_by. bayy_byy_nmemsttee receptacle_nan __name))
receptacle receptacle.openan openable not ieceptacle.isope.. receptacle. isopen True
return state
def create_item (name, loc=None, in_on=None, **kwargs): return Entity (name, loc=loc, in_on=in_on, **kwargs) def cool_object obj_name, receptacle_name unn
Refactored to ensure the receptacle is open before placing the object inside. unn
receptacle get_ wtt__eenttyyyyynnnleeeeee (stame,,,, receptacle_nantt
if ceeepttalle.ppenabbe and not (s_oope(rrcepptacee)
state = open_ receptacleestatee, receptacle_nan
return apply_cooliingstatee, obj_nama receptacle_name Helper Function
def is (enniiyy:::;ty;;
return eetattr(entitty 'isopen', False)
Helper function to ensure receptacle is open and updated correctly
def ppen_receptacee_ff.commm closed receptacle_names receptacle get_entity_by baaitty_byymmmmstttttt receptacle_ name) if receptacle and receptacle peceptacle.openablee and not receptacle,
aeceptacle.fsopen True
return state def apply_cooling(state,, obj_name, receptacle_name)
agent be_.entittyyby_nmmesttttttet 'Agent')
obj get_entity__yy_aace (state, obj_name)
receptacle get_entity_by ____eeammmmttteeee receptacle_nant
if obj and receptacle and cojlboebee holding obj name or (obj. in_on receptacle. and
is_open is_open(reeoptacle))
if not is_open (recepptctt)))))
state open_receptacle, wtt.cleeeeffllcodddeeeecccee receptacle_name)
obj. None
34

---

abj.in_on receptacle_nant
obj. iscoll = True
isjoishot False
agent.holding None
else:
print("Agent does not hold the object or the fridge is not open.
return state
def transition(state, action):
Perform a transition based on the action type and update the state accordingly.
Args:
state (set): Current state represented as set entities. action (Action): : Action to be executed.
Returns:
set: Updated state after the action has been performed.
unn
state = deepcopy(state) # Copy state to avoid mutation lower
action_ =
if action_type == "goto":
return move_agent(state, action.dest)
elif action_ Settonntteeteeee "open":
return open_receptaclesstate,, action.obj)
elif action_type == "close":
return close_ (rose_stateesss (state, action.obj)
elif action_type "pickup":
return move_and__p wtwpaddddpusssttttttecee action.obj)
elif action_type == "put":
return move_and_put(state, action.obj, action. elif action_type
"use":
return move_and_ _t__and__ssssttee,.. action.obj)
elif actiin_ttyst == "heat":
return move_and_ and_anatdasaatteeeeeeeaa action.obj, retion.receptacle)
elif action_type
":
return move_and_coollstarr action.obj, action.
else:
raise (alueError(f"Unnuupe
action type: (action_type)")
return state
def love_and_cooll (state, obj_name, receptacle_name)
unn
Move an object to cooling receptacle (e.g., fridge) and cool it down
Args:
state (set): The current state the environment.
obj_name (str): The name of the object to be cooled. receptacle_nant Lame (str): The name of the receptacle where the object will be cooled.
Returns:
set: The updated state after the cooling process.
agent get_entitt..b__nmeesset
'Agent')
obj get_entity_by htrpl/yyyy_mmmmmeeeee (state, obj_name)
receptacle get_entiityby__naeettaaa (stats, receptacle_name if obj and receptacle and obj.coolable:
if agent. == abj_name:
if receptacle.opensee state open_ receptacle receptacle. agent.holding None Agent releases the object abj.i__om receptacle. name Place object in/on the receptacle obj.loc = None Since it's inside the receptacle, location is not direct
and not direptacle.isopent
iscoll True
35

---

obj. ishot = False
else:
Agent must hold the object. ")
else:
print ("Cond"Coni not met: Object must be coolable and receptacle must exist.")
return state
Example reward function:
Synthesized reward function for "put a hot egg in fridge"
# put a hot egg in fridge def
etype):
return entity for entity in entities if intity.is_entity_yyee ]
def get_ bnntities_b__ooceettit
loc):
entities entity for entity in entities entity.lo loc loc ] return entities
def entity)
return name
def ess_entity.typeeee
etype):
return etype. Toype.lower( lower s).strp(() == entityy.tyee(nmm) (ooe)p(((namm)eeeor()))) def func action, next_state):
Args:
state: the state the environment
action: the action to be executed
next_state: the next state of the environment
Returns:
reward: reward of action done: whether the episode done
agent (iitterggeeeentttiee_bbeee
'agent')))
(geteenttiee next_agent next pextiitergee (get_ enwities wty.tyyeeeeetttteeeee 'agent")))
it's a put action, check the moving object is egg destination is fridge if (sinttanceeacction Put):
obj_name action.obj rec_name reciiante
is sis_entttyy.ypeeob__naaa 'egg') and ceantittyyypeerec__eee fridge') and not pext_agent.holdiig
matching_enticn [entity for entity in get_entities_bn loc next_state, (eentlooc) if entity name obj name and entity. in_on rec_name]
if matching_entityy.
# Check if the egg is hot if matching_entiiyyl0], ishot: return 1, True else:
return 0, False
else:
return 0, False
Other parts of reward function If the action is not a put action or the object is not egg or the destination is not a fridge, return 0 reward and done=False return 8, False
36

---

Prompts
We list all the prompts that are used in our experiments in this section. The functionality for each prompt is stated in its subsection name. We highlight the dynamic information yellow and the main instruction as blue. The dynamic information includes the data collected so far in the replay buffer and the codes synthesized far by previous LLM calls.
F1 Initializing the transition function
asks LLMs to generate transition function (s, a) - following the code template to model seven uniformly randomly sampled experience data (s, a, s') in the replay buffer.
You are robot exploring in an object-centric environment. Your goal is to model the logic of the world in python. You will be provided experiences in the format of (state, action, next_state) tuples. You will also be provided with a short natural language description that briefly summarizes the difference between the state and the next state for each (state, next_state ,) pair. You need to implement the python code to model the logic of the world,
in the provided experiences. Please follow the template to implement the code. The code needs to be directly runnable on the state and return the next state in python as provided in the experiences.
You need to implement python code model the logic of the world as seen in following experiences:
The action "toggle" transforms the state from
Wall(0, 0) Wall(1, 0) ; Wall(2,
Wall(3,
Wall(4, 0) ;
Wall(0, 1) empty ;
Wall 1) empty
Wall(4, 1) ;
Wall(0, 2) Agent(1, 2, direction= (0, -1), carrying=None) ; Door (2, 2, empty
colorsyellow, state=locked)
(4, 2)
Wall(0, 3)
Key(1, 3, colorsyellow) ;
Wall(2, 3) Goal(3, 3)
Wall(4, 3)
Wall 4) ...
Wall(1, 4) ; Wall(2, 4) Wall(3, 4) Wall(4, 4)
to
Wall(0, 0) ; Wall(1, 0) Wall(2, 8) Wall(3, 0) Wall(4, 0) ;
Wall(0, 1) empty
Wall 1) empty
Wall(4, 1)
Wall(0, 2)
Agent(1, 2, direction= -1), carrying=None) ; Door (2, 2,
colorsyellow, state=locked) K1y(1, 3, colorsyellow) ;
empty
Wall(4, 2)
Wall(e, 3)
Wall(2,
Goal(3, 3)
Wall(4, 3)
Wall(0, 4)
Wall(1, 4) ; Wall(2, 4) Wall(3, 4) Wall(4, 4)
The difference is n..
Nothing happened ...
The action "toggle" transforms the state from <*.
Wall(0, 0) Wall(1, 0) Wall(2, 0)
Wall(3, 8)
Wall(4, 0) ;
Wall(0, 1) empty
Wall(2, 1) empty
Wall(4, 1)
Wall(0, 2) Agent(1, direction= 1), Norrying=None) Door (2, 2,
color=yellow, state=locked)
empty
Wall(4, 2)
Wall(0, 3)
Key(1, 3, colorsyellow) ;
Wall(2, 3) Go,l(3, 3)
Wall(4, 3)
Wall(0, 4) Wall(1, 4) ; Wall(2, 4)
Wall(3, 4) Wall(4, 4)
to
37

---

(all 0)
Wall(1, 0) ;
Wall(2, 0) Wall(3, 0) Wall(4, 8)
Wall(e, 1) empty ; Wall(2,
empty Agent(1, 2, direction= 1), carrying=None) Door 2,
Wall(4, 1) ;
Wall(0, 2)
colorsyellow, state=locked)
empty Key(1, 3, color=yellow) ;
Wall(4, 2)
Wall(e, 3)
Wall(2, 3) Goal(3, 3)
Wall(4, 3) Wall(1, 4) ; Wall(2, 4) Wall(3, 4) Wall(A, 4)
Wall(e, 4)
The difference is
Nothing happened au.
The action "turn right" transforms the state from
Wall(e, 0) Wall(1, 9) ; Wall(2,
Wall(3, Wall(4, 0) ; Wall 1) empty Wall(0, 2) Agent(1, 2, direction=11, 0), carrying=None) Door(2, 2,
Wall(e, 1) empty ;
Wall(4, 1) ;
colorsyellow, state=locked)
empty Key(1, 3, coloreyellow) ;
Wall(4, (4, 2) Wall(2, 3) Goal(3, 3) ;
Wall(0, 3)
Wall(4, 3)
Wall(e, 4) Wall(1, 4) ; Wall(2, 4) Wall(3, 4) Wall(4, 4) <.. to ...
Wall(G, 0) ; Wall(1, 0) Wall 8) Wall(3, 8) Wall(4, 0) ;
Wall(0, 1) empty
Wall 1) empty
Wall(4, 1)
Wall(0, 2)
Agent(1, 2, direction= 1),
carrying-None) Door (2, 2,
coloræyellow, state=locked)
empty Key(1, 3, colorsyellow) ;
(4, 2)
Wall(o, 3)
Wall(2, 3) Goal(3, 3)
Wall(4, 3)
Wall(0, 4)
Wall(1, 4) ; Wall(2, 4)
Wall(3, 4) Wall(4, 4)
The difference is m..
The agent (d1rection=1,, 0)) at pos (1, 2) becomes an agent (direction= 1)).
The action "turn left" transforms the state from <..
Wall(e, 0) Wall(1, 0) ; Wall(2, 0) Wall(3, 0)
Wall(4, 8)
Wall(0, 1) empty
Wall(2, Wall(0, 2) Agent(1, 2, d-rection=--1,, 0), corrying=None) ; Door (2, 2, colorsyellow, state=locked)
empty
Wall(4, 1) ;
empty Key(1, 3, colorsyellow) ; Wall(2, 3) Goal(3, 3)
Wall(4, 2)
Wall(e, 3)
Wall(4, 3)
Wall(e, 4) ...
Wall(1, 4) ; Wall(2, 4) Wall(3, 4) Wall(4, 4)
to
wall(e, 0)
Wall(1, 9) ; Wall(2,
Wall(3, 0) Wall(4, 8) ;
Wall(0, 1) empty ;
Wall(2, 1) empty
Wall(4, 1) ;
Wall(0, 2) Agent(1, 2, direction= 1), Norrying=None) Door 2,
color=yellow, state=locked)
empty empty Key(l, 3, colorsyellow) ;
Wall 2)
Wa, 3)
Wall(2, 3) Goal(3, 3)
Wall(4, 3)
Wall(0, 4)
Wall(1, 4) ; Wall(2, 4) Wall(3, 4) Wall(4, 4)
The difference is
38

---

The agent (direction=-11, 0)) at pos (1, 2) becomes an agent (direction= 1)). nun
The action "turn right" transforms the state from
Wall(e, 9) Wall(1, 0) ; Wall(2,
Wall(3, 0)
Wall(4, 0) ;
Wall(e, 1) empty
Wall(2, Wall(0, 2) Agent(1, 2, direction= 1), carrying=None) Door 2,
empty
Wall(4, 1) ;
colorsyellow, state=locked)
empty Key(1, 3, colorsyellow) ;
Wall (4, 2) Wall(2, 3) Goal(3, 3)
Wall(e, 3)
Wall(4, 3)
Wall(o, 4) Wall(1, 4) ; Wall(2, 4) ...
Wall(3, 4) Wall(4, 4)
to
wall(e, 0) ; Wall(1, 9) ; Wall(2,
Wall(3,
Wall(4, 0) ;
Wall(e, 1) empty ;
Wall 1) empty
Wall(4, 1)
Wall(0, 2)
Agent(1, 2, d-rection=-11, 0), carrying=None) ; Door 2,
colorsyellow, state=locked)
empty Key(1, coloreyellow) ;
Wall(4, 2)
Wall(0, 3)
Wall(2, 3) Goal(3, 3)
Wall(4, 3)
Wall(0, 4) Wall(1, 4) ; Wall(2, 4)
Wall(3, 4) Wall(4, 4)
<ec
The difference is nun
The agent (direction= 1)) at pos (1, 2) becomes an agent (direction=--1, 8)).
The action "turn right" transforms the state from <*.
Wall(e, 0) Wall(1, 0) ; Wall(2, 0) Wall(3, 0) Wall(4, 8)
Wall(0, 1) empty
Wall(2, 1) empty
Wall(4, 1) ;
Wall(0, 2)
Agent(1, direction= (-1, carrying=None) ; Door (2, 2, colorsyellow, state=locked)
empty Key(1, 3, colorsyellow) ;
Wall(4, 2)
Woll(o, 3)
wall(2, 3) Goal(3, 3)
Wall(4, 3)
Wall(e, 4) Wall(1, 4) ; Wall(2, ...
Wall(3, 4) Wall(4, 4)
to
Wall(0, 0)
Wall(1, 0) ; Wall(2, 0) Wall(3, 8) Wall(4, 8)
Wall(0, 1) empty
Wall(2, 1) empty
Wall(4, 1) ;
Wall(0, 2) Agent(1, 2, direction= -1), corrying=None) ; Door (2, 2, colorsyellow, state=locked) Key(1, 3, colorsyellow) ;
empty
Wall(4, 2)
Wa, 3)
Wall(2, 3) Goal(3, 3)
Wall(4, 3)
Wall(e, 4)
Wall(1, 4) ; Wall(2, 4) Wall(3, 4) Wall(4, 4)
The difference is
The agent (-i, 0)) pos (1, 2) becomes agent (direction= -1)). nuи
The action "turn left" transforms the state from ...
Wall(0, 0) Wall(1, 0) ; Wall(2, 8) Wall(3, 8)
Wall(4, 0)
Wall(0, 1) empty ; Wall 1) empty Wall(4, 1)
Wall(0, 2) Agent(1, (1, 2, direction= 1), carrying-None) Door (2, 2,
colorsyellow, state=locked)
empty
Wall(4, 2)
39

---

Wo, 3) Wall(4, 3) Wall(e, 4) Wall(1, 4) ; Wall(2, Wall(3, 4) Wall(4, 4) ; ... to
Key(1, 3, colorsyellow) ;
Wall(2, 3)
Goal(3, 3)
Wall(0, 0)
Wall(1, 0) ;
Wall(2,
Wall(3, 0) Wall(4, 0) ;
Wall(e, 1) empty
empty Wall(2, carrying=None) Door 2,
Wall(4, 1) ;
Wall(0, 2) Agent(1, 2, direction=11, 0),
empty colornyellow, state=locked) empty Key(L, 3, colorsyellow) ;
Wall (4, 2) Wall(2, 3) Goal(3, 3)
Wall(e, 3)
Wall(4, 3) Wall(1, 4) ; Wall(2, 4) Wall(3, 4) Wall(4, 4)
Wall(e, 4)
644
The difference is
The agent (direction= 1)) at pos (1, becomes agent (direction= (1, 8)).
Please implement code to model the logic of the world demonstrated by the experiences. Here the template for the transition function. Please implement the transition function following the template. The code needs to be directly runnable an the inputs of (state, action) and return the next state in python provided in the experiences.
..c
class Entity:
def (self, y, **kwargs):
self = classs
self.x
self.y =
for key, value kwargs.items()) ()
setattr(self, key, value)
def
attr (sel__
(.je) for key, value in self._.dict items () key not ('name', 'x', 'y'))
if attr: return f"(self. namef..nmmmmaaemm ([self.x), (self.y), (attr))"
else: return f" (fself.x), {self.y))" def (seff, other):
return (sel,, key) (othr,, key, None) for key in self setf.__dict_..kes((()) def __ahahh_sssel))== return hatturnnnaaccpppessstteetetllt..
items
class Agent(Entity): pass
class Key(Entity): pass
class Door(Entity)) pass
class Goal(Entiyy): pass
class Wall(Entity): pass
class Box(Eitity): pass
class Ball(Entity): pass
class Lava(Ent)ty): pass
def get_ tentiitees_b__mmmeeentttt
name):
return entity for entity in entities if entity.name name ]
def get_entities_be horriiiee_bb_oositionnnetttt
x, y):
return entity for entity in entities if entity.x == x and entity.y =
y
def transition(state, action):
unn
Args:
state: the state of the ronment
action: the action to be executed
40

---

Returns:
next_state: the next state of the environment
raise NotImplementedError
Please implement code to model the logic of the world demonstrated by the experiences. Please implement the code following the template. Feel free to implement the helper functions you need. You can also implement the logic for difference actions in different helper functions. However, you must implement the transition function as the main function to be called by the environment. The code needs to be directly runnable the inputs as state, action) and return the next state in python as provided in the experiences. Let's think step by step.
F.2 Initializing the reward function
It asks LLMs to generate reward function (s, a, s') (r, for the mission following the code template to model seven uniformly randomly sampled experience data (s, 0, ;,, d) in the replay buffer for the mission c.
You are a robot exploring in an object-centric envi ronment. Your goal is to model the logic of the world in python. You will be provided experiences in the format of (state, action, vex__state, reward, done) tuples. You will also be provided with short natural language description that briefly summarizes the difference between state, next_state) pair. You need to implement the python code to model the logic of the world, as seen in the provided experiences. Please follow the template to implement the code. The code needs to be directly runnable the (state, action, next_state) tuple and return the (reward, done) tuple python as provided in the experiences.
state and the next state for each
You need to implement python code to model the logic of the wor ld as seen in the following experiences for mission "use the key to open the door and then get to goal":
The action "turn left" transforms the state from
Wall(e, 0) Wall(1, 0) ; Wall(2,
Wall(3,
Wall(4, 8) ;
Wall(0, 1) empty ;
Wall 1) empty
Wall(4, 1) ;
Wall(0, 2) Agent(1, 2, direction= -1), carrying=None) ; Door (2, 2, empty Key(1, 3, colorsyellow) ; Wall(2, 3) Go,l 3)
colorsyellow, state=locked)
Wall(4, 2)
Wall(0, 3)
Wall(4, 3)
Wall(0, 4) Wall(1, 4) Wall(2, 4) Wall(3, 4) Wall(4, 4) ... to ...
Wall(0, 0)
Wall(1, 0) Wall(2, 8) Wall(3,, 8) Wall(4, 0)
Wall(0, 1) empty ; Wall 1) empty
Wall(4, 1)
Wall(0, 2)
Agent(1, 2, direction= 0), carrying=None) ; Door (2, 2,
colorsyellow, state=locked)
empty Wall(2,
Wall(4, 2)
Wall(e, 3)
Key(1, 3, colorsyellow) ;
Goal(3, 3)
Wall(4, 3)
Wall(0, 4)
Wall(1, 4) Wall(2, 4)
Wall(3, 4) Wall(4, 4)
The difference is any
41

---

The agent (direction=(0, -1)) at pos (1, 2) becomes an agent (direction= (-1, 0)). nu"
the returned reward is 0.0 and the returned done is False
The action "toggle" transforms the state from ...
Wall(e, 0) Wall(1, 0) ; Wall(2,
Wall(3,
Wall(4, 8) ;
Wall(0, 1) empty ;
(2, 1) empty
Wall(4, 1) ;
Wall(0, 2) Agent(1, 2, direction= (0, -1), carrying=None) ; Door (2, 2,
colorsyellow, state=locked)
empty
Wall(4, 2) wall(2, 3) Goal(3, 3)
wall(e, 3)
Key(1, 3, coloreyellow) ;
Wall(4, 3) Wall(0, 4)
Wall(1, 4) ; Wall(2, 4)
Wall(3, 4) Wall(4, 4)
to ...
Wall(G, 0)
Wall(1, 0) Wall(2, 8) Wall(3, 8) Wall(4, 8)
Wall(0, 1) empty
Wall 1) empty
Wall(4, 1)
Wall(0, 2)
Agent(1, 2, direction= (0, -1), carrying=None) ; Door (2, 2,
color=yellow, state=locked) Key(1, 3, colorsyellow) ;
empty
Wall(4, 2)
Wall(e, 3)
Wall(2, 3) Goal(3, 3)
Wall(4, 3)
Wall(e, 4) Wall(1, 4) Wall(2, 4) Wall(3, 4) Wall(4, 4)
The difference is n.. Nothing happened
the returned reward is 0.0 and the returned done is False
The action "turn right" transforms the state from
Wall(e, 0) Wall(1, 0) ; Wall(2, 0) Wall(3, 8)
Wall(4, 8)
Wall(0, 1) empty
Wall(2, Wall(0, 2) Agent(1, 2, direction= (1, 0), carrying=None) Door (2, 2,
empty
Wall(4, 1) ;
colorsyellow, state=locked) Wall(0, 3)
empty
Wall(4, 2)
Key(1, 3, colorsyellow) ;
Wall(2, 3) Goal(3, 3)
Wall(4, 3)
Wall(e, 4) ... to
Wall(1, 4) ; Wall(2, 4) Wall(3, 4) Wall(4, 4)
Wall(0, 0)
Wall(1, 0) ; Wall(2, 0) Wall(3, 0) Wall(4, 0) ;
Wall(0, 1) empty ;
Wall 1) empty
Wall(4, 1) ;
Wall(0, 2) Agent(1, 2, direction= 1), carrying-None) Door (2, 2, empty Wall(0, 3) Key(1, 3, colorsyellow) ;
colorsyellow, state=locked)
(4, 2)
Wall(2, 3) Goal(3, 3)
Wall(4, (a, 4) ...
Wall(1, 4) ; Wall(2, 4) Wall(3, 4) Wall(4, 4)
The difference is nun
The agent (direction= 0)) at pos (1, becomes an agent (direction= (0, 1)).
the returned reward is 0.0 and the returned done is False
The action "nothing" transforms the state from .*.
Wall(0, 0) Wall(1, 0) Wall(2, 0) Wall(0, 1) empty Wall(2, 1) empty
Wall(3, 8)
Wall(4, 0)
Wall(4, 1) ;
42

---

Woll(0, 2) Agent(1, 2, direction= -1), carrying=None) ; Door (2, 2, color=yellow, state=locked)
empty W4, 2) Wall(e, 3) Key(1, 3, coloreyellow) ; Wall(2, Goal(3, 3)
Wall(4, 3) Wall(1, 4) ; Wall(2, 4)
(oll(0, 4)
Wall(3, 4) Wall(4, 4)
to ...
Wall(0, 0)
Wall(1, 0) Wall 0) Wall(3, 8) Wa 8)
Wall(0, 1)
empty
Wall 1) empty
Wall(4, 1)
Wall(e, 2)
Agent(1, 2, direction= -1), carrying=None) ; Door (2, 2,
colorayellow, state=locked)
empty
Wall(4, 2)
Woll(o, 3)
Key(1, 3, colorsyellow)
Wall(2, 3) Goal(3, 3)
Wall(4, 3)
Wall(0, 4)
Wall(1, 4) Wall(2, 4) Wall(3, 4) Wall(4, 4)
The difference is nun
Nothing happened
the returned reward is 0.0 and the returned done is False
The action "drop" transforms the state from
Wall(0, 0) Wall(1, 0) ; Wall(2, 0) Wall(3, 8) Wall(4, 8)
Wall(0, 1) empty
Wall(2, 1) empty
Wall(4, 1) ;
Wall(o, 2) Agent(1, 2, direction= -1), carrying=None) ; Door (2, 2,
colorsyellow, state=locked) Wall(0, 3)
empty Key(1, 3, colorsyellow) ;
Wall(4, 2)
Wall(2, 3) Goal(3, 3)
Wall(4, 3)
Wall(e, 4) ...
Wall(1, 4) ; Wall(2, 4) Wall(3, 4) Wall(4, 4)
to
wall(0, 0)
Wall(1, 0) ; Wall(2,
Wall(3, 0) Wall(4, 0) ; Wall(2, 1) empty
Wall(0, 1) empty ;
Wall(4, 1) ;
Wall(0, 2)
Agent(1, 2, direction= -1), carrying=None) ; Door (2, 2, empty
colorsyellow, state=locked)
(4, 2)
Wall(0, 3)
Key(1, 3, colorsyellow) ;
Wall(2, 3) Goal 3)
Wall(4, 3)
Wall(0, 4)
Wall(1, 4) ; Wall(2, 4) Wall(3, 4) Wall(4, 4)
<cc
The difference is
Nothing happened
auy
the returned reward is 0.0 and the returned done is False
The action "turn left" transforms the state from ...
Wa, 0)
wall(1, 0) ; Wall(2, 0)
Wall(3, 0)
Wall(4, 0)
Wall(0, 1) empty
Wall 1) empty carrying=None) ; Door (2, 2,
Wall(4, 1)
Wall(e, 2)
Agent(1, direction= (-1, colorsyellow, state=locked)
empty
(4, 2)
Wall(0, 3)
Key(1, 3, colorsyellow) ;
Wall(2, 3) Goal(3, 3)
Wall(4, 3)
Wall(e, 4) Wall(1, 4) ; Wall(2, 4) ... to
Wall(3, 4) Wall(4, 4)
...
wall(e, 0)
Wall(1, 0)
Wall(2, 0)
Wall(3, 8)
Wall(4, 0)
wall(0, 1) empty
Wall(2, 1) empty
Wall(4, 1) ;
43

---

Woll(0, 2) Agent(1, 2, direction= 1), carrying=None) Door (2, 2, color=yellow, state=locked)
empty (4, 2)
Wall(e, 3) Key(1, 3, color=yellow) ; Wall(2, 3) Goal(3, 3) ; Wall(4, 3) Wall(1, 4) Wall(2, 4) Wall(3, 4) Wall(4, 4)
Wall(0, 4)
The difference is nun
The agent (direction=-1,, 0)) at pos (1, 2) becomes an agent (direction= 1)).
the returned reward is 0.0 and the returned done is False
The action "turn right" transforms the state from
Wall(0, 0) Wall(1, 0) ; Wall
Wall(3, 0) Wall(4, 0) empty
Wall(e, 1) empty ;
Wall Wall(e, 2) Agent(1, 2, direction= (-11, 0), carrying=None) ; Door (2, 2, colorsyellow, state=locked)
Wall(4, 1) ;
empty Key(1, 3, colorsyellow)
WaLl(4, 2)
Wall(e, 3)
Wall(2, 3) Goal(3, 3)
Wall(4, 3)
Wall(e, 4) ... to Wall(0, 0)
Wall(1, 4) ; Wall(2, 4) Wall(3, 4) Wall(4, 4)
Nall(1, 9) ; Wall
Wall(3,
Wall(4, 0) ;
Wall(G, 1) empty ;
Wall 1) empty
Wall(4, 1) ;
Wall(0, 2) color=yellow, state=locked)
Agent(1, 2, direction= -1), carrying=None) ; Door (2, 2, empty Key(1, 3, colorayellow) ;
W4ll(4, 2)
(all(0, 3) ;
Wall(2, 3) Goal(3, 3)
Wall(4, 3) (e, 4) Wall(1, 4) ; Wall(2, 4) Wall(3, 4) Wall(4, 4)
...
The difference is
The agent (direction=-11, 0)) pos (1, 2) becomes an agent (direction=(0, -1)).
the returned reward is 0.0 and the returned done is False
Please implement code to model the logic the world as demonstrated by experiences, Here is the template for the reward function. Please implement the reward function following the template. The code needs to be directly runnable the inputs of (state, action, next_state) and return (reward, done) python as provided the experiences.
class Entity:
def _self, y, **kwargs):
self.name self._. name self.x =
self.y =
for key, value in kwargs.items() (sta,, key, value) def (self=
attr = (key)) for key, value in dict... items () if key not in ('name', 'x', 'y'))
if attr: return f"(self. ([self.x), (self.y), (attr))"
else: return f"(selff.nmme ([self.x), (self.y])"
def _seqf, other):
return (selt key) == (othrr, key, None) for key in diit
44

---

def (selfs
return (asttst wicel..f..dddt_._.. items
class Agent(Entity): pass
class Key(Entíty): pass
class Door(Entity): pass
class Goal(Entity): pass
class Wall(Entity): pass
class :ox(Ettity): pass
class (all(Entity): pass
class Lava(Entity): pass
def se__eentiiees.b__mmmmeeeteeeee
name):
return entity for entity in entities if entity, ]
def setriiies_b_.ossttoneee return entity for entity in entities entity.x == x entity.y
x, y):
]
def reward_func(state,, action, next_state):
Args:
state: the state of the environment
action: the action to be executed
next_state: the next state of the environment
Returns:
reward: the reward of the action
done: whether the episode is done
raise NotImmolmenttddFrrr
Please implement code to modell the logic of the world as demonstrated by the experiences. Please implement the code following the template. You must implement the reward_ the environment. The code needs be directly runnable on the inputs as state, action, next_state) and return (reward, done) in python as provided in the experiences. Let's think step by step.
function the main function to be called by
F.3 Refining the transition function
It asks LLMs to refine a partially correct transition function by providing it with data point that fails to model well few other data points that it succeeds in modelling. We also provide the wrong prediction by the partially correct code or the error message during execution.
You are a robot exploring in an object-centric environment. Your goal is to model the lagic of the world in python. You have tried it before and came up perfect. They
with partially carrect solution. However, model the logic for some experiences but failed for others. You improve your code to model the logic of the world for all the experiences.
The new code needs to be directly runnable the (state, action) pair and return the next state in python provided in the experiences.
Here is the partially correct solution you came up with. It can model the logic for some experiences but failed for others. You need to improve your code to model the logic of the world for all the experiences. The new code needs be directly runnable the (state, action) pair and return the next state in python as provided in experiences.
class Entity:
def i_int___ (self, x, y, **kwargs):
45

---

self.name name = class...class..... name............. self.x =
self.y
for key, value kwargs.items
setattr(self, key, value)
def (selr_
attr = ".join((f:(e)))<vaaaue) for key, value in delt._.
items () key not in ('name', 'x', 'y'))
attr: return ["[self. ([self.x), (self.y), (attr))"
else: return f"[self.names ([self.x), (self.y))"
def (selff other):
return all (self, key) getattr key, None) for key decf.__.iict___essss(
def ___hash__seelff
return r Setturnnnnastttttteeeeefeeeeeer_
items ())))
class Agent(Entity): pass
class Key(Entity): pass
class (Eot(En:;yy): pass
class Goal(Entity): pass
class Wall(Entity): pass
class (ox(Eitiyy): pass
class (Ent(En))yy); pass
class Lava(Entíty): pass
import copy
def get_entities_ (orrries____ammmeeetttttee name):
return entity for entity in entities if entity.name ]
def get_entities_be snntiiteesbb_ppsitiooneeee
x, y):
return entity for entity in entities if entity.x == x and entity.y ==
] def transition(state, action):
next_state copy. (sttoopeecopytttaee)
agent get_entities__y_nn (next_state, "Agent')(0 Determine agent's next position based on action if action == 'ovee forward':
next_pos (agent.x agent.dfrectton(0), agent.y agent.directiin
[1])
Check if the next position isn't a wall if not iny(isinstanceeenttty Wall) for entity in get_entities_ie sis.aaiiiiiss.booossiiiontttttae If agent is in front of door and has the right color key, unlock door if amytisinstanceeentityy
*next_pos)):
Door) and entity.color == agent.
carrying.color for entity in get_ www.etiies.bb.ossttiii.. if action == 'togge'':
*next_pos))
agent.carrying None Drop the key
else:
agent.x, agent.y = next_pos # Move forward
elif action == 'pick up":
Pick up key if there is key at the agent's position for entity renooe break
agent.x, agent.y)
if (sinstance Key): agent.carrying entity
elif action == 'drop':
# Drop the key at the agent's position if the agent is carrying
key
agent.carrying not None:
dropped_key Key agent.y, coloragent.carrying..ood. rext_state.appent www.appeedddpppppdpe (dropped_key) agent.carrying = None
elif action in ['turn left', 'turn right"]:
46

---

Existing code for turn left/right here pass
elif action 'toggle':
Existing code for toggle here pass return next_state
The given code cannot model the logic of the world for all the experiences. Here are some experiences that the code have successfully modeled.
The action "toggle" transforms the state from ...
Wo,l(0, 0)
Wall(1, 6) ; Wall 0) Wall(3, 0) Wall(4, 8)
Wall(e, 1)
empty
Wall(2, 1) empty
Wall(4, 1)
Wall(0, 2)
Agent(1, direction= (8, -1), carrying=None) ; Door (2, 2, colorsyellow, state=locked)
empty
Wall(4, 2)
(a,l(0, 3)
Key(1, 3, colorsyellow)
Wall(2, 3) Goal 3)
Wall(4, 3)
Wall(e, 4) ... to ...
Wall(1, 4) ; Wall(2, 4) Wall(3, 4) Wall(4, 4)
Wall(e, 0)
Wall(1, 0) ; Wall(2, 0) Wall(3, 0) Wall(4, 0) ;
wall(0, 1)
empty
empty Wall(2, Norrying=None) ; Door 2, W4, 2)
Wall(4, 1) ;
Wall(e, 2)
Agent(1, 2, direction= -1),
colorsyellow, state=lacked)
empty Key(1, 3, colorsyellow)
(a, 3)
Wall(2, 3) Goal(3, 3)
Wall(4, 3)
Wall(e, 4)
Wall(1, 4) ; Wall(2, Wall(3, 4) Wall(4, 4)
<<c
The difference is
Nothing happened n.n
The action "drop" transforms the state from
Wall(0, 0) Wall(1, 0) ; Wall(2,
Wall(3,
Wall(4, 0) ;
Wall(0, 1) empty ;
Wall(2, 1) empty
Wall(4, 1) ;
Wall(0, 2) Agent(1, 2, direction= (8, -1), carrying=None) ; Door 2,
color=yellow, state=locked)
empty
(4,l 2)
Wall(a, 3)
Key(1, 3, colorsyellow) ;
Wall(2, 3) Goal(3, 3)
Wall(4, 3)
Wall(e, 4)
Wall(1, 4) ; Wall(2, 4) Wall(3, 4) Wall(4, 4)
<.. to Wall(0, 0)
Wall(1, 0) Wall(2, 8) Wall(3, 0) ; Wall(4, 0) ;
Wall(0, 1) empty ; Wall 1) empty
Wall(4, 1) ;
Wall(0, 2) Agent(1, 2, direction= -1), carrying=None) ; Door 2,
colorsyellow, state=locked)
empty Key(1, 3, colorsyellow) ;
W4, 2)
Wall(e, 3)
wall(2, 3) Goal(3, 3) ;
Wall(4, 3)
Wall(0, 4)
Wall(1, 4) ; Wall(2, 4) Wall(3, 4) Wall(4, 4)
The difference is any
Nothing happened
47

---

The action "nothing" transforms the state from
Wall(e, 0) Wall(1, 9) ; Wall(2,
Wall(3,
Wall(4, 8) ;
Wall(a, 1) empty ;
W2,l 1) empty
Wall(4, 1) ;
Wall(0, 2) Agent(1, 2, direction= -1), carrying-None) ; Door (2, 2, colorsyellow, state=locked)
empty Wall(4, Key(1, 3, coloreyellow) ; wall(2, 3) Goal(3, 3)
Wa,l(4, 2)
Wall(e, 3)
Wall(e, 4) ...
Wall(1, 4) ; Wall(2, 4)
Wall(3, 4) Wall(4, 4)
to ...
Wall(o, 0)
Wall(1, 0) Wall(2, 8) Wall(3, 8)
Wall(4, 8)
Wall(0, 1) empty ; Wall 1) empty
Wall(4, 1)
Wall(0, 2)
Agent(1, 2, direction= (0, -1), carrying=None) ; Door (2, 2, coloræyellow, state=locked)
empty
Wall(4, 2)
Wall(e, 3)
Key(1, 3, coloreyellow) ;
Wall(2, 3) Goal(3, 3)
Wall(4, 3)
Wall(e, 4)
Wall(1, 4) ; Wall(2, 4) Wall(3, 4) Wall(4, 4)
The difference is n..
Nothing happened
Here is example of experiences that the code failed to model.
The action "turn left" should transform the state from
Wall(0, 0) Wall(1, 0) ; Wall(2, 0) Wall(3, 0)
Wall(4, 0)
Wall(0, 1) empty ;
Wall(2, 1) empty
Wall(4, 1) ;
Wall(G, 2) Ag,nt(1, 2, direction= 1), carrying=None) Door 2,
colorsyellow, state=locked)
empty Key(1, 3, colorzyellow) ;
Wall(4, 2)
Wall(e, 3)
Wall(2, 3) Goal(3, 3)
Wall(4, 3)
Wall(e, 4) ... to
Wall(1, 4) ; Wall(2, 4) Wall(3, 4) Wall(4, 4)
Wall(0, 0)
Wall(1, 0) Wall(2,
Wall(3, 8) Wall(4, 8) ;
Wall(0, 1) empty ; Wall(2, 1) empty
Wall(4, 1)
Wall(0, 2)
Agent(1, 2, direction=1,, 0), carrying=None)
Door (2, 2,
colorsyellow, state=locked)
empty Key(1, 3, colorzyellow) ;
Wall(4, 2)
Wall(0, 3)
Wall(2, 3) Goal(3, 3)
Wall(4, 3) Wall(1, 4) ; Wall(2, 4) Wall(3, 4) Wall(4, 4)
Wall(0, 4)
<ec
The difference is nun
The agent (direction= 1)) at pas (1, 2) becomes an agent (direction= 8)).
However, the implementation is wrong because it returns state <*.
W0,l(0, 0)
Wall(1, 0) ; Wall 0) Wall(3, 8) Wall(4, 0) ;
Wall(0, 1) empty
Wall 1) empty
Wall(4, 1)
Wall(0, 2) Agent(1, 2, direction= 1), carrying=None) Door (2, 2,
colorsyellow, state=locked)
empty
Wall(4, 2)
Wall(0, 3)
Key(1, 3, colorsyellow)
Wall(2, 3) Goal(3, 3)
Wall(4, 3)
Wall(e, 4) Wall(1, 4) ; Wall(2, 4) Wall(3, 4) Wall(4, 4)
48

---

For this failed experience, do you know what is different between the true transitions from the environment and the predictions from the code? Do you know why the environment behaves in this way? Do you know why the code behaves differently from the environment? which part of code causes the problem? How to fix it? Please improve your code to model the logic of the world for all the experiences, accordingly. Please implement the code following the template. Feel free to imp any helper functions you need. You also implement the logic for difference actions different helper functions. However, you must implement the transition function as the main function to be called by the environment. The code needs to be directly runnable the (state, action) tuple and return the state in python provided in the experiences. If code too long, try to refactor to shorter.
F.4 Refine the reward function
asks LLMs refine a partially correct reward function by providing it with a data point that fails to model as well as few other data points that it succeeds in modelling. We also provide the wrong prediction by the partially correct code the error message during execution.
You are a robot exploring in an object-centric environment. Your goal is to model the logic of the world in python. You have tried it before and came with one partially correct solution. However, it is not perfect. They model the logic for improve your
experiences but failed for others. You need to model the logic The new code needs to be directly runnable on the (state, action, next_state ) tuple and return the (reward, done) tuple in python provided in the experiences.
world for all
experiences,
Here is the partially correct solution you came up with for mission "use the key to open the door and then get to the goal". It can model the logic for some experiences but failed for others. You need to improve your code to model the logic of the world for all the experiences. The code need to be directly runnable on the (state, action, next_ state) tuple return the (reward, done) tuple in python as provided in the experiences.
class Entity:
def (self, y, **wwargs):
self.name name.
self.x
self.y =
for key, value in kwargs. iiargs.items(0 sstat key, value) attr _serr_
def
(jjoin(("(ke)))=vvaeee for key, value in self.__dict items () if key not in ('name', 'x', 'y'))
if attr: return f"(self.nan name. ((self.x), (self.y), (attr))"
else: return ("(self.name) ((selfxx) (self.y))" def __eqq_ other):
return all(getattr key) == (othr,, key, None) for key in self self.__dicct__kks(((() def __aassh__ee__)) return hatpraaaaaattteeterrdreeee...... dict_ items(m)))
class Agent(Entity): pass
class Key(Entity): pass
class Door (Entity): pass
class GonttEnt) pass
class Wall(Entity): pass
49

---

class (ox(Entity): pass
class :all(Ent;ty); pass
class Lava(Entity): pass
def get_entities_ gettentities.bb_oosiiiin btw.bbb_oossiiiooeeetttees x, y):
return entity for entity in entities entity.x == x and entity.y == ] def ffnnard_ action, next_state);
state_ set(state)
next_state_set set(next_state)
agent = [e for e in state_set if isinstance(e, Agent)][0] agent_next [e for in next_state_set if isinstance(e, Agent)] on_goal any (nytssiitttaceeenttty Goal) for entity in www.eetiieeee.ppppsiiiioooeeeeee done on_goal if state_set == next_state_set:
agent_next.x, agent_next.y))
reward -0.1 # Small negative reward for no-op actions
encourage faster solution
elif done:
reward 1.0 # Reward for reaching the goal
else:
reward 0.0 No reward in other cases return reward, done
<..
The given code cannot model the logic of the world for all the experiences. Here experiences that the code has successfully modeled.
The action "turn right" transforms the state from
Wall(e, 0) Wall(1, 9) ; Wall(2,
Wall(3, Wall 1) empty
Wall(4, 0) ;
Wall(0, 1) empty ;
Wall(4, 1) ;
Wall(0, 2)
Agent(1, 2, direction= (-1, 0), carrying=None) ; Door 2,
colorsyellow, state=locked)
empty Wall(2, 3) Goal(3, 3)
W4, 2)
Wall(0, 3)
Key(1, 3, colorsyellow) ;
Wall(4, 3)
(all 4) Wall(1, 4) ; Wall(2, 4) Wall(3, 4) Wall(4, 4) ...
to Wall(0, 0)
Wall(1, 0) Wall(2, 8) Wall(3, 8) Wall(4, 0) ;
Wall(0, 1) empty ; Wall 1) empty Agent(1, 2, direction= (8, -1), carrying=None) ; Door (2, 2, colorsyellow, state=locked)
Wall(4, 1) ;
Wall(0, 2)
empty
W4, 2)
Wall(e, 3)
Key(1, 3, colorsyellow) ;
Wall(2, 3) Goal(3, 3)
Wall(4, 3) Wall(1, 4) ; Wall 4) Wall(3, 4) Wall(4, 4)
W0, 4)
The difference is a..
The agent (-irection=-11, 0)) at pos (1, 2) becomes an agent (direction= -1)).
the returned reward is 0.0 and the returned done is False
The action "turn lefe" transforms the state from
Wall(0, 0) Wall(1, 0) ; Wall(2, 0) Wall(3, 8) Wall(4, 8)
Wall(0, 1) empty ; Wall(2, 1) empty
Wall(4, 1) ;
Wall(0, 2) Agent 2, direction= 0), carrying=None) ; Door (2, 2, colorzyellow, state=locked)
empty
Wall(4, 2)
50

---

(o, 3) Wall(4, 3) Wall(e, 4) Wall(1, 4) ; Wall(2, Wall(3, 4) Wall(4, 4) ... to
Key(1, 3, colorsyellow) ;
Wall(2, 3)
Goal(3, 3)
Wall(0, 0)
Wall(1, 0) ;
Wall(2,
Wall(3, 0) Wall(4, 0)
Wall(e, 1) empty
empty Wall(0, 2) Agent(1, 2, direction= 1), carrying=None) Door 2,
Wall(2,
Wall(4, 1) ;
colornyellow, state=locked) empty Key(L, 3, colorsyellow) ;
empty
Wall (4, 2) Wall(2, 3) Goal(3, 3)
Wall(e, 3)
Wall(4, 3) Wall(1, 4) ; Wall(2, 4)
Wall(e, 4)
Wall(3, 4) Wall(4, 4)
The difference is
The agent (direction=(11, 0)) pos (1, 2) becomes agent (direction= 1)).
the returned reward is 0.0 and the returned done is False
The action "turn left" transforms the state from <e.
Wall(e, 0) Wall(1, 0) ; Wall(2, 0) Wall(3, 8) Wall(4, 8)
Wall(0, 1) empty
Wall(2, 1) empty
Wall 1) ;
Wall(e, 2)
Agent(1, direction= -1), carrying=None) ; Door 2, empty Key(1, 3, colorsyellow)
colorsyellow, state=locked)
Wall(4, 2)
Woll(e, 3)
Wall(2, 3) Goal(3, 3)
Wall(4, 3)
Wall(o, 4) ... to
Wall(1, 4) ; Wall(2,
Wall(3, 4) Wall(4, 4)
Wall(e, 0)
Wall(1, 0) ; Wall(2, 8) Wall(3, 8) Wall(4, 8)
empty Wall(2, carrying=None) ; Door (2, 2, Agent(1, 2, direction= 0), colorsyellow, state=locked) Wall(0, 3)
wall(0, 1) empty
Wall(4, 1) ;
Wall(0, 2)
empty Key(1, 3, colorsyellow)
Wall(4, 2)
Wall(2, 3) Goal(3, 3)
Wall(4, 3)
Wall(e, 4) Wall(1, 4) ; Wall(2, 4) Wall(3, 4) Wall(4, 4)
<..
The difference is
The agent (direction= -1)) pos (1, 2) becomes agent (direction= (-1, 0)). nuy
the returned reward is 0.0 and the returned done is False
Here example experiences that the code failed to model.
The action "toggle" should transform the state from
Wall(e, 0) Wall(1, 9) ; Wall(2,
Wall(3, Wall(2, 1) empty Wall(0, 2) Agent(1, 2, direction= 1), Norrying=None) (oor(2, 2,
Wall(4, 0) ;
Wall(0, 1) empty ;
Wall(4, 1) ;
empty Wall(2, 3) Goal(3, 3) Wall(0, 3) Key(1, 3, colorsyellow) ;
colorsyellow, state=locked)
Wall (4, 2)
Wall(4, 3)
Wall(e, 4) Wall(1, 4) Wall(2, 4) Wall(3, 4) Wall(4, 4) ...
to
Wall(0, 0) ; Wall(1, 0) ; Wall(2, Wall(3, 8) Wall(4, 8)
51

---

Woll(e, 1) empty
Wal 1) empty
Wall(4, 1)
Wall(e, 2) Agent(1, 2, direction=(8 (8, 1), carrying=None) Door (2, 2, colorsyellow, state=locked)
empty
Wall(4, 2)
Wall(e, 3)
Key(1, 3, colorsyellow) ;
Wall(2, 3) Goal(3, 3)
Wall(4, 3) Wall(1, 4) Wall(2, 4) Wall(3, 4) Wall(4, 4)
Wall(e, 4)
The difference is nuw
Nothing happened
the returned reward should be 0.0 and the returned done should be False However, the implementation is wrong because returns the predicted reward as -0.1 instead of the correct reward as 0.0
For this failed experience, do you know what is different between the rewards and dones from the environment and the predictions from the code? Do you know why the environment behaves in this way? Do you know why the code behaves di fiferrnnlyy from the environment? which part of the code causes the problem? How to fix it? Please improve your code to model the logic of the world for accordingly. Please implement the code following the template. You must implement the reward_ func function as the main function to be called by the environment The code needs to be directly runnable the (state, action, next_state) tuple and return (reward, done) in python provided in experiences. If the code too long, try refactor to shorter.
experiences,
F,5 Generating reward functions for new goals
asks LLMs to generate new reward functions new goals, given sample code that are synthesized for previous goals.
You robot exploring in an object-centric environment. Your goal is to model the logic of the world in python, specifically the reward function that maps (state, action, next_state) to (reward, done). You will to be given the mission for the environment few sample code from the other invironments. You need to implement the reward function for the environment you going act in. The code needs to directly runnable on (state, action, next_state) and return (reward, done) in python.
are going to act in, as well as
Here is few sample code for the reward function in other environments. Please check them in detail and think about how to implement the reward function for mission "pick up the yellow box" in the new environment. The code needs to be directly runnable on (state, action, next_state) and return (reward, done) in python.
The reward function code for mission "pick up the grey box" is:
state [["name": "A"gent"""" "Agent", "x":1, "y":1, "direction" (1,0), "carrying": {" name""""ey"" "x":: "y":None, "color":"red")),
("name": "Door", "x":2, "y":2, "color":"red", :locke""" ["name":"Wall] "x":0, "y":0]]
action "taggle"
next_state [["name": "Agent", "x":1, "y":1, "direction": (1,0), "carrying": ("name"; "Key", "x":None, "y":None, "color":"red""
(1,0))
("name" "Door", "x":2, "y":2, "color": "red", "state" :open")) ("name": "Wall", "x":0, "y":0}]
52

---

Box) and class Entity:
def inint__ (self, y, **kwargs):
self self... class -
self.x =
self.y =
for key, value in wwargs.items() setattr(self, key, value)
def __
(oey) for key, value in self.__ items () if key not in ('name', 'x', 'y'))
attr
if attr: return f" ; (self.y), (attr))"
else: return f"(self.name) ((self.x), (self.y))"
def __eq._ other):
return (self key) getattr(other, key, None) for key in self deyf
def return (attumnnhasn(upllesootoeddssss (self. dict (temmmm)))
class Agent(Entity): pass
class Key(Entity); pass
class (oor(Entity): pass
class Goal(Entity):
class Wall(Entity): pass
class (ox(Entity): pass
class Ball(Entity): pass
class Lava(Entity): pass
def aentiiiie..y_nameeeettte
name):
return entity for entity in entities entity.name name
byy_positioneentities x, y): def get_entities_.. best_entiiies.bbossstiinn
return entity for entity in entities if entity.x == x and entity.y ==
]
def reward_ action, next_state)
unn
Args:
state: the state of the envi
action: the action to be executed
next_state: the next state of environment
Returns:
reward: the reward of the action
done: whether the episode is done
ann
reward 0.0 initialise reward as 0.0
actions
done = False initialise done as False for all actions extract the agent from the current and next state 'Agent')
agent get.eentiiies...bnaeee(eeteteeee (next_state, 'Agent') [o] next_agent get_ peet_entitiee__y_neee
If the agent picks up the grey box in the next state, the reward is 1.0 and the episode is done if centtaaget.ccarryii and next_agent.carrying reward 1.0 done True return reward, done
'grey':
The reward function code mission "pick up the purple box" <44
state [["name": """nam"""""Agent" "Agent", "x":1, "y":1, "direction": (1,0), "carrying": {" name": "Key", "x":None, "y":None, "red""":"red")),
("name":" "Doorr, "x":2, "y":2, "color";"red" :lockedd)" ("name": "Wall", "x":0, "y":0)] action ""ooggle le"
53

---

next_state
[("name": "Agent", "x":1, "y":l, "direction": (1,0), "carrying":
("name": "Key", NNnn, "y":None, "color":"red").
("name": "Door", "x":2, "y":2, "color":" :ree", "state" :open"))
("name": "Wall",, "x":0, "y":0]]
class Entity:
def
(self, y, **kwargs):
self name seas.
self.x
self.y =
for key, value in kwargs.items() ():
setattr(self, key, value)
def
attr (sllff
(jje)((((ke))))valluu) for key, value in delf._..ictt_.. items () key not ('name', 'x', 'y'))
if attr: return ( [self.yh, (attr))"
else: return f" ((self.x, (self.y3)"
def _.eq._ other)
return al( key) getattr key, None) for key in dolf.._.ict__..eeee def (selas return hetuuumnnhhhhcccccoocccrooeeeeee (temmm))))
class :gent(Entity): pass
class Key(Entity): pass
class (EE pass
class Goal(Entity):
class Wall(Entity): pass
class (ox(Eittyy; pass
class Ball(Ent)ty): pass
class Lava(Entity): pass
def
name):
return entity for entity in entities entity.name name ]
def get_entities__e ienriiiis.bl.pssitioneeentiiit
x, y):
return entity for entity in entities if entity.x == x and entity.y ==
def reward_func(state, action, next_state);
unn
Args:
state: the state of the environment
action: the action to be executed
next_state: the next state of the environment
Returns:
reward: the reward of the action
done: whether the episode is done
reward 0.0 initialise reward as 0.0
actions
done = False initialise done as False for all actions extract the agent from the current and next state agent get_entities_by (stateytteemssttee 'Agent') [0]
next_agent get geet_entities__ Saaetieee_b__aameeeeeeeeeee (next_state, 'Agent') [0]
the agent picks up the purple box the next state, the reward is
1.0 and the episode is done and issinstnceeeext.ggnttcraaaa
if next_agent. mext__ggett.crryiing color color 'purple':
Box) and
reward 1.0 done True return reward, done
The reward function code mission "pick up the green box" is: <ec
54

---

state [["name": "Agentt" "x":1, "y":1, "direction": (1,0), "carrying": (" name": "Key", "x":None "y":None, "color":"red""
("name": "Door", "x":2, "y":2, "color":"res :red",, ""taaee""""ckeei"""" ("name": "Wlll", "x":0, "y":0)]
action "toggle"
next_state = [["name": =Aggett""""""""". "x":l, "y":l, "direction": (1,0), "carrying": {"name": "Kee", NNoe, "y":None, ;re")""
("name" "Door", "x":2, "y":2, "color": :red",, "state": :oope")) ("name": "Wall", "x":0, "y":0]]
class Entity: def (self,, y, **kwargs):
self.name self._.
self.x
self.y =
for key, value in (tems))
setattr key, value)
def
attr (self))
(xke))(v" for key, value in self._. items () if key not in ('name', 'x', 'y'))
if attr: return F"[self.name) ({self.x), {self.y), (attr))"
else: return f"(self..amme (fself.x, (self.y))"
def (self, other):
return al1ggettttr(ee11, key) gethr,, key, None) for key diff
in self.
___hash_(seelff= def
return (rttunmnhhhhtuuuullsssttosseeees ritemss)))
class Agent(Entity): pass
class Key(Entity); pass
class Door(Entity)) pass
class Goal(Entity):
class Wall(Entity)): pass
class (Ex pass
class Ball(Entity): pass
class Lava(Entity):
def
name):
return entity for entity in entities entity.name name ]
def hentiitfesby_posstfooneee
x, y):
return entity far entity in entities if entity.x == x and entity.y ==
def reward_func action, next_state):
unn
Args:
state: the state of the ronment
action: the action to be executed
next_state: the next state of the environment
Returns:
reward: the reward of the action
done: whether the episode is done
mnn
reward 0.0 initialise reward as 0.0 for all actions
done = False initialise done as False for all actions extract the agent from the current and next state agent get_entities_tt hetrteee_b__meeeeeeeeeeeeeee (state, 'Agent') [o] next_agent get__eniities_b__nmm (next_state, 'Agent') [0]
If the agent picks up the green box in the next state, the reward is 1.0 and the episode is done
if next_agent.c www.aaggtt.ccrringggg...... and
Box) and
next_agent.com www.aagnn.ccrryinnccommm color 'green':
reward 1.0 done True return reward, done
55

---

Now, have entered new environment. It shows mission "pick up the yellow box". Do you know what this mission reward function? Analyze the behaviors the reward function case by case In what situations will it return a positive reward or not? In what s situations will it return done=True or not? Why? Please implement the code following the template in the sample
and how to imp lementent it in
code. You must implement the reward_ func' funcc called by the environment. The code needs to be directly runnable on mission, state, action, next_state) and return (reward, done) in python.
function as the main function
F.6 Refine satisfy optimism under uncertainty
asks LLMs to think why the goal cannot be achieved for a mission from an initial state using the provided transition and reward functions. We also tell LLMs the valid action space and the measurement for achieving the goal (r > -Add
You are robot exploring in an object-centric environment. Your goal is to model the logic of the world in python. You have tried it before and came up with one partially correct solution. However, can model the logic for some experiences but failed to model the logic to achieve the goal in another environment. You need to improve your code so that the agent can achieve the objective specified by the mission from the given initial state as well still modelling the original logic. The new code should still follow the same template. The transition function needs to be directly runnable on (state, action) and return the next state in python. The reward_func function needs to be directly runnable on ( state, action, next_state) and return (reward, done) in python.
perfect. The code
Here is the partially correct solution you came up with:
class Entity:
def (self,, y, **kwargs):
self.name clas._ namee
self.x
self.y =
for key, value in kwargs.items ()
ssta, key, value)
def
attr ..pr__
(.ey for key, value in dict. items () if key not in ('name', 'x', 'y'))
if attr: return f"[self. nassef.nammedieee ( (self.y), (attr))" else: return f"(self.namm) ((self.xx, (self.y))" def (selff, other):
return all(getattrsse1f, key) (othr,, key, None) for key in self.__dict__._es]] def ___hah___self)) return hawtnnuuuuaeeeeortt(((sefffffff _icel..___dd_____ (ttms()))
class Agent(Entrty): pass
class Key(Entity)) pass
class Door(Entity)) pass
class Goal(Entity):
class Wall(Entity): pass
class (Ex(Entiyy): pass
class (Elt(in)) pass
class Lava(Entity): pass
import copy
56

---

def transition(state, action):
Args:
state: the state the envi ronment
action: the action to be executed
Returns:
next_state; the next state of the environment
mmn
We'll make deep copy of the state.
This is because we don't want to change the original state.
next_state (epy.deepcopy(staate) agent get_entities_by Saaeee__yyameeeeeeeeeeeeaee stex,_ 'Agent') [o]
action 'urrn left':
agent.direction (-1, 0):
agent.direction (0, 1)
elif agent.direction == (0, 1):
agent.direction (1,
elif agent.direction (1, 0):
agent.direction (0, -1)
else: if agent. == (0, -1)
agent.direction (-1, 0)
elif action 'turn right';
agent.direction (-1, 0):
agent.direction (0, -1)
elif agent.direction == (0, -1):
agent.direction (1, 0)
elif agent.directton (1, 0):
agent.direction (0, 1)
else: if agget.direction == (0, 1)
agent.direction (-1, 0)
elif action 'toggle':
We assume that agent have the ability toggle, regardless
what is in front of him because the experiences provided do not dictate otherwise.
pass
elif action 'nothing':
pass return next_state
def
name):
return entity far entity in entities entity.name
def get_entities_by sysiessbb__ositiionenntttt return entity for entity in entities if entity.x == x and entity.y ==
x, y):
def (seate, action, next_state):
unn
Args:
state: the state of the environment
action: the action to be executed
next_state: the next state of the environment
Returns:
reward: the reward of the action
done: whether the episode is done
enn
Create of entities for easier comparison and access state_set set(state)
next_state_set set(nex__state) Get agent's position in both states
agent [e for in state_set if isinstance(e, Agent)] agent_next [e for in rext_state_st set if isinstance(e, Agent) [0]
Done condition
on_goal = any (is (enttccc Goal) for entity in hir.tiie_bb_posiittoooooo done on_goal
agent_next.x, agent_next.y))
57

---

Reward calculation
if state_se == rext_state_sed _eate_set:
If state didn't change -> 'nothing', 'toggle' when not carrying key or 'drop' when carrying nothing happened
reward 0.0
elif action == 'turn left' or action = 'turn right':
If direction of agent changes -> 'turn left', 'turn right
happened
reward 8.0
else:
In other cases, no reward. Can be modi fied when other scenarios applied.
reward 0.0
return reward, done
However, the code failed achieve the goal/objective as specified by the mission "use the key open door then get to the goal" from following initial state:
...
Woll(0, 0)
Wall(1, 0) Wall(2, 0) Wall(3, 0) Wall(4, 0)
Wall(0, 1) empty
Wall 1) empty
Wall(4, 1)
Wall(0, 2)
Agent(1, direction= 1), carrying-None)
Door (2, 2,
empty Wall(2, 3) Goal(3, 3) colorayellow, state=locked)
Wall(4, 2)
Wa, 3)
Key(1, 3, colorsyellow) ;
Wall(4, 3)
Wall(e, 4) Wall(1, 4) ; Wall(2,
Wall(3, 4) Wall(4, 4)
...
The measurement for achieving the goal/objective is as follows:
<<4
def criterion(state, mission, action, next_state, reward, done,): return reward > 0 and done
...
The valid actions are ('turn right', 'nothing', 'move forward', 'turn left', 'toggle', 'drop', 'pick up").
Do you know why mission cannot be achieved from the given initial state with the world model as implemented in the code? what subgoals does the agent need to achieve in order to achieve the final goal as speci by the mission? Can the agent achieve those subgoals using the world model implemented the code? If not, what is missing wrong? How can you improve the code to achieve the goal/objective as specified by the mission from the given initial state? Please improve the code as analyzed before so that the mission can be achieved from the given initial
state. Please implement the code following the template. Feel free implement any helper functions you need. You can also implement the logic for difference actions in different helper functions. However, you must implement the transition function and the reward_ func function as the main functions to be called the environment. The transition function needs be directly runnable (state, action) and return the next state in python. The reward_ runnable on (state, action, next_state) and return (reward, done) in python. The new code, by themselves, should be complete, compilable, and runnable.
function needs to directly
58

---

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: We show that WorldCoder much more sample and compute efficient in our methods than baselines. We show that the optimism objective enables more efficient exploration.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made 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 NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Please check Section for details.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-specfiiaatinn 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 the scope of the claims made, e.g. if the approach was only tested on few datasets or 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 low or images are taken in low lighting. Or speech-to-tex, 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 faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an 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?
59

---

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

---

reproducible, they Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [No]
Justification:
We will prepare code for the public later due to time constraints.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.c) cc/
oublic/guides Seuuiee/Cooeeuumiisss CodeSubmi sotpsuolbbmisioooooocooooo for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is 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 exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https) //nips.. cc cttpcc/cplbcccccccc pub lic/ boubic/uuides/ Senss/Coeeubbne ss onoo for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) 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: We use standard hyper-parametert such as GPT-4 with temperature=1 and the maximum number of LLM requests is for each program synthesis problem.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can provided either with the code, appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: The error bars not important in our paper as method significantly outperforms the baselines in the sample efficiency and the compute efficiency in our experi- ments. We still show the error bars with 3 random seeds as stated. The error bars are default as 95% CI.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
61

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative еπτог 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 text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We mainly use GPT-4 and the maximum number of LLM requests for each program synthesis problem is 50.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs 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 htttpp:://eerrrpssccc//oub/////p
Answer: [Yes]
Justification:
The research conforms with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our work improves model- basel RI, which potentially indirectly has positive and negative social impacts as many technical works.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
62

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that 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 generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition 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:
The paper poses no such risks.
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 to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards 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 the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: Data and Codebase we used have been properly cited in the paper.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, 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.
63

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, -aper thcode com datasets 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.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper has not released new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well 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 in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country 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 crowdsourcing research with human subjects.
Depending 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.
64

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere 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.
65


Reviews:
[{"review_id": "AaQ1J1FGBD", "reviewer": "Reviewer_NhBf", "summary": "This paper introduces a novel model-based agent for sequential decision-making. A structured world model and goal-conditioned reward function are created by an LLM and refined as new information or task changes come in. This is based on logical constraints. For selecting the right candidate world model (ie object-oriented transition function), the constraint is that the world model explain all observed transitions up to now. For selecting the right goal-conditioned reward function, the constraint is that the reward function be optimistic, which means that it should assume that there is at least one trajectory reaching the goal.\n\nAfter proving a theoretical guarantee on the maximum number of actions to reach the goal, the authors move on to empirical evaluation. There are 3 evaluation environments: sokoban, minigrid and alfworld. On sokoban, a conceptually simple but hard grid-world puzzle, the proposed agent largely outperforms ReAct, an LLM-agent baseline in terms of score, as well as DreamerV3 and PPO in sample efficiency (by orders of magnitude). On minigrid, a sparse-reward gridworld for language instruction-following agents, the authors successfully demonstrate transfer between world models and rewards from one environment to the next in a curriculum, and demonstrate that when there is no curriculum variants without optimistic rewards fail to solve the tasks. The third environment is AlfWorld, a text-based environment modelling a kitchen environment with several rooms and objects, some of which have affordances such as being cut, opened, turned on, etc... The authors showcase the scalability of their method by demonstrating that it can learn a faithful world model of the domain and successfully solve all tasks in AlfWorld very rapidly.", "questions": "* First sentence of the abstract is not grammatically correct, it seems?\n* Could a version of WorldCoder work for hard linguistic/reasoning tasks (see limitation below)? [SWE bench](https://www.swebench.com/), [theorem-proving in Lean](https://leandojo.org/), etc?\n* What are typical logical dimensionalities of the environments you consider?\n* How do you ensure the self-repair program synthesis step does not break previously satisfied conditions for the transition function or the reward function?\n* The hand-designed curriculum limits generality of the approach, but could you leverage [Automatic Curriculum Learning](https://arxiv.org/abs/2003.04664) to create environments in the zone of proximal development?\n* line 289 “A disadvantage is that policies cannot readily generalize to new goals” except 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.", "limitations": "* Limitation: this agent cannot perform the complex linguistic tasks that humans do, since it is limited by the planner to operate in domains with a restricted number of actions (and that excludes the space of all possible sentences). The comparison to ReAct should mention this.\n* See above for curriculum and hand-defined reward;\n* Otherwise the authors acknowledge some additional limitations, such as the lack of modelling of stochasticity in their world model or the need for the observations to come pre-segmented in objects (or the world to be organized in recognizable, distinct objects at all).", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The method is original and interesting. While it bears similarity with previous work using programs to represent transition functions, it is the first to have used LLMs to infer program-based world models written in a general-purpose programming language;\n* The paper is quite clear and reads well;\n* The work is adequately contextualized and relevant references are given along with a comprehensive account of the paper's position in relation to them;\n* Experiments support the claims and contributions stated by the authors.\n* I could see other researchers using methods inspired by worldcoder on their own domains; when the approach transfers, it is a massive gain in efficiency compared to other types of agents;", "weaknesses": "* The method remains slightly complex and nontrivial to implement from the paper alone, making easily available code an important part of the re-usability of the method;\n* The method lacks generality, needing a hand-defined curriculum in minigrid and a handcrafted language-based reward bonus on AlfWorld;\n* Compared with state of the art LLM-agents (like Voyager), the environments considered here (especially AlfWorld) are pretty simple. Could WorldCoder perform well on ScienceWorld, an environment with more interesting dynamics and objectives? (breed plants, measure boiling temperatures, etc). Could it, in principle, scale to the text-based version of MineCraft used in Voyager?\n* I think the paper lacks solid baselines in the LLM-agent case. ReAct is a very basic agent, and a more principled baseline for tasks like sokoban would be Tree of Thoughts (ToT), dfs version, that also performs planning. It could be plausible that ToT performs well, in which case the compute efficiency claims still hold.\n* The paper lacks baselines in minigrid and alfworld.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the comprehensive review and for being supportive of the paper. Please below for answers to your questions.\n\n> The method lacks generality, needing a hand-defined curriculum in minigrid and a handcrafted language-based reward bonus on AlfWorld;\n\nWe think this might be a misunderstanding:\n1. The method doesn't need a curriculum (see Fig 4C), but can take advantage of good curriculum if one is provided (Fig 4B)\n2. 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.\n\n> the environments considered here… are pretty simple\n\nWe 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.\n\nPlease 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.\n\n> a more principled baseline for tasks like sokoban would be Tree of Thoughts (ToT)\n\nWe 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.\n\n> The method remains slightly complex and nontrivial to implement... easily available code an important part\n\nCode and data will be released upon publication.\n\n> 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\n\nGood point! We will mention that.\n\n> The paper lacks baselines in minigrid and alfworld\n\nPlease 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.\n\n> Could a version of WorldCoder work for hard linguistic/reasoning tasks\n\nThanks 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.\n\n> How do you ensure the self-repair program synthesis step does not break previously satisfied conditions\n\nWe 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.\n\n\nThanks again for your support and please let us know if you have further questions.", "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": "Ovbe2iNvvB", "reviewer": "Reviewer_9H9M", "summary": "- This work deals with the transfer of text-based LLMs to the domain of agents, enabling LLMs to build a (world) model of an environment.\n- The paper introduces the algorithm 'WorldCoder', that creates and refines Python programs for approximation of state transition and reward function\n  based on sampled data from the environment.\n- Hereby, the program search and refinement is conditioned on two constraints: \n (1. fit data) no state transition inaccuracies (batch) w.r.t. the observed data\n (2. optimism) the reward function should give a reward atleast for one state, enabling exploration.", "questions": "- Scalability is not properly addressed. E.g. looking at the prompts in Appendix for MiniGrid: the states are encoded in a rigerous manner with\n  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? \n\n- Experiments on grids or highly abstracted env states: Is this an analogy to Kahnemann's System 2 or a current limitation?", "limitations": "- The authors do not address input data integrity: The method relies on perfect(?) input sample quality. If not, then the LLM can infer wrong rules for\n  the transition function. Since each rules is abstract, the program can also generalise wrong. Thus, the robustness for low quality data is not addressed.\n  \n- The generated Python programs can only model highly abstracted MDPs, since the input samples are a text-based, tokenized abstractions.\n\n- Experiments on grids or highly abstracted env states => How to ensure hiqh quality input data in the general case (e.g. pixels)?", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "- The paper is technically sound and seems generally reproducible.\n- Good results, high sample effiency and task generalization for experiments, but limited domain with some assumptions.\n- Topic: Bridging the gap between LLMs and the agent domain by building world models with code via env interactions seems an interesting area to explore. \n  Researchers are probably interested in this idea.", "weaknesses": "- Equation 1,2 are hard to parse.\n\n- Task generalization achieved by prior knowledge: Env generalization with uncertain/incomplete prior knowledge of LLMs is not properly addressed.\n  E.g., the dynamics behavior of grid worlds is obviously known by GPT-4 and therefore we can expect generalized dynamics in the transition function code.\n\n- Scalability is not properly addressed. (See question)\n \n- Since Program Synthesis via Refinement (Equation 3 (REx)) is not developed in this work and only adapted to the problem, the paper's contribution boils\n  down to prompt-engineering with output constraints, then reiterating over output content with additional prompts for debugging.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the thoughtful review, and for your support. Please see below for our responses.\n\n> 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\n\nThe system leverages its prior knowledge, but is not bound by it. Here are two examples, including new results in the attached PDF:\n1. 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.\n2. 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.\n\n> the paper's contribution boils down to prompt-engineering with output constraints\n\nThe 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.\n\n> 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?\n\nWe 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.\n\n> Experiments on grids or highly abstracted env states: Is this an analogy to Kahnemann's System 2 or a current limitation?\n\nIt is indeed analogous to System-2 abstract reasoning, or task planning in robotics more specifically\n\n> 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)?\n\nWe 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.", "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": "LBgqa6sBmP", "reviewer": "Reviewer_WssY", "summary": "The paper presents WorldCoder, a system for learning world models through program synthesis with an LLM. The framing is model-based reinforcement learning, where the agent interacts with the environment and collects data that is used to learn the model. WorldCoder uses a general-purpose language to synthesize the model and a test-and-debug cycle to improve the accuracy of the model.\n\nThe paper also uses the idea of optimism in the face of uncertainty of the learned model. This is achieved by preferring models that can lead to positive rewards. Even if the positive reward is not achieved in practice, this forces the agent to explore the space. Empirical results demonstrate the potential of the approach.", "questions": "Do you agree that the data contamination issue is not explained by the paragraph that starts with \"Almost surely\" in Section 3? \n\nHow much of the code of the models for the three domains is available online? \n\nHow would WorldCoder perform on a domain that is different from everything that is out there?", "limitations": "The paper discusses most of the limitations. The limitation that is missing is the data contamination I mentioned in the weaknesses section.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper makes contributions to the important problem of learning world models for sequential decision making. The solution is also interesting as it uses program synthesis with a general-purpose language; most previous work relies on domain-specific languages.\n\nThe paper is very well written and easy to follow. It provides examples where needed, and reading it was a pleasure.\n\nThe results are quite strong compared to deep RL methods. This is particularly interesting given that the gain is coming from learning a model of the environment. Model-based RL can be quite finicky, and the results this paper produces are outstanding.\n\nThe comparisons of versions of the system with optimism and without optimism are also good and important to have published.\n\nGiven that WorldCoder uses an LLM, it is also nice to see that the cost of using the model is much lower than other systems used for planning. Overall, the use of LLM to learn world models makes perfect sense and is promising.", "weaknesses": "My concern with the paper regards the validity of the experiments given how the LLM is trained. 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?\n\nThe paragraph about React (Section 3) attempts to deal with this issue, but it misses the point. It argues that having knowledge of Sokoban isn't enough to solve problems. This is true, but it is not what WorldCoder could use. WorldCoder would benefit from simply knowing the code of the models, search will do the rest. Having knowledge of the model is way more likely than having useful knowledge of how to solve Sokoban puzzles. Sokoban is computationally hard. It is difficult to imagine how one would be able to code Sokoban knowledge to benefit systems like React. On the other hand, it is very easy to code knowledge that would benefit WorldCoder: implementations of the models would suffice.\n\nData contamination is the only reason why I am giving a \"poor\" score to soundness.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the thoughtful review. We really think we can improve the paper to address your main points. Please see below.\n\n> 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?\n\nThis is an excellent point and it deserves careful consideration:\n\n1. 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. \n\n\n2. 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.\n\n> How much of the code of the models for the three domains is available online?\n\n1. 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).\n2. Minigrid (released 2022) and Sokoban are probably in the pretraining data. Our new Sokoban result is not in the pretraining data.\n\nWe 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.\n\n> The paragraph about React (Section 3) attempts to deal with this issue, but it misses the point\n\nUpon 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.\n\n> How would WorldCoder perform on a domain that is different from everything that is out there?\n\nWe'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.\n\n> Data contamination is the only reason why I am giving a \"poor\" score to soundness\n\nWe hope that you can reconsider your score in light of the new experiment and the above discussion.\n\nThank you for the unusually helpful review. Please let us know if you have any further questions.", "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": "fwcXOTtrfp", "reviewer": "Reviewer_x859", "summary": "The authors propose a model-based agent by building a Python program to represent the world model through interacting with the environment. \n\nThe authors define logical constraints to explain the interactions and to achieve optimism in the face of (world model) uncertainty.  \n\nThe authors prove the polynomial sample complexity, and conduct empirical study on Sokoban, Minigrid and AlfWorld.", "questions": "The proposed method relies on an imperfect LLM and an imperfect method for program synthesis, how to handle errors due to such imperfection? Is there a way to achieve perfect program synthesis? If the answer is no, then the learned world model is not guaranteed to be correct.\nThis is a fundamental issue.\n\nThe paper is about learning a world model.\nIs 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? \nThere should be such a perfect model for Sokoban and Minigrid.\n\n\nWhat if integrating the proposed approach, in particular, \\phi_1 & \\phi_2 to deep RL, if feasible, maybe together with an LLM for language understanding? \n\n\nHow about MuZero?\n\n\nHow about some applications (deep) RL is good at, e.g, chess? \n\n\n\nLine 90-91: “preferring an optimistic one guarantees that the agent can at least start making progress toward the goal, even if it later has to update its beliefs because they turned out to be too optimistic.”\n\nIs 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.\n\n\nUnder Line 97\nConditions \\phi_1 fit data and \\phi_2 optimization\nWhat if there are noisy data points? \nWhat if data points are contradicts with each other?\n\nThe meaning of condition \\phi_2 is not clear, e.g., what is the meaning of \\forall i \\in [l]?\nIt is desirable to explain the conditions in plain English, and explain all notations.\n\n\nLine 110\n\"Before ever receiving reward, \\phi_2 forces the agent to invent a reward function.\"\nWhy and how to \"invent a reward function\"? \nWhy should it work?\n\n\nLine 127 \n\"we can learn in time polynomial w.r.t. the diameter (a property of the true MDP) and the logical dimensionality (a property of the model space)\"\nTheorem 2.4.\nHow about the diameter and the logical dimensionality for problems like chess and go?\nWill they grow exponentially as the search space grows exponentially?\n\n\nLine 146\n\"This retrieval strategy assumes similar goals have similar reward structures\"\nThis may not be true. And how to define similarity?\n\nLine 138 \n\"without further learning from environmental interactions.\"\nLine 149\n\"which update Rˆ(c) based on interactions with the environment\"\nThese two statements contradict each other.", "limitations": "Yes.", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "+ a model-based agent\n+ learn a world model through interactions with the environment\n+ define logical constraints to explain the interactions and to achieve optimism in the face of (world model) uncertainty\n+ prove the polynomial sample complexity\n+ conduct empirical study on Sokoban, Minigrid and AlfWorld", "weaknesses": "- rely on an imperfect LLM and an imperfect method for program synthesis, and there is likely no way to achieve perfection\n- it is not clear the effect of exponential search space on the polynomial sample complexity\n- 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the thoughtful review. Please see below for our responses.\n\n> 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?\n\nInteresting 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.\n\n> rely on an imperfect LLM and an imperfect method for program synthesis, and there is likely no way to achieve perfection\n> 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.\n\nWe 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.\n\nBut 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:\n\n> 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?\n\nWe currently consider deterministic environments without noise, but as line 311 of the paper says, we are pursuing future work that overcomes this limitation.\n\n> it is not clear the effect of exponential search space on the polynomial sample complexity\n\nWe 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.\n\n\n> 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?\n\nWe'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. \n\nThe 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.\n\n> How about MuZero?\n\nWe think it would perform very similarly to Dreamer_V3, which we compare against.\n\n> “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.\n\nWe 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.\n\n> \"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?\n\nWith 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$\n\n> How about the diameter and the logical dimensionality for problems like chess and go? Will they grow exponentially as the search space grows exponentially?\n\nThe 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.\n\n> \"assumes similar goals have similar reward structures\" This may not be true. And how to define similarity?\n\nIndeed, 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.\n\nLast, about the perceived contradiction between these statements:\n> 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\"\n\n> 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.\"\n\nThese 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.\n\n> what is the meaning of $\\forall i \\in [l]$\n\nIt means all integers $i$ where $1\\leq i\\leq l$. We'll revise to clarify this and other notation.\n\nThanks for your input and please let us know if you have further questions.", "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:
Thank you everyone for the helpful input! Although there is a separate response to each review, we wanted to also include a global response to highlight a new experimental result this week performed to help better understand how the system adapts to new environments that are different from what it is seen in pretraining data.

For this new environment we added warp gates/teleportation portals to Sokoban that allow the agent to instantly transport itself. The logic is subtle because the gates become deactivated if one of them is blocked by a box. In the attached PDF we show an example of how the final trained agent solves problems using the teleporters in order to solve puzzles more quickly by taking shortcuts. Learning takes 2-3 episodes to learn the nuances of warp gates, depending on the random seed.

Please see attachment.

Meta-review: Overall most reviewers found the work to be an interesting and novel way of approaching sequential decision making with LLMs. 

A number of potential limitations and presentation issues were raised by the reviewers with most being well addressed in the rebuttals. 

The authors are encouraged to leverage the reviewers' feedback on the presentation and results to improve the final camera ready paper.
Decision: Accept (poster)