Title: PoE-World: Compositional World Modeling with Products of Programmatic Experts
Venue: NeurIPS 2025
Accepted: yes
Keywords: program synthesis

compositionality

world modeling
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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).\n\nSTRENGTHS AND WEAKNESSES:\n## 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.\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nLearning 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.\n\nSTRENGTHS AND WEAKNESSES:\n**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?\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nThe 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.\n\nQUESTIONS:\n(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).\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThey 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\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\n\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\nI 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nThe 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\n\nQUESTIONS:\n* 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?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-03"}