PaperlayerView exact revision on arXiv

Source-linked AI summary

PoE-World: Compositional World Modeling with Products of Programmatic Experts

Wasu Top Piriyakulkij, Yichao Liang, Hao Tang, Adrian Weller, Marta Kryven, Kevin Ellis

arXiv:2505.10819v4cs.AIcs.LG

TL;DR

Deep neural world models require extensive data, while prior programmatic world models have not scaled beyond simple text-based and gridworld environments. PoE-World addresses this gap by synthesizing probabilistic products of small programmatic experts, achieving efficient learning and zero-shot generalization in Atari planning tasks. Its scope is limited by reliance on symbolic observations and focus on world modeling rather than exploration, reward learning, or planning itself.

  • Problem

    Existing neural world models require enormous training data, while programmatic world models have been limited beyond simple text-based and gridworld environments.

    Source: §1 · pp. 1–2

  • Method

    PoE-World decomposes symbolic world modeling into many small LLM-synthesized programs whose probabilistic product predicts future observations and supports planning.

    Source: §1 · pp. 1–2, §3.2 · p. 5, §3.3 · p. 6

  • Results

    PoE-World learns complex Atari world models from brief demonstrations, synthesizes 4000+ line programs that generalize zero-shot to novel levels and variations, and performs best across evaluated environments in the low-data regime.

    Source: §1 · p. 2, §4 · p. 7

  • Takeaways & Limitations

    Programmatic expert composition enables scalable symbolic world models for complex Atari environments while supporting planning and simulated pretraining for deep reinforcement learning.

    Source: §1 · p. 2, §3.3 · p. 6, §4 · p. 8

    PoE-World assumes symbolic inputs from an object detector, does not address exploration, decision-making, or reward-function learning, and requires custom object-detection patches for the evaluated Atari games.

    Source: §4 · p. 6, §6 · p. 10

Abstract

from arXiv · show

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. We release our code and display the learned world models and videos of the agent's gameplay at https://topwasu.github.io/poe-world.

1 Introduction

PoE-World addresses the difficulty of learning flexible, data-efficient world models beyond gridworlds by decomposing world-model learning into modular programmatic experts. The resulting symbolic models support stochastic, partially observable Atari environments and planning from brief demonstrations.

1 Introduction

  • Programmatic world models are more data-efficient than neural models but previously struggled beyond simple gridworlds because they search for a single large program.Neural world models demand enormous training data, while symbolic models face discrete combinatorial search over monolithic programs.Source: §1 · p. 1
  • PoE-World decomposes world-model learning into hundreds of small programs encoding different causal laws, rather than searching for one monolithic program.The programs are probabilistically aggregated to predict future observations, making the representation more modular and learnable.Source: §1 · pp. 1–2
  • PoE-World models stochasticity and partial observability and supports planning and reinforcement learning from brief demonstrations in complex Atari games.Its product of programs is probabilistic, and the learned models can support planning in environments such as Montezuma’s Revenge.Source: §1 · p. 2, Figure 1 · p. 2
  • The system models fine-grained pixel-level movement but assumes symbolic observations from an object detector rather than modeling pixel-level visual appearance.It focuses on learning from demonstrated trajectories and on world modeling, not efficient exploration or the computational problem of planning itself.Source: §1 · p. 2
  • PoE-World synthesizes 4000+ line programs that generalize zero-shot to novel Atari levels and game variations, outperforming deep RL in learning efficiency and improving symbolic-model scalability.The empirical study covers Pong and Montezuma’s Revenge.Source: §1 · p. 2

2 Background: World Model Learning

World-model learning approximates unknown environment dynamics from observed trajectories so agents can plan or train policies. Prior approaches use neural networks or LLM-synthesized programs, trading flexibility and data efficiency against sample efficiency, generalization, or scalability limits.

2 Background: World Model Learning

  • The learned world model approximates unknown environment dynamics from observation-action histories and is later used for lookahead planning or reinforcement-learning training.The formulation treats the environment as a sequential decision-making problem with observations, actions, dynamics, and rewards.Source: §2 · p. 3
  • World-model learning is formulated as empirical risk minimization over observed trajectories, using losses such as negative log likelihood.The loss evaluates predicted next observations conditioned on prior observations and actions.Source: §2 · p. 3
  • Deep neural world models optimize parametric models with gradient descent but can suffer poor sample efficiency and generalization, including physically implausible predictions after extensive training.One cited failure mode is imagining that a player can walk through walls or teleport after almost 100 hours of observations.Source: §2 · p. 3
  • LLM-based code world models are more sample-efficient and extrapolate more systematically than deep-learning approaches, but had not succeeded beyond simple text-based and gridworld games.These methods search over programmatic world models rather than neural parameterizations.Source: §2 · p. 3

3 Modeling the World as a Product of Programmatic Experts (PoE-World)

PoE-World represents world dynamics as a weighted product of modular programmatic experts, learns them from sparse trajectories, and uses the resulting model for simulation and hierarchical planning.

3.1 World Model Representation: Product of Programmatic Experts (PoE-World)

  • PoE-World represents world dynamics as an exponentially weighted product of programmatic experts, enabling modular composition of many small programs into a full world model.Each expert expresses particular environmental dynamics, while associated scalar weights determine their contribution.Source: §3 · p. 3, §3.1 · p. 3
  • A factored object-centric state and full-history conditioning make inference tractable and allow independent mechanisms to be learned without globally entangling experts.Feature-level conditional independence provides separate normalizing constants, while history conditioning avoids updating every expert when new latent variables are introduced.Source: §3.1 · p. 4

3.2 World Model Learning: Program Synthesis and Weight Optimization

  • LLMs synthesize simple deterministic Python experts from small transition batches, which are interpreted as noisy attribute distributions and iteratively weighted, pruned, and updated as trajectories grow.The procedure fits expert weights with a gradient-based optimizer, removes weights below threshold δ, and repeats after new observations.Source: §3.2 · p. 5

3.3 World Model Usage: RL in Simulation and Planning with World Model

  • The learned world model supports reinforcement-learning simulation and lookahead planning, including a hierarchical planner that converts abstract contact-state plans into Atari button presses.The planner first searches an abstract graph of object-contact states, then lowers the selected path into low-level actions using the learned model.Source: §3.3 · pp. 5–6, Figure 3 · p. 6

4 Experimental Results

PoE-World is evaluated on partially observed Pong and Montezuma’s Revenge using brief demonstrations, alternate game levels, prediction tests, and planning-based agents. It outperforms baselines in low-data performance, generalizes without demonstrations to altered environments, and benefits from pretraining and hard constraints.

Experimental Setup

  • The evaluation uses fewer than 1000-frame, unsuccessful demonstrations, including a Montezuma’s Revenge demonstration with no positive reward.OCAtari required game-specific manual patches, limiting application to the full Atari suite.Source: §4 · p. 6

Compositional Generalization

  • The alternate environments recombine objects by adding balls and enemies in Pong-Alt and enemies, ladders, and a rearranged map in Montezuma’s Revenge-Alt, without new demonstrations.These tests evaluate compositional extrapolation from models initialized on the corresponding base games.Source: §4 · p. 6, Figure 4 · p. 7

Agent Results

  • PoE-World + Planner performs best across all four environments in the low-data regime and is the only method earning positive reward on both Montezuma’s Revenge versions.PPO requires over one million training steps to surpass PoE-World + Planner and never achieves positive Montezuma’s Revenge score after 20M steps.Source: §4 · p. 7, Table 1 · p. 7

Policy Pretraining

  • Policy pretraining inside PoE-World reduces the training needed to outperform a random agent from 1M PPO steps to 200k, while both policies converge asymptotically.The result characterizes world-model pretraining as a warm start rather than a change in final policy value.Source: §4 · p. 8, Figure 6 · p. 8

Prediction Accuracy

  • PoE-World outperforms symbolic baselines on most next-observation and object-attribute prediction settings, except random-frame Pong where all methods perform similarly.The Pong exception occurs because random agents rarely hit the ball, so game mechanics are not comprehensively tested.Source: §4 · p. 8, Table 2 · p. 9, Table 3 · p. 9

Hard Constraints

  • Removing hard constraints lowers planning performance on both Montezuma’s Revenge versions without significantly changing next-observation accuracy.The authors interpret hard constraints as damage control that makes imperfect predictions usable for long-horizon planning.Source: §4 · pp. 8–9, Table 4 · p. 9

World-Model Quality

  • PoE-World models exceed WorldCoder’s scalability and causal fidelity, using 4000+ lines for Montezuma’s Revenge versus fewer than 100 and avoiding unconstrained or hallucinated mechanics.Low-weight experts can be pruned, yielding more precise world models than a single monolithic program.Source: §4 · pp. 7, 9

5 Related Work

PoE-World extends programmatic world modeling beyond monolithic or restricted representations by combining symbolic programs with modular and hierarchical planning ideas.

  • WorldCoder [3] and CodeWorldModels [24] synthesize single Python transition programs, whereas PoE-World uses substantially more code to model complex environments with partial observability and non-determinism.Source: §5 · p. 9
  • Recent symbolic world-modeling work also synthesizes high-level programs for robotic planning, offering potential synergy with PoE-World’s low-level object-motion models.Source: §5 · p. 9
  • Earlier systems learn world models in restricted domain-specific languages, often on benchmarks co-designed with those languages, while other work generates programs primarily from natural-language instructions.Source: §5 · pp. 9–10
  • PoE-World’s abstract planning uses object-contact relations to define symbolic states, connecting it to hierarchical and task-and-motion planning approaches [46] [47].Source: §5 · p. 10
  • Its specialized programmatic experts and hierarchical representations align with cognitive accounts of interacting agents, event segmentation, human concept learning, and motion-based state spaces [52].Source: §5 · p. 10

6 Discussion and Limitations

PoE-World argues that compositional factoring makes symbolic world models scalable, while acknowledging that its current scope excludes pixel-level learning, exploration, and reward learning.

Discussion

  • Compositionality lets PoE-World recombine small experts to generalize to novel entity and relational compositions, whereas monolithic program synthesis is viewed as intractable even for Atari.Source: §6 · p. 10

Limitations

  • PoE-World does not learn directly from pixels, because its symbolic programs require symbolic inputs from an object detector.Source: §6 · p. 10
  • The system addresses world modeling rather than the full model-based reinforcement-learning problem, omitting exploration, decision-making, and reward-function learning.Source: §6 · p. 10
  • The authors speculate that interpretable program-structured beliefs could support future exploration methods based on software testing and program analysis.Source: §6 · p. 10

A.1 PoE-World algorithm details

The PoE-World algorithm converts trajectory transitions into modular LLM-synthesized experts, interprets them probabilistically, and optimizes their weights before pruning weak components.

Optimization

  • The algorithm alternates batch expert synthesis with gradient-based optimization of predicted next-observation distributions.Source: §A.1 · pp. 15–16

Programmatic expert synthesis

  • Each transition sequence is verbalized as an initial object list, action sequence, and subsequent object-list changes so LLMs can propose causal explanations and programs.Source: §A.1 · p. 15
  • Restricting each textual representation to one changing object type enforces modularity, with experts assigning non-uniform distributions to attributes of that type.Source: §A.1 · p. 15
  • PoE-World uses ten LLM-based synthesis modules per object type to generate experts for distinct environmental aspects, including passive motion, interactions, and object creation or deletion.Source: §A.1 · p. 15

Optimization

  • Expert weights are optimized with L-BFGS and L1 regularization, without incorporating hard constraints into the gradient optimization.Source: §A.1 · p. 16

Pruning

  • After optimization, programs below δ = 0.01 are pruned, while contradictory or minimally explanatory constraints are removed.Source: §A.1 · p. 16

A.2 Planner

The planner combines an abstract contact-based graph with low-level motion search, then replans online while accounting for stochastic outcomes and long Atari action sequences.

Planning motivation

  • Montezuma’s Revenge requires at least 100 actions before the first positive reward, creating an 8^100 primitive-action search space.Source: §A.2 · p. 16

Hierarchical planning

  • The hierarchical planner builds an abstract graph over object-contact states, searches it for a goal path, and executes the resulting subgoals through a low-level planning agent.Source: §A.2 · p. 16
  • Graph construction tests traversability between node pairs with a low-level planner, while breadth-first search selects a shortest abstract path and false edges are removed after failed execution.Source: §A.2 · pp. 16–17

Low-level motion planner

  • MCTS uses Manhattan-distance heuristics, maximizes child values during backpropagation, and varies its exploration parameter across games and iterations.Source: §A.2 · p. 17
  • Sticky actions extend the search space with repeated primitive-action chunks of lengths 1, 4, and 8, yielding 3n actions for n primitive actions.Source: §A.2 · p. 17

Low-level planning agent

  • The online agent executes plans and replans as needed, but stochastic planning can produce different scores from the same world model, so experiments run the planner three times.Source: §A.2 · p. 17

A.3 Domain details

The experiments use object-centric Atari environments and construct alternative levels to test compositional generalization.

  • Experiments use ALE with frameskip 3 and OCAtari object parsing, including categories, bounding boxes, and velocities.Source: §A.3 · p. 17
  • The implementation fixes OCAtari object-detection and bounding-box bugs so detected interactions correspond to touching or overlap.Source: §A.3 · p. 17
  • Pong-Alt and Montezuma’s Revenge-Alt are introduced as alternative environments for evaluating rearranged game configurations.Source: §A.3 · p. 18
  • Pong-Alt layers three Pong environments, producing one synchronized player, three enemies, and three balls.Source: §A.3 · p. 18
  • Montezuma’s Revenge-Alt stacks and flips lower-room sections to create platforms and stairs requiring jumps over skulls.Source: §A.3 · p. 18

A.4 Baseline details

The baseline comparison includes PPO, ReAct, and WorldCoder, with each adapted to the Atari evaluation setting using specified preprocessing or planning choices.

PPO

  • PPO uses four-frame stacking and is fine-tuned on alternative environments from models pretrained for 20M base-environment steps.Source: §A.4 · p. 18

ReAct

  • ReAct prompts an LLM with object positions, pairwise interactions, and the four most recent observations, thoughts, and actions.Source: §A.4 · p. 18

WorldCoder

  • WorldCoder uses tailored prompts for object-centric Atari frames and shares PoE-World’s hierarchical planner for fair comparison.Source: §A.4 · p. 18
  • WorldCoder’s implementation uses the MIT license.Source: §A.4 · p. 18

A.5 Compute resources and execution time

The appendix reports compute budgets, execution times, synthesized constraints, prompts, example experts, and implementation code fragments used by PoE-World.

Compute Resources

  • PoE-World and WorldCoder use four CPUs, 64 GB memory, and $20 of OpenAI credit per world-modeling run.Source: §A.5 · p. 18

Execution time

  • PoE-World world modeling takes around 8 hours without planning, while most planner runs finish under 24 hours.Source: §A.5 · p. 18

Synthesized prompts

Examples and constraints

Constraint implementation

  • Constraint code checks object types and touch relationships, including player alignment with ropes, conveyor belts, platforms, ladders, and other objects.Source: §A.5 · p. 23

NeurIPS Paper Checklist

The checklist reports that the paper’s claims, reproducibility information, experimental details, limitations, ethics, licensing, and LLM usage disclosures meet the stated requirements.

Limitations and theory

Reproducibility and experiments

Open access

Ethics and assets

  • The paper reports no direct societal impact, no high-risk release safeguards are required, and existing assets are properly credited and licensed.Source: NeurIPS Paper Checklist · pp. 27–28

LLM usage

Related papers

Loading 2505.10819v4…