PoE-World: Compositional World Modeling with Products of Programmatic Experts
Authors: Wasu Top Piriyakulkij
Yichao Liang
Hao Tang
Adrian Weller
Marta Kryven
Kevin Ellis
Venue: NeurIPS 2025 spotlight ()
Keywords: program synthesis
compositionality
world modeling
Primary area: general_machine_learning
TL;DR: We introduce a novel program synthesis approach to output world models of complex, non-gridworld domains by representing world models as products of programmatic experts.

Abstract: Learning how the world works is central to building AI agents that can adapt to complex environments. 
Traditional world models based on deep-learning demand vast amounts of training data, and do not flexibly update their knowledge from sparse observations. 
Recent advances in program synthesis using Large Language Models (LLMs) give an alternate approach which learns world models represented as source code, supporting strong generalization from little data. 
To date, application of program-structured world models remains limited to natural language and grid-world domains. We introduce a novel program synthesis method for effectively modeling complex, non-gridworld domains by representing a world model as an exponentially-weighted product of programmatic experts (PoE-World) synthesized by LLMs.
We show that this approach can learn complex, stochastic world models from just a few observations. 
We evaluate the learned world models by embedding them in a model-based planning agent, demonstrating efficient performance and generalization to unseen levels on Atari's Pong and Montezuma's Revenge.

PoE-World: Compositional World Modeling with Products of Programmatic Experts
Wasu Top Piriyakulkij
Yichao Liang?
Hao Tang'
Adrian We
Marta Kryven'
Kevin Ellis!
Cornell University?
University of Cambridge"
The Alan Turing Institute
Dalhousie University
Abstract
Learning how the world works is central building AI agents that can adapt complex environments. Traditional world models based on deep learning demand vast amounts of training data, and do not flexibly update their knowledge from sparse observations. Recent advances in program synthesis using Large Language Models (LLMs) give alternate approach which learns world models represented as source code, supporting strong generalization from little data. To date, applica- tion of progran-structuree world models remains limited to natural language and grid-world domains. We introduce novel program synthesis method for effectively modeling complex, non- soiddrrdwoordd domains by representing world model an exponentially-weihhee product of programmatic experts (PoE-World) synthesized by LLMs. We show that this approach can learn complex, stochastic world models from just a few observations. We evaluate the learned world models by embedding them in model-based planning agent, demonstrating efficient performance and generalization to unseen levels on Atari's s Pong and Montezuma's 's Revenge. We release our code and display the learned world models and videos of the agent's gameplay at uthttps://topaaat git..b.......................... io/ wwwpbbiooopppeeeeeeee world
Introduction
How should an intelligent agent represent the dynamics of the natural world? We want a representation that is efficiently learnable, yet flexible enough to handle stochasticity and partial observability, and which supports planning and decision Docision-making, Neural network world models such as Dreamer [T] are radically flexible, but demand enormous training data (compared to humans (2).. Symbolic world models such as WorldCoder 3 instead generate a Python program to represent how the world works. These programmatic world models are data-efficient, because program synthesis requires less data than neural network training- but struggle to scale beyond simple gridworlds, as they do a discrete combinatorial search to find single large program describing everything about how the world works.
We take inspiration from a longstanding view in philosophy and cognitive science of the mind as community of interacting experts 6 77. This modular organization is evident across multiple scales in natural intelligence, from functional specialization of brain systems (8) to the distinct learning trajectories of specific skills 9 [0]. We integrate this modular perspective with the computational paradigm of learning as Program Synthesis, which models learned concepts as symbolic programs 14 [5| [6]. We extend this line of work by proposing new computational account of learning world models: as the acquisition of context-specific expert programs, which are refined through practice and reused compositionally to support flexible, gial-directed behavior.
[2 13
Algorithmically, our key idea is to decompose the problem of learning a world program into learning hundreds of small programs. Each of these learned programs encodes different causal law, which we probabilistically aggregate to predict future observations (Figure D Th). This makes our world
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

(a) World Model Representation observation o
(b) World Model Learning Demonstration
observation o'
action a
LLM-based code generation
Srayertouchessplatooee playerv,
P(o" lo,a)
plavertouches((edder) DOWN:
Programmatic experts
Optimize weights the programmatic experts
playervv
W2
playertouchess(eeeee (nnmmy) playendelcted
W2
Wn
(c) Planning using the world model
(d) Generalization to novel, made-up levels
player louching a skull
player can't move through walls
A player can climb up and down ladders
action LEFT
action
NOOP action RIGHT
player can move on platforms
Model-based agent acting novel environment
World model
Planner
Figure (a) World models predict the next state given state-action history. We do this with product of experts of many small programs. (b) The learner is given a short (<l min) demonstration of gameplay as input, and uses synthesize initial world-model programs. These programs are refined online in later environment interactions. (c) World models support planning by imagining future states. (d) Symbolic programs encode abstract knowledge that generalizes to new game levels.
knowledge more modular, and also more learnable, because we no longer search for single monolithic program handling everything once. The resulting system, which call PoE-World (Product of programmatic Experts), can build elaborate world models that accurately support planning and reinforcement learning (RL) from even a brief demonstration in complex Atari games, such as Montezuma's ' Revenge. Wor-d handles stochasticity because the product of programs is probabilistic, and, as we show, further handles partial observability. To the best of our knowledge, this is the first time symbolic world model has been learned for environments of this complexity.
Importantly, although level visual appearance, instead assuming symbolic observations from an object detector. Unlike model-based reinforcement learning, PoE-World does not attempt efficient exploration, but focuses on faithfully learning from a demonstrated trajectory (Figure[p) Lastly, while the world models are ultimately used for planning (Figure [P.) PoB-World fundamentally focuses world modeling, and not on solving the challenging computational problem of planning itself.
models fine-grained pixel-level movement, it does not model pixel-
Despite these limitations, view BoE-dorld addressing central learning problem: Given limited demonstrations of new environment, quickly assemble a working world model that compositionally generalizes to situations (Figure It). We highlight the following contributions:
The PoR-World representation and learning algorithm for symbolic world models.
2. An empirical study of PoE--Wolld Revenge, demonstrating its superior learning efficiency compared to deep RI, and improved scalability compared to state-of-the-art symbolic model-based RL.: PoE-World 4000+ line programs that generalize zero-shot to novel game levels and game variations.
representative Atari games, Pong and Montezuma's
synthesize
Demonstration of how to use World's world models for planning-baan based decision making, and simulated pretraining environment for deep RL.

---

2 Background: World Model Learning
A sequential decision-makingg problem can be described as (O. where O is an observation space, is an action space. an environment dynamics P Pene Pen 01:t, Q1:t) where o E O and a E A, and R is a reward function. This full-history environment formulation is mathematically equivalent to a Partially Observable Markov Decision Process (POMDP) formulation [17] In the setting where P is unknown to an agent, the agent learns by interacting with the environment and observing transitions (Ot, 0t, 0t+1, T2) at timestep
We focus on learning the world model P Pmodet (0++1 01:r, (1:r) which approximates the true, unknown environment dynamics P == Penv (0r+1 10 01:)) given observed trajectories D )r where each trajectory a sequence of observations and actions T = (01:T+1, 1 (1:T) for some T (ri:T is also part of a trajectory, but we drop it for simplicity). The learned world model will later be used by a model-bar based agent, either through lookahead planning or RL training, to act in the environment.
We treat the world modeling problem as an optimization problem, specifically empirical risk mini- mization:
T
Proodet = arg min
Σ
Σ
E(pooded 01at+1 (1sst)
Pmodel
where l is a loss function, such as the negative log likelihood function {(Pmodeli 01:1+1, Q1:))
log Prodei(Or+11 (0:0::1111:::000001011 a1:1).
Previous works in the model-based reinforcement learning literature 18
have used various deep neural network architectures, including convolutional and recurrent neural networks, transformers, and diffusion models, to define Pmode! as a parametric model Proaed = Po- Then, 9 is optimized via gradient descent. A weakness of such approaches is poor sample efficiency and generalization. For example, Diamond (22) has failure modes such as imagining that the player can walk through wall teleport, even after training on almost 100 hours of observations.
Other works leverage Large Language Models (LLMs) to synthesize code world model 3 24 25
which can be seen searching for programmatic Pmodel- These LLM-based code generation algo- rithms are more sample efficient and extrapolate more systematically than deep learning approaches. However, they have yet to succeed beyond simple text-based and gridworld games.
3 Modeling the World as a Product of Programmatic Experts (PoE-World)
To address the limitations of existing world model learning methods, we propose representing world models as exponentially-weighhe
Products of programmatic Experts (PoE- world, enabling sample- efficient and scalable learning of probabilistic world models that leverages L.MM code generation. Figure visualizes both the representation and learning algorithm discussed this section.
3.1 World Model Representation: Product of Programmatic Experts (PoE-World)
World models represented as exponentially weighted Products of programmatic Experts (PoE-World) can be described mathematically as follows:
po(or41lo1a a12) & I expert (01+1
expertt are programs and 0g are their associated scalar weights. where p.
This representation enables modularity and compositionality. It allows composing many small programs into full world model capture complicated environmental dynamics (see Figure 2 We can think of each program as an expert which expresses opinions about particular aspects of the world. For example, in the context of modeling a video game environment, one expert might encode "if a player touches skull, then the player dies," while another encodes "if an action is LEFT when the player on a platform, the player's x-axis velocity is -2." The former expert does not express opinions on the player's movement, while the latter does not mention conditions for player's death.
3

---

If action RIGHT ...
p expect with w 0.2
player.tocches(plaaffmm and action RIGHT:
Most likely player's position at next timestep
expezt
player.Vx 2
P expect with
if playertouchees(lltform) and action == RIGHT: player.v,
with 0.4:
(layer.touchessladdeee) and action RIGHT: player.
P expert
expart
Figure 2: A "heat map" illustration of how simple Python programs are interpreted as distributions, and how they are combined into single distribution over next-timestep object locations.
factored state representation yields tractable inference. We assume an object-centric state where each object has a bounding box and velocities, each stored as separate slot, attribute, feature. We treat each such feature as conditionally independent, given the history so far. This independence compute separate normalizing constant
assumption makes Equation (2) tractable, because for each feature. Concretely, let f index the different object features the next observation 0++1 and let of be the indexing of o with feature f. Then,
Po ( ( a1a) = II II f
expert (of lova, ana)) Z, = EI eapert (o
lols, @1:))
0:
Pr
Pr
of
Benefits of the full-history formulation over POMDP formulation. The full-history environment formulation is formally equivalent to a POMDP which instead compresses the history into a Markov latent state. We the history formulation because it makes world models more modular. Learning global latent variables would entangle all the experts, because every expert would condition on the latent state. Therefore, learning new latent variable (e.g. "how long the player has been falling") changes the input/output space of every expert, necessitating global joint updates to the structure of every program (e.g. an expert for "is the player dead"' would need to be updated). The history formulation allows independent learning of independent mechanisms.
Hard constraints. Atari (and the real world) is too complex perfectly simulate with any ef- fectively learnable program. Therefore, our probabilistic model tends to over-approximate the set of possible futures, giving fuzzy approximate predictions. For example, in Montezuma's Revenge, instead of perfectly modeling the physics of falling downward and landing on the ground, we predict generic downward trajectory. Ideally that trajectory would perfectly enforce the constraint that the player lands flat on the ground, and never sinks into the ground, but a fuzzy stochastic expert for falling downward could violate that constraint. Therefore, to sharpen the model's 's outputs, we further learn collection of hard constraints, {cg }. where O + {0, 1}:
po(0441l01a,
1 expert
(0141/01a, a1
(g((o41))
(3)
We further discuss hard constraints and provide concrete examples in Appendix A.I
Multi-timestep predictions. represent the programmatic experts and the world distributions over the next-timestep observations in eq. (2). It is likewise possible to reformulate multi-timestee expert p expert (0r+1:++H H (0t,,, a1:)) as product of next-timester experts by assuming that the predic- tions of the multi-step expert at different timestep are independent:
"
p expert (0++1:++ (01ss a1:))
p eppert (0441
(4)

---

3.2 World Model Learning: Program Synthesis and Weight Optimization
We begin with a demonstration trajectory, learn world model, and then begin act in the world according to that model. As the agent acts, it collects more trajectory data, which it uses to update or "debug" its model. Concretely, learning proceeds follows:
Step 1: Synthesize programmatic experts {p. (p
Ji=1 given observed trajectories D { )=
Step 2: Fit the weights 0 of the experts according to eq. (5) with a gradient-based optimizer
Step 3: Remove the experts with weights below threshold
Step 4: Repeat Step 1-3 every time the observed trajectories get updated
we adopt Large Lan- guage Models (LLMs) as our Python program generator. We input small batch of transitions (Ot:+H+1, #1:: to the LLM prompt to produce the programmatic experts (p.epent).
Generating program experts. Following previous works
3
24 25),
Figure 2 shows how we interpret small Python programs distributions over the observations. While we could ask LLMs to synthesize probabilistic programs to specify the distributions, we find it much more effective have LLMs synthesize simple, deterministic Python programs, presumably because generic Python code is far more prevalent in LLM training data. We assume that an observation represented as a list of objects, where each object has the following attributes: x/y velocity and visibility. Then, a distribution over observations is distribution over each object's ' attributes. To interpret a Python program as a distribution, we assume that all object attributes are conditionally independent given full history, as mentioned in Section 3.1 and convert all object attributes set by the program to distributions with single peaks at the given values. We then add noise to the distributions to ensure non-zero probabilities over alternative values. Any object attributes whose values are not set by the program follow uniform distributions over all possible values: consequently, experts with a single if-condition (fig. 2 yield uniform distribution when the if-condition is not satisfied.
Gradient-based Weights Optimization. likelihood estimation obtain 0:
Once we have expert pm Ji=1. we can perform maximum
{pr
0 arg max
Σ
Σ
log
(5)
(DLT+IALTTEDD
This equation instantiates eq. (T) by letting Proadel have a parametric form Pmodet po and choosing negative log likelihood as the loss function. We can use any gradient-based optimizer to optimize the weights. We use L-BFGS (26), which worked better than Adam 27 and SGD, because we have small data and few parameters.
Finally, the experts whose weights are below threshold are removed from the world model. We repeat this loop every time there are new observations. Appendix A.T contains full algorithm details.
3.3 World Model Usage: RL in Simulation and Planning with World Model
An important goal of world modeling is to aid decision-making. We consider two such ways of using world models. First, world model can serve simulator for reinforcement learning. This makes RL policy learning more sample efficient, because we can quickly learn a world model from real environment interactions, which then substitutes the actual environment. Subsequent policy learning can take place in the world model, obviating the need for further interaction with the real environment. In practice, we also continue reinforcement learning after pretraining in the world model. Formally, we learn a policy π O' + A that inputs an observation history and outputs an action.
Alternatively, world models plan for a horizon of H timesteps by searching for an optimal action sequence, given our previous observations 01: and actions 01.t-1:
be used for lookahead planning. Given reward function R,
G:44 = arg max E (014
Ro (1444,
Gu:44
where po aloaa, (14+H) pe peloc-R41001cc (144))
I11
5

---

3.
To get the key (goal), the agent needs to perform very long sequence (length 100) with no intermediate reward
The task is much easier if it is broken down into a sequence of sub-tasks (high-level plan).
Thus, build abstract graph. The nodes are defined by object contact, and the edges connect pairs of nodes that low-level planner can find plan to traverse between.
4.
5.
6.
Thus, search in this graph for possible high-level plans. An agent can now choose high-level plan follow.
Some high-level plans might not work. When an agent acts in the environment and discovers that plan does the false edge from the graph and update the world model.
Eventually, we find high-level plan that works
not work,
Figure 3: sequence illustrations that demonstrates how our hierarchical planner works.
To help our world models guide long- long-horizon decision-making inspired by task and motion planning (TAMP). It first plans in high-level abstract state space defined by object contact, then lowers those plans into actual Atari button presses using the learned world model (Figure Appendix A.2).
implement hierarchical planner
The resulting agent we denote PoE-World Planner.
Experimental Results
Domains and Evaluation. We evaluate our agent, PoE-World Planner, against other methods on Atari's Pong and Montezuma's Revenge (MR), using the Arcade Learning Environment (28). We OCAtari [29] parse each image frame list of objects, each with object category, a bounding box, and velocities. Both games are partially observed: the current state and action cannot uniquely determine the next state. demonstration of fewer than 1000 frames created for each game, but these demonstrations are not successful gameplays: they serve only to illustrate the core causal mechanics. In Montezuma's Revenge, our demonstration never achieves positive reward.
To test compositional extrapolation, we created alternative versions of both games, called Pong-Alt and Montezuma's 's Revenge-Alt (Figure both of which recombine and rearrange the types of objects seen in the training demonstration. Pong-Alt increases the number of objects (3 balls and 3 enemies). Montezuma's Revenge-Alt adds more enemies (which the player has to jump over) and ladders, while changing the map to resemble the game Kangaroo (see Appendix A.3 for details), We do not provide demonstrations for these alternative versions of the games.
'OCAtari does not reliably extract objects so we have make custom changes OCAtari for every single Atari game we apply our method to, which makes it infeasible to run on the full Atari suite. We choose Pong and Montezuma's is Revenge (arguably one of the hardest games in Atari [30) calls it "grand challenge") as representative games for which manually patch OCAtari reliably detect objects.

---

Method
Score
Pong
Pong-Alt
MR
MR-Alt
Random Agent
-20.67 + 0.33
-20.00 0.58
0.00 = 0.00
0.00 + 0.00
PPO @ 100k env steps BT
-21.00 + 0.00
-18.66 = 0.88
1.00
0.00 0.00
LLM as Agent (ReAct)
32
-20.00 + 0.00
-20.67 0.33
0.00 0.00
WorldCoder Planner
-17.00 + 3.00
-19.00 = 1.00
0.00
0.00
PoE- World + Planner (Ours)
-12.83 + 0.88
-13.67 + 0.67
100 0.00
100 0.00
PPO @ 20m env steps 31
17.00 + 0.58
1.33 + 2.03
0.00 + 0.00 0.00 + 0.00
Table Scores on Pong and Montezuma's 's Revenge (MR) and their alternate versions. For PoE-Workd and WorldCoder, brief demonstrations on Pong and MR are given to initialize the world models. Their agents then train for at most 3k steps before evaluation.
Baselines. PPO 33 is a go-to, widely-used model-free RL algorithm. It optimizes a lower bound of the policy's performance using gradient descent. LLM Agent (ReAct) [32] directly uses an LLM as a policy. ReAct prompts LLMs to use extra Shain-of-thoughh [33] "thinking" actions before selecting an action. WorldCoder program. It uses an LLM code generation and repair algorithm called REx [34] to refine its world
S is LLM agent that models the world
single Python
model to achieve high predictive accuracy on observed trajectories. More details in Appendix A.4
Agent Results. Figure 5 and Table [ show the scores of different agents on Pong, Pong-Alt, Montezuma's Revenge, and Montezuma's Revenge-Alr. The scores on Pong and Pong-Alt indicates the difference in points achieved by the player and the enemies when the game ends at points. The scores on Montezuma's 's Revenge and Montezuma's Revenge-Alt become positive if and only if the agent succeeds in collecting the key. shown in T our agent, PoE-World + Planner, performs best across all environments in the low-data regime, particularly when PPO baseline is allowed to have 100, 000 training environment interactions, the standard budget for sample-efficient agents on Atari 2600 200. In Figure 5 we keep training PPO for more steps, finding that it takes over a million steps for PPO to surpass PoE- World + Planner. Moreover, PPO with 20M training steps never achieves positive score on Montezuma's 's Revenge. PoE-World Planner is the only method that manages to obtain positive reward on Montezuma'ss 's Revenge in both base and alternative versions.
Pong-Alt
Montezuma's Revenge-Alt
Figure 4: Screenshots of the alternative environments. Pong-Alt's objective is to hit three balls past three enemies to score points, while Montezuma's 's Revenge-Alt requires players to avoid moving skulls by jumping over them and climbing up ladders to reach and collect the key.
7

---

Max Score Difference
Max Score Difference
over Random Baseline (Pong)
over Random Baseline (Pong-Alt)
Max Score Difference
PPO
Max Score Difference
PPO
ReAct
ReAct
WorldCoder
WorldCoder
15
PoE-World (Ours)
PoE PoE-World (Ours)
20
10
10
5
(demonstration)
100
102
10
10
10a
10a
10c
Training Steps log scale)
Training Steps (in log scale)
Max Score Difference over Random Baseline (MR)
Max Score Difference
over Random Baseline (MR-Alt)
200
200
Max Score Difference
PPO
Max Score Difference
PPO
ReAct
ReAct
150
WorldCoder
WorldCoder
PoE-Worke (Ours)
PoE-World (Ours)
100
100
50
50
(demonstration)
10a
10a
104
10c
100
10a
104
106
Training Steps (in log scale)
Training Steps (in log scale)
Figure 5: Maximum score differences over random baseline achieved at different number of training steps (in log scale). The "(demonstration)) annotates the length of demonstration used initialize WorldCoder and FoE-World, which is not displayed for the Alt games since we use the world models initialized by the demonstrations of the corresponding base games.
Instead of planning, can use PoE-World to learn a policy? Training a policy avoids the test-time compute cost of planning. In Figure 6 we show that pre-training a policy inside Pol- World's world model accelerates policy learn- ing: we first run PPO "in simulation" (in our world model), and then fine-tune in the actual Atari environment. The fine-tuned PPO achieves significantly higher score than vanilla PPO at most training steps, and while the vanilla PPO takes training steps to do better than a ran- dom agent, the fine-tuned PPO takes only 200k training steps. Asymptotically the pretrained and randomly-initialin
Max Score of PPO W/ and w/o Pre-training (Pong)
40
PPO w/o Pre-training
PPO Pre-training
PoE-World's world model
Max Score
20
-20
5M
10M
15M
20M
Training Step
policies converge to the same value: World-mode pretraining is ef- fectively way of warm-starting policy training.
Figure PPO with and without pre-training within our world model.
Next Observation Prediction Results. Table 2 and Table 3 show next observation and next observation's object attributes prediction accuracies under the symbolic world modeling approaches, respectively. PoE- World outperforms the baselines on most settings, except for the test observation (random frames) for Pong where all methods perform similarly because random agent rarely succeeds in hitting the ball, and so knowledge of game mechanics is not comprehensively tested.
The role of hard constraints. We investigate the role of hard constraints our world model representation in Table 4 In practice, hard constraints act to rule out "physically impossible" scenes. PoF- World ended up enforcing constraints just for MR, since most possible scenes Pong are already physically possible. As shown in Table 4 removing hard constraints causes the agent's performance

---

Method
Next Observation Prediction Accuracy
Train
Test (1000 Random Frames)
Pong
MR
Pong
Pong-Alt
MR
MR-Alt
WorldCoder
0.33 + 0.01
0.36 + 0.02
0.06 = 0.01
0.01 = 0.01
0.10 + 0.01
0.12 0.03
PoE-World (ours)
0.51 + 0.01
0.75 + 0.00
0.06 = 0.01 0.03 + 0.00
0.31 + 0.00
0.43 0.00
Table 2: Next observation prediction accuracy the training demonstration and 1000 random frames.
Method
Next Observation's Object Attributes Prediction Accuracy
Train
Test (1000 Random Frames)
Pong
MR
Pong
Pong-Alt
MR
MR-Alt
WorldCoder
0.82 + 0.00
0.75 = 0.01
0.66
0.01
0.67 = 0.02
0.58 + 0.00
0.65 + 0.02
PoE-Worke (ours)
0.88 + 0.00
0.93 = 0.00
0.66 0.00
0.76 = 0.00
0.76
0.83 = 0.00
Table Next observation's 's object attributes prediction accuracies on the training demonstration and 1000 random frames. Non-moving objects, such as walls, platform, etc. are excluded.
Method
Planning successes
Next Observation Prediction Accuracy (Test)
MR
MR-Alt
MR
MR-Alt
PoE-World + Planner w/o Hard Constraints
2/9
2/9
0.30 0.01
0.43 + 0.00
PoE-World + Planner
5/9
4/9
0.31
0.00
0.43 + 0.00
Table 4: Planning successes (out of 9 tries), where the goal is for the player to grab the key. and next observation prediction accuracies on 1000 random test frames of our agent with and without hard constraints on Montezuma's 's Revenge (MR) and its alternate version (MR-Alt).
drop on both the base and alternative versions of Montezuma's 's Revenge. Interestingly, we found no significant changes in next observation prediction accuracies. We hypothesize that hard constraints do not necessarily turn bad predictions into good ones. Instead, they perform damage control -hey refine poor predictions just enough to make them usable for Long-horizon planning.
Qualitative difference between WorldCoder's and PoE-World's world models. World models produced by PoE-World consists of 4000+ lines of code for Montezuma's 's Revenge compared to WorldCoder's ' less than 100 lines. This difference reflected in their ability to capture the underlying models accurately represent important causal laws like character movement constraints with respect to platforms and ladders, while WorldCoder fails to model these mechanics, predicting that the player can "fly"" around the map without any constraints. Moreover, WorldCoder models tend to hallucinate, e.g., imagining nonexistent bullet firing abilities in Montezuma's 's Revenge, potentially because there is no granular way to downweigh buggy parts of world model. in contrast, prunes irrelevant experts with low weights, yielding more precise world models.
causal laws:
5 Related Work
World models as programs has been explored in several recent works that motivate PoE- World. Similar to our work, WorldCoder ( and CodeWorldModel 24 use LLMs to write Python transition function. Unlike our work, they learn a single monolithic program. which severely limits scalability: world models have order of magnitude more code needed to model complex environments, as well as handle partial observability and non-determinism. Recent works have also explored synthesizing high-level abstract world models programs [35] 36 to support robotic
planning by integrating visual perception with symbolic reasoning. They could synergize with our work, as we have focused on learning low-level world models describing the motion of objects.
learn world models in restricted (non-Turing complete) domain specific programming languages, performing well on benchmarks co-designed with their Domain Specific Language. Even earlier, Schema Networks learn conceptually related factor-graph world model [39]. AIXI [40] is a theoretical model of reinforeement learning which considers all possible Turing
Earlier works 37 38

---

machines, mathematically related to all these works. Interestingly, AIXI also works with a history space as conceptually elegant way of modeling partial observability. Other studies generate programs for world models primarily from natural language instructions, rather than from example interactions with the environment
Hierarchical Planning. Our approach to segmenting complex continuous tasks into symbolic states inspired by hierarchical [45] and task-and motion [46 47 planning. Hierarchical representations in Reinforcement Learning (RL) are often expressed as options-ttempral abstractions that allow agents reason at multiple time scales by grouping sequences of primitive actions [48]. Options can be manually specified by design or learned from experience [49]. Task and Motion Planning (TAMP) automatically discover symbolic states and action abstractions that enable generalizable symbolic plans grounded in continuous motion, however the applications of TAMP to Atari-style domains are
still lacking due to complex dynamics. Similar to recent work [29 50], we define a high-level state space in terms of object contact relations- which is key to supporting symbolic abstraction.
Alignment with human cognition. Our system, composed of simpler, specialized programmatic experts that give rise to complex behavior, is inspired by view of mind as a community of interacting agents - a recurring theme in philosophy and cognitive science [ 6 52 7 Our modeling approach to modeling objects and actions aligns with empirical studies of event segmentation [53] [4] and
hierarchical planning [55] [6]. Hierarchical state-spaces based on motion cues, such as ours, predict how people draw [13 and interpret social interactions [57], attesting to the cognitive aligmment of our approach. Likewise, our programmatic representations of actions align with studies that demonstrate human concept learning to be akin to mental programs 58 59
12], which recent work models by
LL.M-based program synthesis.
Discussion and Limitations
Compositionality Recombining pieces of knowledge to generalize and extrapolate is a core feature of symbolic systems and, arguably, also of human cogmition, ranging from natural language to Our approach in this spirit. It factors its knowledge into small experts whose predictions can combine to extrapolate to scenes with more objects recomposed into new arrangements. More formally, our approach generalizes to novel "entity compositions" and "relational compositions, terms used by [63]. It should be noted, however, that our compositional factoring is orthogonal to our use of symbolic code: Programs be monolithic (3), and neural nets can be factored [64] Nonetheless, compositionality proved critical to scalably learning the underlying symbolic program. Synthesizing a single monolithic program is, in our view, intractable not just for the real world, but even for Atari.
abstract reasoning everyday thought [60] 52 61 62].
Limitations. We make important assumptions, and only address part of the full model-based reinforcement learning problem. Symbolic programs expect symbolic inputs: We do not learn straight from pixels. RL involves exploration, decision-maine making, and reward function learning, but we do not address those problems here. However, we speculate our approach could unlock better methods for exploration: program-structure world model exposes an interpretable interface for describing beliefs about how the world works, and efficiently exploring the world is analogous to testing the program that encodes the world model. Therefore ideas from software testing and program analysis could, in theory, be brought to bear, enabling new approaches to model- exploration.
Acknowledgments
We thank Edward Gu and Taha Jafry for their contributions to earlier versions of this project. We also grateful to Athary Sonwane, Yilun Du, and Tom Silver for their valuable discussions and feedback on the paper, and we appreciate the OCAtari team's 's assistance and support with the OCAtari package. This work was supported by NSF CAREER grant.
10

---

References
[1] Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301 04104, 2023.
[2] Pedro Tsividis, Thomas Pouncy, Jaqueline Xu, Joshua Tenenbaum, and Samuel J Gershman. Human learning in atari. In AAAI Spring Symposia, volume 1, 2017.
[3] Hao Tang, Darren Yan Key, and Kevin Ellis. Worldooder, model-based LLM agent: Building world models by writing code and interacting with the environment. In The Thirtt-eighth Annual Conference on Neural Information Processing Systems, 2024.
[4] Jemry A Fodor. The modularity of mind. MIT press, 1983.
[5] Marvin Minsky, Society of mind. Simon and Schuster, 1986.
[6] Daniel Clement Dennett and Daniel C Dennett. Consciousmess explained. Penguin uk, 1993.
Gary F Marcus. The algebraic mind: Integrating connectionism and cognitive science. MIT press, 2003.
[8] Michael S Gazzaniga, Richard B Ivry, and GR Mangun. Cognitive neuroscience. the biology of the mind. sixth edition, 2025.
[9] John R Anderson. Acquisition of cognitive skill. Psychological review, 89(4):369, 1982.
[10] Sannyuya Liu, Qing Li, Xiaoxuan Shen, Jianwen Sun, and Zongkai Yang. Automated discovery of symbolic laws governing skill acquisition from naturally occurring data. Nature Computational Science, 4(5):334-345, 2024.
[11] Brenden M Lake, Ruslan Salakhutdinov, and Joshua B Tenenbaum. Human-level concept learning through probabilistic program induction. Science, 13366266)1332--111 1338, 2015,
[12] Brenden Lake and Steven T Piantadosi. People infer recursive visual concepts from just few examples. Computational Brain & Behavior, 3(1):54-65, 2020.
[13] Lucas Tian, Kevin Ellis, Marta Kryven, and Josh Tenenbaum. Learning abstract structure for drawing by efficient motor program induction. Advances Neural Information Processing Systems, 33:2686-2697, 2020.
[14] Matthias Hofer. Tuan Anh Le, Roger Levy, and Josh Tenenbaum. Learning evolved combinatorial symbols with neuro-symbolic generative model. arXiv preprint wrXiv:2104.08277 2021.
[15] Joshua S Rule, Steven Piantadosi, Andrew Cropper, Kevin Ellis, Maxwell Nye, and Joshua Tenenbaum. Symbolic metaprogram search improves learning efficiency and explains rule learning in humans. Nature Communications, 15(1):6847, 2024.
[16] Yichao Liang, Josh Tenenbaum, Tuan Anh Le, et Drawing out distribution with neuro- neuro-symbolic generative models. Advances in Neural Information Processing Systems, 35:15244-1522 15254, 2022.
[17] Karl Johan Åström. Optimal control of markov processes with incomplete state information Journal of mathematical analysis and applications, 10:174-205, 1965.
[18] David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv:1803 180..00130010. 2018
[19] Julian Schritwieserr Toannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Gracpel, et al. Mastering atari, go, chess and shogi
by planning with a leamedd model. Nature, 58847839)6604-609, 2020.
[20] Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blažej Osiński, Roy Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, Afroz Mohiuddin, Ryan Sepassi, George Tucker, and Henryk Michalewski. Model based reinforcement lcarning for atari. In International Conference Learning Representations, 2020.
[21] Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample-efficient world models. In The Eleventh International Conference on Learning Representations, 2023.
[22] Eloi Alonso, Adam Jelley, Vincent Micheli, Anssi Kanervisto, Amos Storkey, Tim Pearce, and François Fleuret. Diffusion world modeling: Visual details matter atari. Thirty-eighth Conference on Neural Information Processing Systems, 2024.
[23] Guangyao Zhou, Sivaramakrishnan Swaminathan, Rajkumar Vasudeva Raju, J Swaroop Guntupalli, Wolfgang Lehrach, Joseph Ortiz, Antoine Dedieu, Miguel Lizzro-Gredilla, and Kevin Murphy. Diffusion model predictive control. arXiv preprint arXiv:2410.053e6 2024.
[24] Nicola Dainese, Matteo Merler, Minttu Alakuijala, and Pekka Marttinen. Generating code world models with large language models guided by monte carlo tree search. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
[25] Zergham Ahmed, Joshua B Tenenbaum, Christopher J Bates, and Samuel J Gershman. Synthesizing world models for bilevel planning, arXiv preprint arXiv:2503.2011 2025.

---

[26] Dong C Liu and Jorge Nocedal. On the limited memory bigs method for large scale optimization. Mathematical programming, 45(1):503-528, 1989.
[27] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
[28] M. G. Bellemare, Y. Naddaf, J. Veness, and M. Bowling. The arcade learning environment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47:253-279, jun 2013.
[29] Quentin Delfosse, Jannis Blüml,, Bjarne Gregori, Sebastian Sztwiertnia, and Kristian Kersting. Ocatari: Object-centric atari 2600 reinforcement learning environments. arXiv preprint arXiv:2306.088: 230.00640,, 2023.
[30] Adrien Ecoffet, Joost Huizinga, Joel Lehman, Kenneth O Stanley, and Jeff Clune. First return, then explore. Nature, 590(7847):580--33333 08600055666288 2021.
[31] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint
1707.06347 2017.
[32] 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.
[33] Jason Wei. Xuezhi Wang. Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc Le. Denny Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in Neural Information Processing Systems, 35:24824-2485 24837, 2022.
[34] Hao Tang. Keya Hu, Jin Zhou, Si Cheng Zhong, Wei-Long Zheng, Xujie Si, and Kevin Ellis. Code repair with Ilms gives an exploration-explooiatio tradeoff. Advances in Neural Information Processing Systems, 17:117954-11999 179996, 2024.
[35] Yichao Liang. Nishanth Kumar, Hao Tang, Adrian Weller, Joshua B Tenenbaum. Tom Silver, João F Henriques, and Kevin Ellis. Visualpredicator: Learning abstract world models with neuro-symbolic predicates for robot planning. arXiv preprint arXiv:2410. 10. 23156, 2024.
[36] Tom Silver, Rohan Chitnis, Joshua Tenenbaum, Leslie Pack Kaclbling, and Tomás Lozano-Pérez. Learning symbolic operators for task and motion planning. In 202/ IEEE/RS/ International Conferencee on Intelligent Robots and Systems (IROS), pages 3 3182-3189. IEEE, 2021.
[37] Pedro A. Tsividis, Joao Loula, Jake Burga, Nathan Foss, Andres Campero, Thomas Pouncy, Samuel J. Gershman, and Joshua B. Tenenbaum. Human-level reinforcement learning through theory- modeling. exploration. and planning. 2021.
[38] 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(POPL), January 2023.
[39] Ken Kansky, Tom Silver, David A Mély, Mohamed Fldawy, Miguel Azaro-Grediill, Xinghua Lou, Nimrod Dorfman, Szymon Sidor, Scott Phoenix, and Dileep George. Schema networks: Zero-shot transfer with generative causal model intuitive physics. In International conference on machine learning, pages
1809-1818. PMLR, 2017.
[40] Marcus Hutter. Universal artificial intelligence: Sequential decisions based algorithmic probability. Springer Science Business Media, 2005.
[41] Fan-Yun Sun, SI Harini, Angela Yi, Yihan Zhou, Alex Zook, Jonathan Tremblay, Logan Cross, Jiajun Wu, and Nick Haber. Factorsim: Generative simulation via factorized representation. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
[42] Feiyu Zhu and Reid Simmons. Bootstrapping cognitive agents with a large language model. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pages 655 663,6663222 2024.
[43] Lionel Wong, Jiayuan Mao, Pratyusha Sharma, Zachary S Siegel, Jiahai Feng, Noa Korneev, Joshua Tenenbaum, and Jacoh Andreas. Learning adaptive planning representations with natural language guidance. arXiv preprint arXiv: 2.t.2.//2.08666 2023.
[44] Lin Guan, Karthik Valmeekam, Sarath Sreedharan, and Subbarao Kambhampati. Leveraging pre-trained large language models to construct and utilize world models for model-based task planning. In Thirty- seventh Conference on Neural Information Processing Systems, 2023.
[45] Leslie Pack Kaelbling and Tomás Lozano- Pérez. Hierarchical task and motion planning in the now. In 2011 I IEEE International Conference on Robotics and Automation, pages 1470-1477. IEFF, 2011.
[46] Aidan Curtis, George Mathcos, Nishad Gothoskar, Vikash Mansinghka, Joshua Tenenbaum, Tomás Lozano- Pérez, and Leslie Pack Kaelbling. Partially observable task and motion planning with uncertainty and risk awareness. arXiv preprint arXiv:2403.110455 2024.
[47] Aidan Curtis, Tom Silver, Joshua B Tenenbaum, Tomás Lozano-Péree and Leslie Kaelbling. Discovering state and action abstractions for generalized task and motion planning. In Proceedings of the AAAI conference artificial intelligence, volume 36, pages 5377-5384, 2022.
12

---

[48]
Richard S Sutton, Doina Precup, and Satinder Singh. Between mdps and semi-mdps: A framework for temporal abstraction in reinforcement learning. Artificial intelligence, 112(1-2)1811-21, 1999.
[49] Pierre- Luc Bacon, Jean Harb, and Doina Precup. The option-critic architecture. In Proceedings of the AAAI conference artificial intelligence, volume 31, 2017.
[50]
Jannis Blüml, Cedric Derstroff, Bjarne Gregori, Elisabeth Dillies, Quentin Delfosse, and Kristian Kersting. Deep reinforcement learning via object-centric attention. arXiv preprint arXiv:2504.03024 2025.
[51] Satinder Singh, Michael James, and Matthew Rudary. Predictive state representations: A new theory for modeling dynamical systems. arXiv preprini arXiv: 4167, 2012.
[52] Brenden M Lake, Tomer D Ullman, Joshua B Tenenbaum, and Samuel J Gershman. Building machines that learn and think like people. Behavioral and brain sciences, 40:c253, 2017.
[53] Tan T Nguyen, Matthew A Bezdek, Samuel Gershman, Aaron Bobick, Todd Braver, and Jeffrey M Za- cks. Modeling human activity comprehension at human scale: prediction, segmentation, and categorization. PNAS nexus, 3(10) 3(10):gae,59 2024.
[54] Jeffrey M Zacks. Event perception and memory. Annual review of psychology, 71(1):165-19))) 1916 2020.
[55] Marta Kryven, Suhyoun Yu, Max Kleiman- Kleiman-Weiner, Tomer Ullman, and Joshua Tenenbaum. Approximate
planning in spatial search. PLOS Computational Biology, 10(11)cc112252, 2024.
[56] Jan Balaguer, Hugo Spiers, Demis Hassabis, and Christopher Summerfield. Neural mechanisms of hierarchical planning in a virtual subway network. Neuron, 90(4):893-903, 2016.
[57] Tianmin Shu, Marta Kryven, Tomer D Ullman, and Joshua B Tenenbaum. Adventures in flatland:
Perceiving social interactions under physical dynamics. 42d Proceedings of the Annual Meeting of the Cognitive Science Society, 2020.
[58] Marta Kryven, Cole Wyeth, Aidan Curtis, and Kevin Ellis. Cognitive maps are generative programs. Cognitive Sciences Society, 2025.
[59] Sugandha Sharma, Aidan Curtis, Marta Kryven, Josh Tenenbaum, and Ila Fiete. Map induction: Composi- tional spatial submap learning for efficient exploration novel environments. International Conference of Learning Representations, 2022.
[60] Nouha Dziri, Ximing Lu, Melanie Selar, Xiang Lorraine Li, Liwei Jiang, Bill Yuchen Lin, Sean Welleck, Peter West, Chandra Bhagavatula, Ronan Le Bras, Jena Hwang, Soumya Sanyal, Xiang Ren, Allyson Ettinger, Zaid Harchaoui, and Yejin Choi. Faith and fate: Limits of transformers on compositionality. In shirth-seventh Conference on Neural Information Processing Systems, 2023.
[61] Brendea M Lake and Marco Baroni, Human-like systematic generalization through meta-learning neural network. Nature, 623(7985): 15-121, 2023.
[62] Elizabeth S Spelke and Katherine D Kinzler. Core knowledge. Developmental science, 10(1):89-9)) 2007.
[63]
Atharva Sehgal, Arya Grayeli, Jennifer J. Sun, and Swarat Chaudhuri. Neurosymbolic grounding for compositional world models, 2023.
[64] Anirudh Goyal, Alex Lamb, Jordan Hoffmann, Shagun Sodhani, Sergey Levine, Yoshua Bengio, and Bernhard Schölkopf. Recurrent independent mechanisms. In International Conference on Learning Representations, 2021.
[65] Marlos C. Machado, Mare G. Bellemare, Erik Talvitie, Joel Veness, Matthew Hausknecht, and Michael Bowling. Revisiting the arcade learning environment: Evaluation protocols and open problems for general agents. Journal of Artificial Intelligence Research, 61:523-562, 2018.
[66] Ruocheng Wang. Eric Zelikman, Gabriel Poesia, Yewen Pu, Nick Haber, and Noah Goodman. Hypothesis search: Inductive reasoning with language models. In The Twelfth International Conference on Learning Representations, 2024.
[67] Linlu Qiu, Liwei Jiang, Ximing Lu, Melanie Sclar, Valentina Pyatkin, Chandra Bhagavatula, Bailin Wang. Yoon Kim, Yejin Choi, Nouha Dziri, and Xiang Ren. Phenomenal yet puzzling: Testing inductive reasoning capabilities of language models with hypothesis refinement. arXiv preprint arXiv:2310.08559 2023.
[68] Top Piriyakulkij. Cassidy Langenfeld, Tuan Anh Le, and Kevin Ellis. Doing experiments and revising rules with natural language and probabilistic reasoning. Advances in Neural Information Processing Systems, 37:53102-533 53137, 2024.
[69] Samuel Acquaviva. Overcoming the expressivit-elliciennyy tradeoff in program induction. Master's thesis, Massachusetts Institute of Technology, 2024.
[70] Jake Bruce, Michael D Dennis, Ashley Edwards, Jack Parker Holder, Yuge Shi, Edward Hughes, Matthew Lai, Aditi Mavalankar, Richie Steigerwald, Chris Apps, et al. Genie: Generative interactive environments. In Forty-first International Conference on Machine Learning, 2024.
13

---

[71] Geoffrey E Hinton. Training products of experts by minimizing contrastive divergence. Neural computation. 14(8):1771-1800, 2002.
[72] Caelan Reed Garrett, Rohan Chitnis, Rachel Holladay, Beomjoon Kim, Tom Silver, Leslie Pack Kaelbling, and Tomás Pozano-Pérez. Integrated task and motion planning. Anmual review of control, robotics, and attonomous systems, 4(1):265 293, 2021.
[73] Siyuan Zhou, Yilun Du, Jiaben Chen, Yandong Li, Dit- Dit-Yan Yeung. and Chuang Gan. RoboDreamer: Learn- ing compositional world models for robot imagination. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp, editors, Proceedings of the 4 International Conference Machine Learning. volume 235 of Proceedings of Machine Learning Research, pages 61885-61896. PMLR, 21-27 Jul 2024.
[74] Yilun Du, Shuang Li, and Igor Mordatch. Compositional visual generation with energy based models. In Advances in Neural Information Processing Systems, 2020.
[75] Zhutian Yang, Jiayuan Mao, Yilun Du, Jiajun Wu, Joshua B. Tenenbaum, Tomás Lozano-Pérez, and Leslie Pack Kaelbling. Compositional Diffusion-Based Continuous Constraint Solvers. In Conference on Robot Learning. 2023.
[76] A Paszke. Pytorch: An imperative style, high-performance deep learning library. arXiv preprint arXiv:1912.01703, 2019.
[77] Rémi Coulom. Efficient selectivity and backup operators monte-carlo tree search, International conference on computers and games, pages 72-83. Springer, 2006.
[78] Cameron B Browne, Edward Powley, Daniel Whitchouse, Simon M Lucas, Peter Cowling, Philipp Rohlfshagen, Stephen Tavener, Diego Perez, Spyridon Samothrakis, and Simon Colton. survey monte carlo tree search methods, IEEE Transactions on Computational Intelligence and Al in games, 4(1):1-43, 2012.
[79] Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dormann. baalle-baselins.:: Reliable reinforcement learning implementations. Journal of machine learning research, 22 22(268):1-8, 2021.
14

---

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: The claims made in the abstract and introduction accurately reflect the paper's contributions and scope.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not 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 are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
The paper discusses the limitations of the work performed by the authors.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested 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 is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to 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 fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
15

---

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

---

Answer: [Yes]
Justification: The code is available at https://anonymoun how.aoooooouuuuoop....ooo science//eeeeeeeccccccc // poe- she-world-neuse neurips-- 6
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.cm ec/
wetgggdddeeeeeeubiiiieeeeeeeeeee (ossiccyyyy)) for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https for more details.
//nips. //nips.cc/pil wwcppbblc/guuiicuccccccccoocoopp detggiddddeeeoeebbbiiinn
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, 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 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: Algorithm details are deseribed in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: The paper reports error bars computed over 5 seeds.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall 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).
17

---

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 in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We have provided OpenAl API cost for our experiments and compute resources used for our experiments.
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 as well as 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 failed experiments that didn't make into the paper).
9. Code ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification: The research conducted the paper conforms, in every aspect, 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 (c.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 is foundational research and has no direct societal impact.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
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.
18

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology 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 ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, 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 is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: The creators of original data are all properly credited, and the license and terms of use of the data are explicitly mentioned and properly respected.
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, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the Datasets.papersstt paperswi thcode. cww/ datasets package should be provided. For popular 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 it has changed) should be provided.
19

---

If this information is not available online, the authors are encouraged to reach out to the asset'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 currently does release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmoel 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 is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
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 in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question:
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor 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.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
20

---

Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [Yes]
Justification: This paper describes the usage of LLMs in method.
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy http:////eeuripp.cc/ooofeeeeeeee for what should or should not be described.
21

---

A Appendix
A.1 PoE-World algorithm details
As mentioned in Section 3.2] our learning algorithm PoE: World alternates between two main steps: programmatic expert synthesis and maddiett-based weight optimization. Given trajectory, do batch processing of size 10-we first learn pe (011 (01:00 41:10), then pe(021 (1:20), and so on.
Programmatic expert synthesis. To synthesize programmatic experts, first need to turn the observations into natural language observations that can be input into LLMs. A sequence consecu- tive observation transitions [0+.+++++1, (t.++11)) 11) be transformed into text showing the first input object list, the list of actions, and the changes to the input object list for the following timesteps, all in natural language. Table shows example text representation.
We note that in the text representation, only changes to a single object type (which is 'player' in Table 5) is displayed at time. This choice is to enforce modularity between the object types: each experts put non-uniform distribution only on single attribute of objects with a specific object type. can be further rewritten
This means Equation (2
po (0 l01dt (1:1)
II II
Pj obj typetypeeee expert 10141 01:);
(6)
obj-type
where {p; phpeetpppppppo hwpptteppeeeeeee expert are the experts associated with each object type.
We then implement multiple (10 in total) LLM--ased synthesis modules. For each object type. these modules take as input sequence of consecutive observation transitions and output set of programmatic experts. The output programmatic experts are pooled into a single set once all modules finish running. Having multiple synthesis modules allows each module to focus on different aspects of the environment, e.g., how objects move passively, how object moves when interacts with objects, how objects get created and deleted, etc. We describe one of our modules, ActionSynthesizer, below.
ActionSyn focuses synthesizing experts that explain how each action affects objects when the objects are interacting (touching) other objects (see example experts of this kind in Figure 2 takes in just a single transition (0r:2+1,00 az). It turns this transition into text representation as discussed above (Table 5). Then, it prompts a LLM to output causal explanations for the object changes (see the prompt T). An example causal explanation "the player objects that touch an ladder object set their y-axis velocity to -4". With set of causal explanations, we prompt LLM to turn each of them into program (see the prompt Table 7 and Table
Other modules follow this template of first prompting a LLM for natural language causal explanations and then prompting LLM to turn the explanations into programs. We refer the reader code at httpp::/ github...... com/ http///twwwwueeeeee poe- world for the implementations of other L.M--based module.
The programs synthesized by LLMs use our manually-written helper classes: Obj, ObjList, RandomValues, and SeqValues. Their docstrings are passed to the prompt. Obj provides method Obj touches which is the function we manually build in to help determine object contact. Obj
Example input list of objects:
player object (id = 0) with x-axis velocity = +4,
Interaction player object (id 0) is touching ladder object (id 2),
Interaction -- player object (id 0) is touching unknown object (id 4),
Example list of actions: NOOP, NOOP, RIGHT
Example output list of object changes: The player object (id = 0) sets x-axis velocity to [+0, +0, +2]
Table 5: Example text representation of a sequence of consecutive transitions.
22

---

contains object attributes: object category. x/y position, x/y velocities. It also contains properties including center_x, center-y, beft_side right_side, which are calculated based on the posi- tion and velocities, and their setter methods actually modify the velocities under the hood. ObjList class that represents an object list. has a method ObjList. Besiissttee__bb3_ _eet__bj__br _ and ObjList. create_ create_object. Lastly, RandomValus and SeqValues are the classes used to mark values set by LLM. The purpose these two classes is further discussed below.
After the programmatic experts are synthesized (see examples in Table 2). we interpret them as distributions as discussed in Section 3.2 We tell LLMs in the prompt to set attribute values as instances of BandomValues as opposed to integers. This marks the attributes whose values changed by LLM so that we can write an program- coauumtt--iiee oo-distribution interpreter that puts a single-peak distribution on the attribute whose value is set and uniform distributions on all other attributes. SeqValues is similar to RandomValues, but is used in multi-timestep predictions scenario, as discussed in Section 3.1
The LLM used in the steps above is 2pt-40-2024-08-066 We implemented a disk cache for the LLM responses to avoid paying multiple times for the same prompts and seeds.
Hard constraints. The hard constraints (Section 3.1 are also synthesized similar to how we synthesize the programmatic experts. Table 10 shows the constraints learned for Montezuma's Revenge. In Equation (3), we choose to use a disjunction rather than a conjunction because the physics in video games can be peculiar and unrealistic player might have their body overlap with a platform when climbing down a ladder attached to that platform. In the real world, we believe a conjunction would be a better choice.
Gradient-based weight optimization. Once we have the expert distributions {p. expert we can optimize their weights {w;} according to Equation (5). We use the L-BFGS optimizer (26 imple- mented in PyToorch (76) with strong Wolfe line search, learning rate = 1, number of epochs = 4, and without mini- bathib. We also include a Ll regularization loss with weight so that the weights do not get too big.
We note that the weight optimization is done without taking into account the hard constraints since we would like to use gradient- based approaches.
We prune programs with weights lower than 6 = 0.01 after the weight optimization is done, and we prune constraints that contradict with the observations or explain less than 1% of the observations.
A.2 Planner
In a game like Montezuma's 's Revenge, planning in the actual, low-level action space is hard because the number of actions required to get to the first positive reward is very high at least 100. This means the size of search space is g100 as there are 8 possible actions: NOOP, UP, DOWN, LEFT, RIGHT, FIRE, LEFTFIRE, RIGHTFIRE.
Thus, inspired by task and motion planning (TAMP), our hierarchical planner interleaves planning in the low-level motion action space with planning in a high-level abstract action space. It learns an abstract graph where the nodes are abstract states defined by object contact, and the edges represent whether the world model believes the player can traverse between the two nodes. Then, we search for a path in the abstract graph that takes the player to the goal. Atari, a goal usually for a player to touch goal object (key, ball, platforms, etc.). The discovered path in the abstract graph is high-level plan-a sequence of subgoals-ffrr a low-level planning agent. A low-level planning agent performs online planning to choose action and then execute it in the environment.
The planning algorithm can be described step-by-step as follows:
Step 1: Learn an abstract graph by running a low-level motion planner in simulation on all pairs of nodes. A (ordered) pair of nodes has an edge between them if we can find at least one low-level plan to traverse between them.
Step 2: Search for path in the abstract graph. We use breadth-firse search (BFS) here to find path to the goal with the shortest length. The path our high-level plan. If there is no path to the goal, go back to step 1.
Step 3: Attempt to follow the high-level plan, completing each subgoal in order, with a low-level planning agent.
23

---

Step 4: If there is "false" edge. update the world model, remove that edge from the graph, then go back to step 2. Otherwise, the algorithm stops, and the agent has achieved the goal.
Figure 3 shows simplified illustrations of how our hierarchical planner works.
We now discuss the implementation of our low-level motion planner:
Low-level motion planner. We implemented two low-Ievel motion planners. The first a variant of Monte Carlo Tree Search (MCTS) [77] It follows the same set of procedures as vanilla MCTS with two differences: first, similar to the MCTS algorithm used in [3] approximates the value of a node using a heuristic function instead of doing random rollout in the simulate step. The heuristic function is the Manhattan distance between the current position of the player object and the position of a goal object. We find that this function is good estimate of how good an observation is when trying to achieve a goal. Second, the value nodes, instead of the expected value, in the backpropagation step. Intuitively, using the maximum value encourages the planner to be more optimistic. The exploration parameter for MCTS is initially equal to for Montezuma's 's Revenge and Montezuma's 's Revenge- Alt and 10 for Pong and Pong-Alt, and it increases by 10 times every 1000 iterations of MCTS.
node is updated as the maximum value of its children
We also implement "sticky actions" by extending the action space that MCTS searches on so that it includes repeated sequences of primitive actions with lengths 1, 4, and 8. Thus, the extended action space has 3n actions where n is the number of primitive actions. We include these repeated action chunks in our action space to make planning easier. Playing Atari games rarely requires players to change actions at every timestep, so repeated action chunks can be helpful.
The second low-level motion planner is greedy search with the same heuristic function used in MCTS. At each iteration, it greedily finds the best repeated action chunk of length 8 and includes it in the plan. It backtracks if the current state leads to death no matter which action chunk of length the planner chooses. It returns a plan when the player achieves the goal (touches the goal object).
Pong, Pong-Alt, and Montezuma's 's Revenge-All Alt agents only use the greedy planner, while Mon- tezuma's ' Revenge uses both: it first tries to find a plan with MCTS for 4000 iterations and falls back to the greedy planner if MCTS fails.
As discussed earlier, the low-level motion planner is used in two steps: to build the abstract graph by finding plan to traverse between two abstract nodes in the world model, and to help inform low-level planning agent that is trying to complete a subgoal.
Low-level planning agent. The agent performs online planning: it uses the low-level planner to find a plan to the goal, and then it takes a sequence of actions and replans. The agent replans when the ccurent plan no longer takes the agent to the goal, so the agent may take several actions before replanning. We further optimize the agent by letting it replan only 40% of the times when the current plan no longer works in Montezuma's ' Revenge.
Because of this introduced stochasticity, however, we find that our whole planning pipeline can give different scores in different runs even with the same initial world model, so we treat running the hierarchical planner multiple times with the same initial world model as part of training, and we run the planning algorithm 3 times on the same initial world model to get results on Montezuma's Revenge and Montezuma's 's Revenge-Alt. Our work focuses on world modeling, and we leave it to future work to increase the efficiency and performance of the hierarchical planner.
A.3 Domain details
We evaluate methods Atari's Pong and Montezuma's Revenge using the Arcade Learning Environment (ALE). The frameskip parameter is set to 3 for both games. We use OCAtari to parse each image frame as a list of objects, each with object category, a bounding box. and velocities. OCAtari rengieess the RAM values of Atari games to get the bounding boxes of each object.
We modify OCAtari to fix number of issues in its handling of Pong and Montezuma's 's Revenge. This includes fixing bugs in object detection so that it fully detects every object in play, correcting bugs in the the bounding boxes so that object interactions correctly correspond to when object bounding boxes touch/overlap. etc. We refer the reader to our code implementation https: : /// github. com/ topwasu/poe- world for full details.
24

---

Details on how we create Pong-AI Alt and Montezuma's Revenge-Alt are below:
Pong-Alt is created by layering three Pong environments. We syne the player's location in all three environments, so that it appears as if we have only one paddle. The enemies and balls, on the other hand, are all different locations. We end up with one player, three enemies, and three balls.
Montezuma'ss 's Revenge-Alt is created by stacking the lower section of the first room in the original Montezuma's 's Revenge make three platforms of different heights, connected by stairs. We stack up three lower sections of three different Montezuma's 's Revenge environments. The first and the third section the stack are flipped horizontally so that the stairs are on different sides of the room, requiring the player to jump over the skulls to reach the stairs.
ALE code uses GPL-2.0 license, and OCAtari code uses MIT license.
A.4 Baseline details
PPO. PPO uses frame stacking = 4. We use the same hyperparameters as the PPO paper 31
and OCAtari 29 We the stable-baselines? implementation of PPO with MLP backbone ('MIpPolicy') 79.. For the Alt environments, we take the PPO model pretrained 20M steps on the base environments and finetune it on the Alt environments. This pretraining process is done for fiar comparison with our method, since our method assumes demonstrations from the base environments when evaluating on the corresponding Alt environments.
ReAct. We write prompts that would alternate hetween thinking and taking actions, one for Pong and Pong-Alt and another for Montezuma's 's Revenge and Montezuma's Revenge-Alt. The frame observation is transformed into text where we provide each object's x and position and list of all pairwise object interactions. This text observation is input as part of the prompt. In the prompt, only include the 4 most recent observations (along with the most recent thinking actions and 4 most recent taken actions) as each observation is quite long in text.
WorldCoder. We use the official WorldCoder implementation 3 but replace the existing prompts with new ones that are tailored towards our text representation of object-centric Atari frames. The prompts can be found in our codebase ation of WorldCoder in B actually has its own planner implemented, but for fair comparison with our method, we use only the world modeling part of that system and combine it with our own planner so that the planner is the same for both WorldCoder and our method. We choose to use our own planner instead of theirs since ours is hierarchical.
www/topwwsu//peeecooooo wtwpaas/ppoewwool//// world The instanti-
WorldCoder code MIT license.
A.5 Compute resources and execution time
Compute Resources. For the world modeling part, our experiments are run on 4 CPUs (Cascade- Lake, IceLake, SaphireRapids) with 64 GB memory. PoE- World and WorldCoder uses budget of $20 worth of OpenAl credit per run. For the planner. we also run mostly on 4 CPUs, but for the part where need to build abstract graph by rumning many low-level planners, parallelize on multuple compute jobs on a job scheduling cluster, so we might be using 100 CPUs at a time.
Execution time. PoE-World World alone without the planner tends to take around 8 hours to (this includes the time we need to wait for OpenAl LLM requests). The planner running time varies, but most runs finish under 24 hours.
25

---

I'll give you an input list objects and an output list of object changes, and I want you to list possible reasons for the effects
Here's example with player objects:
Example input list of objects:
player object (id = 0) with x-axis velocity = +0 and y-axis velocity +2, Interaction -- player object (id 0) is touching ladder object (id 2),
Interaction player object (id 0) is touching unknown object (id 4),
Example output list of object changes:
The player object (id = 0) sets x-axis velocity to +0 The player object (id 0) sets y-axis velocity to -4
Example reasons:
The player objects that touch an unknown object set their x-axis velocity +0
1.
2. The player objects that touch an unknown object set their y-axis
velocity -4
3. The player objects that touch an ladder object set their x-axis velocity
4. The player objects that touch an ladder object set their y-axis velocity
Please output a list of 4 reasons of the tobj_type} objects for the following input and output list of objects.
Input list of objects:
(input}
Output list of object changes: (effects)
Please follow these rules for your output:
1. make sure each reason only talks about one object change
2. not talk about IDs
Table prompt in ActionSynthesizer used get LLM to provide causal explanations of the object changes.
26

---

We observe that the possible effects of {action) on {obj_ (tpe) objects include {oba_lat_txtt)
want aynthesize python functions that implementa these effects. The format of the functions should be
ObjList, action: str) -> ObjList: fobj_type)_ooje = obj list bet..obje.y_ooj..rpp........ ("objj. # get all Dbj of obj_type '[obj_type)'
def hoeer_(bb_tttpe_oojjcctttso
for tobj_type>_ob in [obj_typee # tybj_type)_ott _obj of type Obj # You can assume {obj_type)_ood _bj velocity_x and tybj_type))_bp velocity_y are integers.
pass return obj list
And here are the docstrings for relevant classes:
class RandomValues:
Use this class to express the possibility of random values. Example x RandomValues ([x + 2, x 2])
Attributes:
values (list [ints)) list of possible values
Methods:
init__ ival__e)
Initialize instance
class Obj:
Attributes:
id (int); id of the object
obj_type (string)) type of the object
velocity_x (int | RandomValues)) x-axis velocity of the object
velocity_y (int I RandomValues) y-axis velocity of the object deleted (int RandomValues): whether this object gets deleted (1 if it does and if it does not)
Methods:
(obj: Obj) -> bool:
Returns whether this Obj is touching the input obj (True/False)
class ObjList:
Attributes:
obja (list of Obj)
Methods:
get_objs_be bas_by__bbjjytp type(obj str) -> 1ist [Obj]: Returns list of objects with the input tyje
create_objest (obbee__bbecctbb_tyype str, int, y: int) ObjList:
Returns a new instance ObjList with the new object (with obj_type, x,
y) added
Table 7: First half of prompt in ActionSynthesizer used to turn the natural language causal explanations into programs.
27

---

Please output {n} different alter_ Aater_(obj_typp _bbects functions that explains each of the {n} possible effects of action '(action'' following these rules:
1. Each function should make changes to one attribute this could be the x-axis position, y-axis position, creation of object, or deletion of object.
2. Always use RandomValues to set attribute values. there are conflicting changes to an attribute, instantiate RandomValues with a list all possible values for that attribute.
3. Use Obj touches to check for interactions.
4. Avoid setting each attribute value
each {obj_type) object more than once. For example, use 'break' inside a nested loop.
5. You can assume the velocities of input objects are integers.
6. Please use if-condition to indicate that the effects only happen because action "(action''
Format the output
numbered list.
Table 8: Second half of a prompt in ActionSynthesizer used the natural language causal explanations into programs.
Pong
Mon Montezuma/ ss Revenge
Pong-Alt
Montezuma
Revenge
Figure Screenshots of both base and alternative environments for Pong and Montezuma's Revenge.
28

---

Example program 1 def htw_pplayyr_ojectsoool Beuch_percenttO.00 6) -> ObjList:
ObjList, action: str, touch_side-3,
if action 'NOOP":
player_objs obj_list.get_obee http:a..eee..ooobbboooooopoppppp lonveyer_belts for player_ in player_objs:
moj_list.gel beb.at..ett.bjj.bbbbb........ ort.://o.....oooo..ooooorrrrrrrr "conveyer_belt")
for conveyer_belt in conveyer_beil belts: if player_obj touches touchee touches (conveyer_belt, touch_side,
touch_percent):
player_ob] velocity_x RandomValues ([-1]] break
([-1])
return obj_list
Example program 2
def alter_player jobjeess__jcccsbboo_o (ibj_list: ObjList, action: str, touch_sidess
souch_percent=1.0) ObjList:
if action == 'FIRE":
player_objs obj ooo__.ese...g....o...ooooovvvvrr
platform_objp = obj _ibiis..ett_bbbb_boooo....pppp ('platform')
for player_ob player_objs:
for platform_obj in platform_objs: player_ob] touches (platform touch_side,
touch_pprreett
player_ velocityyy RandomValues ([ ([-6]) break
return obj_list
Example program 3
def touch__ercent -> ObjList:
ObjList, action: str, touch_side=3,
if action == 'RIGHTFIRE':
player_objs [b___la__eet.ojjjs....oo.ooooood
platform_oby obj obtj_laat.eet_bbjjbbb........... type (typttypapppoon platform') for player_ob] in player_objst
for platform_obj in platform_objas if player_obj] touches (platformooo) touch_side,
touch_percent):
player_obj velocity_y = SeqValues -7, -4, 0, 2, 6,
9])
break
return obj_list
Table Example synthesized programmatic experts.
29

---

Constraint 1: Player's body must align with the center of the rope def cl clboj_list: ObjList, souch_ touch.percent-0.3) tuple: touch_ids, satisfied_ide ids O player_objs - htt.ssssgeet.oos.......oe...._ee ('player') # get all Obj of type pplayer'
for player_ in lil_iisttte__bjeee.. bybj__tyttp__yppeeccecc
rope_obja abj__ist.__
get all Obj of type 'rope' _aaee_obj is of type Obj
for rope_obj in rope_objs: # rope_obj is of type Obj
if tow.ces (rope_obj, touch_side, touch_percent);;
touch_ids (ppeendp
player_ob) center_x rope_obj. center_x: satisfied, iddi idd.appendd (ppayerrllayr_dll obj.id) id)
return touch_ ids, satisfied_ics ids
Constraint 2: Player's feet must align with the top of the conveyer belt def c2(obj_list: ObjList, touch_side=3, pouch_.erc)nt=0.1) ObjList: touch_ids, satisfied_ids [), player_objs - het.aast.get.cojs....obo..oooeee ('plaaayerr""" 'player'
# get all Obj of type
let.ltttttt.bjasa.... by_ type Obj type 'conveyer_belt"
conveyer_belt,
conveyerr wtt.oo/oppeccocceeeeeeee..eeeee
get all
for player_oby in player_ objs: # player_ole _bjj is of type Obj for conveyer_belt_og in conveyer_belt_objs::
conveyer_belt_oty is of type
Obj
player_obj
ooyyer_bbjttuuches touches (conveyer__ell_oob) touch_side, touch_percent):
touch_ids app.ddd..pppdddcccccc (opperttcovveyer_bbtt.bbjii
if player_oby bottom, side == conveyer belt_ob] wbw.oooojo.tittteeeeeeoee sid_side:
satisfied_ide ittpeddddddaappaddddaeeeeaa booooveeerbblt..eee id)
return touch_ids, satisfied_ids
Constraint 3: Player's feet must align with the top of the platform def c3 _lstt: ObjList, -, touch_side=3, perch.perce) -> ObjList: touch_ids, satisfied_ids [] 0 player_objs obj_list. oot.sst.geet.ooj.....ooeo.eeeeer player'
get all Obj type
objj = obj aaj_liat.get_bjeee............. wtw.ggeett.bbbbbjoaaaeepe obbj__typee pplatform') # get all Obj of type 'platform'
for player_oby player_obja
player_oby
type Obj
for platform_obj in platform_objs: platform_ob, is type Obj
if player_obj touchhs player_ (plathooplatffomm.o. touch_ids souch_ids.append (pateoorppattooo.oo...ii
touch_side touch_percent):
if player_ob) cttr_oobbbbottmmeee.. satisfied append (platform_obj..rk
bottooootsssiooo.ooeeeeoo siide == platform_od cttprrmmooootopssidd.o... tideside:
return touch_ida ids, satisfied_ids
Constraint 4: Player's body must align with the center of the ladder liat: ObjList, touch_ids, satisfied_ids D. player_objs obj.list..et............. byt..ooe._bbb..bbb..rrrr... (ytpebbbrr_..ttppppppaaa"" ('player') # get all Obj of type 'player'
def co
side__ touch_percent=1100) ObjList:
ladder_objj = abj_list.get_oby get.sitteeebbjs.b......... www.byybbjjtyyp. ('ladder') # get all Obj of type
"addder'
for player_ok in player_oby
player_ob obj is of type Obj
for ladder_ob] in ladder_objas
ladder_obj is of type Obj
if player_obj- touches (ladder_ob) (lucees touch_ids (Ladder_ooor
touch_side, touch_percent):
player_ob) center_x ladder_obl center_x: satisfied_idh iepeeddddds.pppddddeeaeeae append (ladder_oby id) return couch_ids, satisfied_ids
Table 10: The collection of constraints synthesized for Montezuma's Revenge. Function names have been shortened to space.
30


Reviews:
[{"review_id": "xB4T4pUk9D", "reviewer": "Reviewer_aKQa", "summary": "This paper introduces a method for learning world models (distributions over future observations given the past) as a product of experts, where each expert is a small Python program that attempts to capture a particular aspect of the world dynamics. The experts are generated by LLMs prompted with a small batch of trajectories. Then, a gradient-based optimizer fits weights to the experts in order to maximize the likelihood of the observed data. Experts with small weights are pruned, and the procedure can repeat while the agent collects more observations from the environment. Experiments on two Atari games, Montezuma Revenge and Atari, show that the learned world model, combined with a planning algorithm, performs significantly better than baselines under the small data regime. The world model is also shown to be useful to pre-train a PPO agent, improving its sample efficiency.", "questions": "* In (3) (L119), is the aggregation of hard constraints not supposed to be a conjunction (and) as opposed to disjunction (or)? If it is really a disjunction, why are they \"hard\" constraints if you only need one of them to be satisfied? (for instance, leaving the screen downwards wouldn't necessarily be a \"hard\" constraint unless all the constraints prevent it?)\n* Did the authors observe any experts that become too slow to compute as the trajectory grows? Was the unpredictability of running times ever a problem for the scalability of the learned world model?\n* I'm confused about the \"Multi-timestep predictions\" description. The text describes the idea here of assuming independence to phrase Multi-timestep predictions as a product of next-timestep experts. But Equation 4 contains $p^{expert}(o_{t+k} | o_{1:t}, a_{1:t})$. When k > 1, this is not a next-timestep expert anymore. Don't you need to also condition on an accumulating sequence of observations and actions (up to $t+k-1$)?\n* L157: serve \"as\" a simulator?\n* How often do the synthesized experts overlap in behavior? Is there any pressure for them to be orthogonal? What happens if the same (useful) expert is generated multiple times - do all of its copies get kept by the weight-based pruning procedure, since they'd have similar weights?\n* Are there other existing environments with less challenging perception where PoE-World would be able to scale better than WorldCoder in terms of the complexity of the world model? I think those could be compelling cases to study, even if the observations are not at the pixel level, and would help address concerns about the generality of the method.\n* How large does the world model get on Pong? Intuitively, this looks like an environment where the dynamics are simple enough for a monolithic program to capture (even if it might not generalize as well to Pong-Alt). Is the main advantage here against WorldCoder the fact that the learning method is compositional, even if the total size of the world model might not be as large as in MR?", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "The paper proposes a novel method that significantly improves the scalability of learning programmatic world models. The scale of the learned world models (4k+ lines of code) is encouraging for this line of work. The idea of using deterministic expert programs, but treating their output as probability distributions and using gradient-based optimization to learn how to weight the experts, and prune them, is technically interesting. The experiments are well-motivated, and the choices of baselines make sense.\n\nMy main high-level concerns are about how specific to Montezuma Revenge (and Pong) were some of the choices that the authors made.\n\nFirst, the paper assumes an existing object detector. This seems quite hard to obtain for most rich domains, and certainly will be a major challenge if we want to eventually extend this line of work to naturalistic settings. The authors note that even for Atari games having an accurate scene parser is difficult. Thus, while I understand that the paper focuses on world modeling and not on perception, the dependency on an object detector might limit how broadly applicable the method ends up being.\n\nIt also seems like the planner might be hard to generalize across domains. In Step 1 in A.2, it's unclear how you obtain the list of all possible abstract states. Is it assumed that all of the relevant objects are given a priori, and that they do not change (e.g., objects appear or disappear)? I'd have liked to see a discussion of which of these assumptions might be violated in other cases (e.g., even in other Atari games), and which are supposed to be reasonable in a wider range of environments.\n\nOverall, I like the idea and the set of experiments, but I'm unclear on how applicable the method currently is in other environments (Atari and otherwise) - how hard it would be to satisfy its assumptions, and whether it could bring gains in cases where perception is not necessarily a significant challenge (but world modeling can still be, regardless). This can be improved by expanding the evaluation in terms of new environments", "was_revised": true, "final_justification": "The authors provided a clear response to most of my concerns. I agree that the world modelling results are interesting and potentially inspiring for other domains, even though their application in particular to Atari problems creates complications that are orthogonal to world modelling per se (as the authors discussed), and thus make the results a bit harder to disentangle than I think is ideal. That said, these are challenging domains and the method is sound, so I recommend acceptance.", "rebuttal": "Thank you for your feedback! We address the raised concerns below:\n\n> Concern 1: The reviewer is concerned that we assume an object detector and a planner, limiting generalization to other domains. The reviewer wonders how hard it is to satisfy these assumptions. (See concern 1.1 and 1.2 for specific assumptions)\n\nWe agree with the reviewer’s concerns on the limitations of used object detectors and planners. It is indeed hard to obtain these assumed components. **We note, however, that these assumptions are not specific to our work and should be seen as generic limitations of symbolic world models line of research.** We elaborate on each assumption below:\n\n> Concern 1.1: The paper assumes an existing object detector\n\nObject detectors are an important assumption, and one often made in robot decision-making. All our baselines use the same structured input representation, rather than pixels. Under this comparison, we make significant progress relative to previous symbolic world models (Worldcoder NeurIPS 24 [1], Code World Models NeurIPS 24 [2]). **While WorldCoder and CWM, both published at NeurIPS 24, work on gridworlds and text domains, we work on OCAtari (Atari + Object detector), which is much more complex.**\n\nOther works learn symbolic world models from pixels [3], but do not scale to dynamics as complex as those we consider. We hope that publishing our paper paves the way for combining our method with such works, which would help address this concern.\n\n> Concern 1.2: seems like the planner might be hard to generalize across domains\n\nPlanning is a very hard problem with its own research community. **The planner used is not part of our method/contribution: we use the same planner for both our work and the baseline (WorldCoder) we are comparing against. In the paper, we denote our agent as PoE-World + Planner to emphasize this point.** We further note that a planner is not always needed: As figure 5 shows, we can run model-free RL (PPO) in the world model and see performance gains. We acknowledge that we use a more domain-specific planner, but this was necessary to check that our world models were accurate enough for long-range planning.\n\n\n\n> Concern 2: Whether it (this work) could bring gains in cases where perception is not necessarily a significant challenge (but world modeling can still be, regardless)\n\nThanks for pointing this out: We focus on world modeling because it still presents a significant challenge even when perception and planning are taken as given, as evidenced by the poor performance of the baselines. Our main experiments (Table 1 and Figure 4) serve to demonstrate that **given reliable perception, PoE-World outperforms other methods, even in the presence of fine-grained physics, stochasticity, and partial observability.**\n\n> Q1: Are there other existing environments with less challenging perception where PoE-World would be able to scale better than WorldCoder in terms of the complexity of the world model?\n\nYes! Robot decision-making, specifically Task and Motion Planning (TAMP), is a prime example. We think Atari is adequate for a first paper, though: Montezuma’s Revenge has served as the centerpiece of several well-known papers [4, 5].\n\n> Q2: How large does the world model get on Pong? ... Is the main advantage here against WorldCoder the fact that the learning method is compositional, even if the total size of the world model might not be as large as in MR?\n\nPong has surprisingly complex physics, requiring a world model with ~3000 lines of code. Compositional factoring allows learning independent programs for each object, and even breaking the dynamics of a single object up into several programs. \n\nTo illustrate the complexity of Pong, we include below two example trajectories from Pong. Here is a trajectory (displaying just the paddle’s velocity y and the action) from a real Pong gameplay (Note that NOOP stands for NO OPeration):\n\n```\npaddle velocity = 0 \n-> action = RIGHT \npaddle velocity = -6 \n-> action = RIGHT \npaddle velocity = -8 \n-> action = NOOP \npaddle velocity = -15 \n-> action = NOOP\npaddle velocity = -2 \n-> action = NOOP\npaddle velocity = -2 \n-> action = NOOP \npaddle velocity = 0\n```\n\nAnd here is a trajectory in the same gameplay but at a different timestep:\n\n\n```\npaddle velocity = 0 \n-> action = RIGHT \npaddle velocity = -3 \n-> action = RIGHT \npaddle velocity = -18 \n-> action = NOOP \npaddle velocity = -8 \n-> action = NOOP\npaddle velocity = -6 \n-> action = NOOP\npaddle velocity = -1 \n-> action = NOOP \npaddle velocity = 0\n```\n\nNotice how the two trajectories lead to different sequences of velocities despite the same action sequences. In fact, in the first trajectory, performing NOOP when velocity = -2 can lead to two possibilities: velocity = 0 or -2. Performing NOOP also does not immediately stop the paddle from moving; the paddle only stops moving (velocity = 0) after a few timesteps.\n\n> Q3: How often do the synthesized experts overlap in behavior? Is there any pressure for them to be orthogonal? What happens if the same (useful) expert is generated multiple times - do all of its copies get kept by the weight-based pruning procedure, since they'd have similar weights?\n\nAll copies do get kept, so there might be redundant experts if similar experts are all synthesized at the same timestep. However, in the next timesteps, the corresponding behavior would be well-explained, so our algorithm would not try to synthesize more experts to explain that behavior.\n\n> Q4: I'm confused about the \"Multi-timestep predictions\" description. The text describes the idea here of assuming independence to phrase Multi-timestep predictions as a product of next-timestep experts. But Equation 4 contains $p^{expert}(o_{t+k}|o_{1:t}, a_{1:t})$. When $k > 1$, this is not a next-timestep expert anymore. Don't you need to also condition on an accumulating sequence of observations and actions (up to $t+k-1$)?\n\nYes, for multi-timestep experts, we make an extra assumption, so we do not condition the predictions on the sequence of observations and actions between time $t$ and $t + k$.\n\n> Q5: Did the authors observe any experts that become too slow to compute as the trajectory grows? Was the unpredictability of running times ever a problem for the scalability of the learned world model?\n\nRuntime is linear with the number of experts, so it is quite predictable. The experts tend to be quite simple, so they run fast. Appendix A5 includes runtime.\n\n> Q6: In (3) (L119), is the aggregation of hard constraints not supposed to be a conjunction (and) as opposed to disjunction (or)? If it is really a disjunction, why are they \"hard\" constraints if you only need one of them to be satisfied? (for instance, leaving the screen downwards wouldn't necessarily be a \"hard\" constraint unless all the constraints prevent it?)\n\nAs discussed in appendix A.1, in Equation (3), we choose to use a disjunction rather than a conjunction because the physics in video games can be peculiar and unrealistic—a player might have their body overlap with a platform when climbing down a ladder attached to that platform. The constraints are still hard in the sense that each of them is deterministic, not probabilistic. In certain situations, we agree that a conjunction could also be a reasonable choice.\n\n**References:**\n\n*[1] Tang, H., Key, D., & Ellis, K. (2024). Worldcoder, a model-based llm agent: Building world models by writing code and interacting with the environment. Advances in Neural Information Processing Systems, 37, 70148-70212.* \n\n*[2] Dainese, N., Merler, M., Alakuijala, M., & Marttinen, P. (2024). Generating code world models with large language models guided by monte carlo tree search. Advances in Neural Information Processing Systems, 37, 60429-60474.*\n\n*[3] Liang, Y., Kumar, N., Tang, H., Weller, A., Tenenbaum, J. B., Silver, T., ... & Ellis, K. (2024). Visualpredicator: Learning abstract world models with neuro-symbolic predicates for robot planning. arXiv preprint arXiv:2410.23156.*\n\n*[4] Kulkarni, T. D., Narasimhan, K., Saeedi, A., & Tenenbaum, J. (2016). Hierarchical deep reinforcement learning: Integrating temporal abstraction and intrinsic motivation. Advances in neural information processing systems, 29.*\n\n*[5] Ecoffet, A., Huizinga, J., Lehman, J., Stanley, K. O., & Clune, J. (2021). First return, then explore. Nature, 590(7847), 580-586.*"}, {"review_id": "JqUQxkEQpU", "reviewer": "Reviewer_RHFT", "summary": "They introduce a world model that uses multiple LLMs to synthesize programs, then optimize on the weights of these programs to describe the rule set behind the observed trajectories. This allows to compose many small programs, as opposed to a large monolithic program to describe the domain, which allows for greater compositionality and extends to more complex domains.", "questions": "- An ask would be to push on compositionality in terms of more rigorous definitions as well as another domain more applicable to evaluating it.\n- A detailed analysis in a domain with more complex underlying rules or is less object-centric and seeing what representation the programs will use in that case would allow me to build confidence about its ability to extend to more domains. \n- Please expand on the partial observability point more. Especially having experiment results to back it up.", "limitations": "I would mention the reliance on object-centricity with one more line. \"Symbolic programs expect symbolic inputs\" is valid, but can be made more clear with a direct mention to where this approach does and doesn't apply.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n- The algorithm is simple and clearly explained.\n- There are clear applications and a need for an approach like this. In terms of program-structured world models, it seems to be the logical next step. I enjoy the references to human cognition.\n\nWeaknesses:\n- There is a lack of novelty of this approach. I see its promise but at its core its WorldCoder but with multiple smaller synthesizers that gets aggregated together. \n- Even though compositionality is a major motivation behind the approach, there were no formal definitions for compositional generalization and experiments really emphasizing that. I wouldn't argue that the alterations made to the two environments is sufficient to show that. However, I do believe the algorithm could extend to more difficult, composition-specific domains.\n- This approach would extend well to domains with large amounts of simple rules that are object-centric. However, when the rules cannot be broken down into such modular, small components, when the scene is not as distinct in terms of its labeled objects, or when the size of the rule set is less than m, I am unsure of how this approach will perform.\n- There were mentions to partial observability but no justification or expansion on it. This is a key point that got completely overlooked", "was_revised": true, "final_justification": "I had concerns about novelty and the difference between WorldCoder, but the rebuttal's addressed my concerns and have updated my view on the originality of this paper. Thus, I have updated my score to an Accept.", "rebuttal": "Thank you for your feedback! We address the raised concerns below:\n\n> Concern 1: There is a lack of novelty of this approach. I see its promise but at its core its WorldCoder but with multiple smaller synthesizers that gets aggregated together.\n\nQuick clarification before the actual reply: it is the programs that get aggregated, not the synthesizers, but this is probably a typo from the reviewer. The actual reply is below:\n\n**We believe our novelty lies not in the idea of aggregation itself**, as we indeed cite many works advocating the idea of composing smaller parts to obtain a sophisticated whole. **Our novelty is in how the aggregation is done:** we turn small, deterministic programs into simple distributions and aggregate them stochastically to obtain sophisticated, stochastic world models. We highlight key differences between WorldCoder and PoE-World in a table below:\n\n| | WorldCoder | PoE-World (ours) |\n|---|---|---|\n| Learning mechanisms | Code generation (with refinement) with LLM | Code generation with LLM AND Gradient-based weight optimization |\n| Representation | Discrete | Mixed Discrete/Continuous |\n| Stochasticity | No, deterministic world models | Yes, fine-grained stochastic world models |\n| Scalability | Low, world models with ~100 lines of code | High, world models with ~4000 lines of code |\n| Most complex domains handled | Gridworlds, text domains | Non-gridworlds, i.e., OCAtari |\n\n> Concern 2: Even though compositionality is a major motivation behind the approach, there were no formal definitions for compositional generalization and experiments really emphasizing that.\n\nBy “compositional generalization”, we mean two things:\n1. Generalization to object combinations not seen during training: what Cosmos, ICLR’24 [1] calls “entity composition”\n2. Factorization of the world model into rules conserved across object types: What Cosmos calls “relational composition”, and what RIMS [2] calls “independent mechanisms”.\n\nExperiments in Table 1 and Figure 4 evaluate both forms of compositionality: The Alt environments evaluate entity composition, and the comparison with the WorldCoder baseline evaluates relational composition. We will revise the paper to clarify this point about compositionality.\n\n> Concern 3: This approach would extend well to domains with large amounts of simple rules that are object-centric. However, when the rules cannot be broken down into such modular, small components, when the scene is not as distinct in terms of its labeled objects, or when the size of the rule set is less than m, I am unsure of how this approach will perform.\n\n**Although the experts in this work are relatively simple, they could be more complex: we can in fact even use WorldCoder wholesale to generate experts for PoE-World.** So even if the environment cannot be described by modular rules, our approach can just fall back to WorldCoder. We think though that many environments are mostly modular, and note that experts can model inter-object interactions.\n\n> Concern 4: There were mentions to partial observability but no justification or expansion on it. This is a key point that got completely overlooked\n\n**OCAtari, the domain we use in this work, is partially observed**: the current scene and the current action alone cannot describe the next scene. This is true *even for Pong.* Looking to the future, the real world is partially-observed, so handling this complexity is an important step for program world models. We will revise the paper to clarify that OCAtari is partially observed.\n\n**References:** \n\n*[1] Sehgal, A., Grayeli, A., Sun, J. J., & Chaudhuri, S. (2024). Neurosymbolic Grounding for Compositional World Models. ICLR.*\n\n*[2] Goyal, A., Lamb, A., Hoffmann, J., Sodhani, S., Levine, S., Bengio, Y., & Schölkopf, B. (2019). Recurrent independent mechanisms. arXiv preprint arXiv:1909.10893.*"}, {"review_id": "QMAyI9kLeI", "reviewer": "Reviewer_nRCp", "summary": "This paper puts forth the idea of modeling the world (environment) as a product of (expert) programs.\nEach program is some deterministic python code, and it predicts one attribute, or state variable, of \nthe environment. A collection of programs generates (zero or more) predictions for each state variable,\nand these predictions can be combined to predict the most likely state of the environment. The weights\nthat are assigned to each Python program can be learned by finding which combination of weights generates\nthe predictions that best match the demonstrations.\n\nEach Python program is generated using a LLM. The weights are learned by gradient descent. I am guessing\nthat some human effort goes into decomposing the world into its constituent state variables.\n\nThe paper shows that their approach of world modeling can yield useful models from just a few demonstrations,\nso, there is a massive saving in the number of data points required. The environment here is taken from \ntwo Atari games. The paper also shows benefits of having a product-of-symbolic-experts model for several\ntasks, such as planning, or predicting the next observation, or RL.", "questions": "(1) Is the summary of the paper above accurate, or have I misrepresented the work?\n(2) Are there cases where decomposing the world model as done in this paper is not a good idea? Even in classical control theory, when the world is modeled as, say, a linear system, dx/dt = Ax + Bu, where x is a vector of state variables, the world model is decomposed into constituent parts since each x_i is predicted separately using a different equation; however, the predictions are still conditioned on the current estimate, x, of the full state of the system. What is the input to the Python program -- I think it is a history of past observations, but I am not sure how long a history. I also understand that the input does not include the current estimate of the state of the system. Is that right?\n(3) The same question as (1), but for the evaluation part of the review (titled the strengths and weaknesses).", "limitations": "Yes.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "The paper describes an interesting piece of work. The claims and results in the paper look reasonable. \nThe evaluation is only on a very small set of worlds, but the gains are clearly seen. \nIn some ways, this paper does a U-turn and goes back to the original way of designing systems in \nsystems theory. We first learn a plant (environment) model, and then use the plant model to design a \ncontroller (planner). The first step is often referred to as system identification. Typically, one starts\nwith linear or nonlinear, discrete-time or continuous-time or hybrid, parametric models for the plant, and then\none learns the parameters from some demonstrations (system identification). Once we have a model of the \nenvironment, one uses it to control or plan or do a host of other things. This paradigm has existed for a long\ntime and served us well. In contrast to this paradigm, a second approach has been to just learn a model that \nperforms the end-to-end task -- bypassing the step of learning a world model and directly learning to control.\n\nThe main innovation here is the idea of using multiple deterministic programs as the world model. There is also\nsome precedence to this idea. David Harel's work on scenario-based programming had a similar flavor. Of course,\nit was before the time we had LLMs that could generate Python programs from observations.\n\nThe results on the \"the role of hard constraints\" are intriguing. I am not surprised that hard constraints\ndo not affect next observation prediction -- I can see that -- but it is surprising that they contribute to planning\nsuccess. Is the planner only using the learned model (learned using the hard constraints), or is it using\nthe learned model plus the hard constraints separately? In other words, are the hard constraints used in two \ndifferent places in the 2nd row of Table 2?\n\nMy overall impression is that the paper makes an interesting contribution that can be potentially useful\nin some cases.", "was_revised": true, "final_justification": "No change to original score. My score is the most positive, perhaps it is a bit lenient. I found the approach of learning multiple small programs and combining them to get a world model interesting, and something that makes sense. The idea also has some validation from the past research going back many years. That is the main reason for keeping the positive rating. I do recognize the challenges in taking this work forward.", "rebuttal": "Thank you for your positive feedback! We are glad the reviewer finds this work interesting. We answer the reviewer’s questions below:\n\n> Question 1: Are there cases where decomposing the world model as done in this paper is not a good idea? Even in classical control theory, when the world is modeled as, say, a linear system, dx/dt = Ax + Bu, where x is a vector of state variables, the world model is decomposed into constituent parts since each x_i is predicted separately using a different equation; however, the predictions are still conditioned on the current estimate, x, of the full state of the system. What is the input to the Python program -- I think it is a history of past observations, but I am not sure how long a history. I also understand that the input does not include the current estimate of the state of the system. Is that right?\n\nWe think there are two different questions here. First, whether decomposing the world model is always a good idea; and second, whether conditioning predictions only on the past history (but not the other state variables) is too strong of an assumption.\n\nFor the first question, we think decomposing the world model should never really hurt – having a single program as a world model is a special case  of PoE-World (with number of experts = 1). \n\nFor the second question, this relates to a research direction we’re currently exploring. Factoring the state such that each state variable does not depend on the current values of other state variables works in this work, because we’re given the ground truth state variables (x/y position/velocity and visibility of each object). \n\nHowever, if we do not have access to these ground truth state variables and we have to *learn* the state variables, then state variables can become entangled (like “player y-velocity” correlates with “player’s feet supported by platform”), so the factored state assumption no longer holds. In our ongoing work, we’ve relaxed the independence assumption and added experts that model correlations between state variables. This gives interesting technical puzzles, as we cannot easily normalize the final probability distribution, so we need to use training techniques from the area of energy-based models (EBM) [1]; we’re currently exploring this direction in follow-up work.\n\n\n> Question 2: I am not surprised that hard constraints do not affect next observation prediction -- I can see that -- but it is surprising that they contribute to planning success. Is the planner only using the learned model (learned using the hard constraints), or is it using the learned model plus the hard constraints separately? In other words, are the hard constraints used in two different places in the 2nd row of Table 2?\n\nThe planner does not use the hard constraints separately; it only uses the world models, whose outputs are affected by the hard constraints. We also find this somewhat surprising, and we suspect that hard constraints tend to play an important role when the world model has to generate a multi-timestep sequence of observations (as opposed to one-step prediction) where errors can accumulate.\n\n> Question 3: (1) Is the summary of the paper above accurate, or have I misrepresented the work?\n\n> And the same question as (1), but for the evaluation part of the review (titled the strengths and weaknesses).\n\nThe written summary is a good and concise summary of our work :)\n\nThe evaluation part is also accurate. We were not aware of the connection to older systems theory work, but the learning algorithm you described does sound very similar to what we have in this paper. We see this work as relating back to many older works but done in the modern age where LLMs can be used as powerful program synthesizers, and we demonstrate how such methods can work on a domain more complex than what symbolic world models have previously considered.\n\n**References:**\n\n*[1] Song, Y., & Kingma, D. P. (2021). How to train your energy-based models. arXiv preprint arXiv:2101.03288.*"}, {"review_id": "hRwa7vflaI", "reviewer": "Reviewer_SJZQ", "summary": "Learning a world model is crucial for building capable AI agents. This work focuses on how to leverage large language models to synthesize a world model utilizing weighted program experts. The program experts are synthesized utilizing a large language model. With only a few observations, this approach can capture the underlying rules of a complex, non-grid environment. The authors demonstrate their methodology on Atari’s Pong and Montezuma’s Revenge.", "questions": "1. How well does the learned world model reflect the actual dynamics of the environment?\n2. How does the world model help downstream policy learning? Can you expand the scope of the Figure 5 results to show this more comprehensively?\n3. Why does the method show different learning behaviors on Pong vs. Montezuma’s Revenge?\n4. What does the discrepancy indicate for the general applicability and limitations of the approach?\n5. In real-world settings, how would one identify or acquire the required objects, actions, and possible interaction results to build the program experts?\n6. How are the expert programs synthesized?\n7. How is the world domain defined and incorporated into the LLM?\n8. What prompt strategy is used to guide the synthesis?", "limitations": "Yes.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Originality: 4/5**\n\nThis paper presents an innovative and refreshing approach to synthesizing an environment using a product of program experts. The generated rules effectively capture the semantics of the environment and enable learning from only a few observations. \n\n**Quality: 2.5/5**\n\nI have some concerns about the experimental setup. The main goals should be to demonstrate (1) how well the world model reflects the actual dynamics of the environment, and (2) how the world model supports the learning process of downstream models. However, I interpret Figure 4 primarily as evidence of data efficiency, without providing clear insights into research questions (1) or (2). On the other hand, Figure 5 presents interesting results relevant to research question (2), but it is limited in scope, as it shows only one learning policy in a single learning environment. Table 1 illustrates efficient adaptation to new environments, which is promising, but I believe Figure 5’s setup is a more direct response to research question (2). Since you have access to the ground truth setup, adding an explicit evaluation addressing research question (1) should be feasible. Additionally, extending the results shown in Figure 5 would strengthen the paper’s overall illustration of its core ideas. \n\nFurther, the Pong game and Montezuma’s Revenge exhibit very different behaviors in terms of learning: the PPO performs better in Pong, whereas in Montezuma’s Revenge the other methods fail to learn. While adding another experiment on a different game may not be feasible during the rebuttal period, I would appreciate more insight into why this discrepancy occurs and how it reflects on the general applicability and limitations of your approach.\n\n**Significance: 3/5**\n\nThe world modeling is a crucial topic for AI agents building. While program synthesis could be an important building block. However, applying this work to general case could be challenging on how do people know the objects, actions, and potential interaction results? Can you elaborate further on how to get these building blocks in a real world scenario? \n\n**Clarity: 3/5*\n\nThe paper reads well and the images are very helpful for illustration. However, can you elaborate more on how the expert programs are synthesized? Specifically, how is the world domain defined and passed into the large language model? What is the prompt strategy?", "was_revised": true, "final_justification": "The authors has addressed all my concerns. I have raised my score accordingly.", "rebuttal": "Thank you for your feedback! We address the raised concerns below:\n\n> Concern 1: The main goals should be to demonstrate (1) how well the world model reflects the actual dynamics of the environment… However, I interpret Figure 4 primarily as evidence of data efficiency\n\nWe believe **Table 2** (Next observation prediction accuracy vs baseline) and **Table 3** (Next observation’s object attributes prediction accuracy vs baseline) is what the reviewer is looking for.\n\n> Concern 2: (2) how the world model supports the learning process of downstream models… Figure 5 presents interesting results relevant to research question (2), but it is limited in scope\n\n**We consider a broader range of ways world models can support decision-making: Both learning (RL) AND planning.** Our main experiments, Table 1 and Figure 4, analyze planning, while Figure 5 shows how the world model supports downstream learning via RL.\n\nWe focus more on planning than RL learning because planning is usually seen as having stricter requirements on the accuracy of world models, as errors can accumulate when an agent plans a long sequence of actions before acting. Pre-training model-free RL on world models (like Dreamer [1]) is usually easier to get to work and see gains of world models.\n\n> Concern 3: Further, the Pong game and Montezuma’s Revenge exhibit very different behaviors in terms of learning: the PPO performs better in Pong, whereas in Montezuma’s Revenge the other methods fail to learn.\n\n**This is because Montezuma’s Revenge is much harder than Pong. There is a consensus among RL researchers that Montezuma’s Revenge is one of the hardest Atari games for RL.** Go-Explore, Nature 2021 [2] calls it a *grand challenge*. Many papers explicitly mention Montezuma’s Revenge in their title/abstract [3, 4]. Reviewer LjAF notes that the Montezuma’s Revenge result is impressive.\n\nModel-free RL generally fails on Montezuma's Revenge, unless trained on an exorbitantly large number of environment interactions.\n\n> Concern 4: program synthesis could be an important building block. However, applying this work to a general case could be challenging on how do people know the objects, actions, and potential interaction results? Can you elaborate further on how to get these building blocks in a real world scenario?\n\nAll these building blocks (perception, planning/acting, and world modeling) have their own dedicated research areas, and one paper cannot do everything at once. We believe that building AI agents that can learn, reason, and adapt flexibly in the real world requires progress on all of these building blocks, but in this work, we target the world modeling problem.\n\nThen, to get each of these building blocks to work in a real world scenario, we think researchers should gradually work on more and more complex domains over time. Previous work in symbolic world models only considered gridworlds and text worlds. We go beyond that and work on non-gridworlds like OCAtari: Easier than raw Atari, but much more complex than gridworlds. Ultimately, these games serve as testing grounds for research, and this work should be seen as a stepping stone toward more general agents.\n\n> Concern 5: Can you elaborate more on how the expert programs are synthesized? Specifically, how is the world domain defined and passed into the large language model? What is the prompt strategy?\n\nWe have put details on how programs are synthesized and how input is turned into natural language to be passed into LLMs in Appendix A.1.\n\n> Q1: What does the discrepancy [between Pong and Montezuma’s Revenge] indicate for the general applicability and limitations of the approach?\n\nMontezuma’s Revenge is a long-horizon sparse-reward problem (see our reply to concern 3). The discrepancy indicates that model-free RL struggles in such settings (relative to planning). Although this is well known, it strengthens the argument for learning a world model: Only with such a model is planning possible.\n\n**References:**\n\n*[1] Hafner, D., Pasukonis, J., Ba, J., & Lillicrap, T. (2025). Mastering diverse control tasks through world models. Nature, 1-7.*\n\n*[2] Ecoffet, A., Huizinga, J., Lehman, J., Stanley, K. O., & Clune, J. (2021). First return, then explore. Nature, 590(7847), 580-586.*\n\n*[3] Salimans, T., & Chen, R. (2018). Learning Montezuma's revenge from a single demonstration. arXiv preprint arXiv:1812.03381.*\n\n*[4] Aytar, Y., Pfaff, T., Budden, D., Paine, T., Wang, Z., & De Freitas, N. (2018). Playing hard exploration games by watching youtube. Advances in neural information processing systems, 31.*"}, {"review_id": "dSt7ZLr8Qy", "reviewer": "Reviewer_LjAF", "summary": "The authors propose a program-synthesis approach to building a world model (i.e. transition function). Instead of attempting to write one large program to capture all the rules of the environment under consideration, they divide up the problem by having many \"experts\" (i.e. smaller programs) together define a distribution over the next state. They evaluate their approach on Atari's Pong and Montezuma's Revenge by looking at performance in these environments as well as next observation and next observation object attribute accuracies. They find the proposed approach to generally outperform all baselines, at least in the limited interaction setting (i.e. 100k env steps).", "questions": "1. In the definition of $Z_f$ in the equation after line 105, is it tractable to iterate over $o^f$? I guess the authors are using some discretized values with some range for velocity etc? And if not, how are continuous values handled?\n2. On line 132, how expensive is it to repeat step 1 - 3 every time the observed trajectories get updated? \n3. Line 140: \"x/y velocity and visibility\". Do the observations not include bounding boxes / positions of objects? And what does visibility mean?\n4. In Table 2 and 3, why are the prediction accuracies typically better on the alt environments? This feels a bit counterintuitive.\n5. On line 226, are the 4000+ lines of code \"active\" or is most of it pruned? And also, how many experts does this typically consist of? It would be great if the authors could give some insight / statistics like this on the resulting world model programs.", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "## Quality\nI generally found the quality of the paper to be good. I verified all equations in the paper are correct, and I appreciate the experimental reporting includes error bands and other measures of uncertainty (though it would be great if the authors clarified what this is). In addition, I appreciate the authors going in and altering the environments to test for generalization.\n\nSome things that could further improve the quality of the paper:\n1. I understand getting this method to run even on a single Atari game requires substantial work. That being said, only two Atari games does feel on the low side. Maybe the authors could add a few environments from the Atari-5 subset: https://arxiv.org/abs/2210.02019. \n2. In table 1, there are no baselines that also use the demonstration that is used for PoE. It would be good if the authors could include a behavioral cloning baseline as well as maybe one offline RL baseline that uses this demonstration.\n\n## Clarity\nI generally found the clarity of the paper to be lacking. One of the biggest concerns I have is that a lot of the core parts of building the programmatic world model is left to the appendix. I understand the authors would do this for the planning part which is not core to the contribution, but the way to construct the programs that are used as experts is completely missing in the main body of the paper. This feels strange as it is core to the method. \n\nIn addition, I found the paper lacking in clarity in the following (sometimes more minor) places:\n1. Line 105: does \"projection of $o$ onto feature $f$\" mean \"indexing $o$ with $f$\"? From what I understood, $o$ is a vector of state attributes (e.g. velocity) and so $f$ just serves as an index into this vector? If this is the case, it might make sense to reword this to improve clarity.\n2. I found the paragraph starting on line 106 a bit confusing. Could the authors make the argument being made here a bit more precise / formal? Do the authors mean instead of conditioning on $o_{1:t}$, using some latent variable $z_t$ to represent this? And then the issue would be that the dimensions of $z_t$ might not be interpretable anymore?\n3. The hard constraints paragraph starting on line 114 doesn't feel expansive enough to understand what's really going on or why the hard constraints are necessary. Including an example of a hard constraint here would go a long way.\n4. Are the multi-timestep predictions being discussed in equation 4 used anywhere? If not, I'd recommend leaving this part out and use the space it provides to address some of the other comments.\n5. Line 144: \"We then add noise to the distributions ...\". What kind of noise?\n6. Why does the equation on line 166 not use the same independence assumptions as the one on line 4?\n7. Line 172: how exactly are the demonstrations created? Are they random agents? And if not, how much work went into creating these demonstrations? The details here could reflect on the practicality and assumptions of the method.\n8. Figure 3: it would be helpful here if screenshots of the original environments were shown as well so it's easy to see the differences.\n9. Are the numbers after +- in the various tables standard errors? Or standard deviations? Or something else? And how many seeds were used to compute these? And does each seed regenerate the programmatic experts?\n10. In Table 2 and 3, how are the prediction accuracies defined? Are they exact match?\n11. Table 4 reports \"planning successes\" but as far as I could tell this is never defined. What does \"planning successes\" mean? \n\n## Significance\nThe results seem good, especially the ones on Montezuma's revenge are impressive! However, my biggest concern is that while the authors evaluate on Atari games, they end up working on a simplified state space, not on the direct pixel space. So part of my concern here is: how different is working on this \"simplified\" Atari setting from working on a complex grid world? Is it the dynamics? The number of objects? It would be good to get some sense here in order to understand how big of a step the authors are taking away from toy grid worlds.\n\n## Originality\nWhile there has been work on building world models with programs, it is certainly a challenge to get them to work on more complex environments. This combined with the modular expert program approach makes the paper feel pretty original. I haven't seen any other work trying to build world models in this way.", "was_revised": true, "final_justification": "I thank the reviewers for their careful response, which has cleared up many confusions. While I think the paper has substantial merit, I maintain my current score (which is leaning positive) because of the following standing limitations:\n1. **Limited evaluation**: as pointed out by other reviewers as well, the paper only uses two Atari games to perform all of their evaluation. To convincingly show the paper's assumptions are not too narrow to be at least somewhat broadly applicable, I think a few more Atari games would go a long way.\n2. **Clarity**: while I understand the reviewers cannot change the current pdf, I would argue the paper would greatly benefit from another pass on the writing. The biggest ticket here would be to incorporate the core parts of building the programmatic world model, which is currently in the appendix. Reviewer SJZQ has also pointed this out. In addition, things like including the original screenshots in Figure 3, clarifying how exactly the demonstrations were collected, etc. would all help push the paper to the next score.", "rebuttal": "Thank you for your positive feedback! We are glad the reviewer finds our Montezuma's Revenge result impressive. We address the raised concerns below:\n\n> Concern 1: My biggest concern is that while the authors evaluate on Atari games, they end up working on a simplified state space, not on the direct pixel space. So part of my concern here is: how different is working on this \"simplified\" Atari setting from working on a complex grid world? Is it the dynamics? The number of objects? It would be good to get some sense here in order to understand how big of a step the authors are taking away from toy grid worlds.\n\n**OCAtari is much “noisier” than gridworlds, and much higher-resolution, involving fine-grained motion and physics. It is also partially observed.** Handling partially observed stochastic environments with fine-grained physics is an important step toward making these techniques applicable in the real world.\n\nTo illustrate the complexity of OCAtari, we include below two example trajectories from Pong. Here is a trajectory (displaying just the Pong paddle’s velocity y and the action) from a real Pong gameplay (Note that NOOP stands for NO OPeration):\n\n```\npaddle velocity = 0 \n-> action = RIGHT \npaddle velocity = -6 \n-> action = RIGHT \npaddle velocity = -8 \n-> action = NOOP \npaddle velocity = -15 \n-> action = NOOP\npaddle velocity = -2 \n-> action = NOOP\npaddle velocity = -2 \n-> action = NOOP \npaddle velocity = 0\n```\n\nAnd here is a trajectory in the same gameplay but at a different timestep:\n\n\n```\npaddle velocity = 0 \n-> action = RIGHT \npaddle velocity = -3 \n-> action = RIGHT \npaddle velocity = -18 \n-> action = NOOP \npaddle velocity = -8 \n-> action = NOOP\npaddle velocity = -6 \n-> action = NOOP\npaddle velocity = -1 \n-> action = NOOP \npaddle velocity = 0\n```\n\nNotice how the two trajectories lead to different sequences of velocities despite the same action sequences. In fact, in the first trajectory, performing NOOP when velocity = -2 can lead to two possibilities: velocity = 0 or -2. Performing NOOP also does not immediately stop the paddle from moving; the paddle only stops moving (velocity = 0) after a few timesteps.\n\n> Concern 2: One of the biggest concerns I have is that a lot of the core parts of building the programmatic world model is left to the appendix. I understand the authors would do this for the planning part which is not core to the contribution, but the way to construct the programs that are used as experts is completely missing in the main body of the paper. This feels strange as it is core to the method.\n\nThank you for pointing this out. If accepted we will use our extra page to move those details to the main text.\n\n> Concern 3: Some things that could further improve the quality of the paper:\n> 1. I understand getting this method to run even on a single Atari game requires substantial work. That being said, only two Atari games does feel on the low side. Maybe the authors could add a few environments from the Atari-5 subset\n> 2. In table 1, there are no baselines that also use the demonstration that is used for PoE. It would be good if the authors could include a behavioral cloning baseline as well as maybe one offline RL baseline that uses this demonstration.\n\nWe agree with the reviewer. We were not aware of Atari-5 when we wrote the paper, and it would have been a good addition to the paper. In table 1, WorldCoder also uses the same demonstration that is used for PoE-World, but yes, we use an on-policy algorithm (PPO) as a model-free RL baseline, so this baseline could not take advantage of the demonstration. An offline RL baseline may help with Pong, as the demonstration shows the player scoring once, but for Montezuma’s Revenge, the demonstration never shows positive score, so we do not expect off-policy RL to help in this case. We believe that the experimental results will convey the same message even with an off-policy RL baseline that uses the demonstration. \n\n> Concern 4: I found the paper lacking in clarity in the following (sometimes more minor) places:\n> 1. Line 105: does \"projection of $o$ onto feature $f$\" mean \"indexing $o$ with $f$\"? From what I understood, $o$ is a vector of state attributes (e.g. velocity) and so $f$ just serves as an index into this vector? If this is the case, it might make sense to reword this to improve clarity.\n> 2. I found the paragraph starting on line 106 a bit confusing. Could the authors make the argument being made here a bit more precise / formal? Do the authors mean instead of conditioning on $o_{1:t}$, using some latent variable $z_t$ to represent this?\n> 3. The hard constraints paragraph starting on line 114 doesn't feel expansive enough to understand what's really going on or why the hard constraints are necessary. Including an example of a hard constraint here would go a long way.\n> 4. Are the multi-timestep predictions being discussed in equation 4 used anywhere? If not, I'd recommend leaving this part out and use the space it provides to address some of the other comments.\n> 5. Line 144: \"We then add noise to the distributions ...\". What kind of noise?\n> 6. Why does the equation on line 166 not use the same independence assumptions as the one on line 4?\n> 7. Line 172: how exactly are the demonstrations created? Are they random agents? And if not, how much work went into creating these demonstrations? The details here could reflect on the practicality and assumptions of the method.\n> 8. Figure 3: it would be helpful here if screenshots of the original environments were shown as well so it's easy to see the differences.\n> 9. Are the numbers after +- in the various tables standard errors? Or standard deviations? Or something else? And how many seeds were used to compute these? And does each seed regenerate the programmatic experts?\n> 10. In Table 2 and 3, how are the prediction accuracies defined? Are they exact match?\n> 11. Table 4 reports \"planning successes\" but as far as I could tell this is never defined. What does \"planning successes\" mean?\n\n1. Yes, we will reword it.\n2. Yes, The POMDP formulation would condition on a latent variable $z_t$ instead of $o_{1:t}$. The issue is that $z_t$ would be an output of a learned component; it would be an output of $z_t = f_\\phi (o_{1:t-1}, a_{1:t-1})$ where we are learning $f_\\phi$. Since $f_\\phi$ would be changing, $z_t$ would also change, and so all $p^{expert}_i$, which takes $z_t$ as input, needs to change correspondingly\n\n3. We display the learned constraints in Table 10. An example is something like player_obj.bottom_side == conveyer_belt_obj.top_side\n4. It is used in the PoE-World algorithm evaluated in the paper.\n5. Small epsilon value: 1e-6.\n6. Line 166 is meant to formalize how world models can be used in general (not specific to our work), so we did not include the independence assumption we use in equation 4 there.\n7. They are not random agents. We do a few iterations of demonstrations to make sure it covers all the causal rules of the environments. The important thing here is that the demonstration shows examples of the causal rules needed to do well in the environment (for example, the player falls when its feet are not supported, touching the skull kills the player, etc.). Creating this demonstration is not hard for humans, but a random agent would not be sufficient.\n8. Thank you, that is a good point. We will try to include the original screenshots in Figure 3 if there is enough space, if not, we will put them in the appendix.\n9. They are standard errors. We use three seeds. Each seed does regenerate the programmatic experts :) \n10. Yes, they are exact match\n11. Thank you for pointing it out, we will fix it. Planning successes for Montezuma’s Revenge mean successfully grabbing the key.\n\n> Q1: In the definition of $Z_f$ in the equation after line 105, is it tractable to iterate over $o_f$? I guess the authors are using some discretized values with some range for velocity etc? And if not, how are continuous values handled?\n\nIn OCAtari, the velocities are discrete, corresponding to per-pixel movement\n\n> Q2: On line 132, how expensive is it to repeat step 1 - 3 every time the observed trajectories get updated?\n\nNot expensive – we optimize the code so that it only synthesizes experts for newly added observations, and also we skip synthesizing experts for observations that are already well-explained by the current world model. \n\n> Q3: Line 140: \"x/y velocity and visibility\". Do the observations not include bounding boxes / positions of objects? And what does visibility mean?\n\nThe world model takes as input an observation with bounding boxes and positions, and so it only needs to predict velocities and visibility (whether the object gets deleted). Object sizes do not change in Pong and Montezuma’s Revenge, so world models do not need to predict bounding boxes. \n\n> Q4: In Table 2 and 3, why are the prediction accuracies typically better on the alt environments? This feels a bit counterintuitive.\n\nThe prediction accuracies are on random frames (generated by random actions). It is possible that random actions in the original environment may lead to unseen/unusual situations easier than in the alternative environments.\n> Q5: On line 226, are the 4000+ lines of code \"active\" or is most of it pruned? And also, how many experts does this typically consist of? It would be great if the authors could give some insight / statistics like this on the resulting world model programs.\n\nThey are all active (the pruned programs are deleted from the world model and do not count words in the 4000+ line of code). There are usually around 300-500 experts."}]

Author rebuttal:

Meta-review: This paper focuses on an interesting direction - program-based world models. Specifically, the paper proposes to construct many small programs which are composable and uses an expert model with learnable weights to predict the future. The paper proposes a clearly novel paradigm and distinguishes itself from other standard world models. The symbolic nature can make it more data efficient. There are potential concerns of the paper such as bespoking and generalizability (Reivewer aKQa), experiments setups (reviewer SJZQ) and its evaluation on the symbolic world (reiviewer LjAF). With the more capable LLMs, exploring a programmable world model can have some nice impacts, hence I recommended acceptance.
Decision: Accept (spotlight)