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

TL;DR

PoE-World presents a modular, symbolic world model built as an exponentially-weighted product of many small programmatic experts synthesized by LLMs. This compositional approach improves data efficiency and scalability, enabling accurate modeling of stochastic and partially observable dynamics and supporting planning and RL through a hierarchical planner. Empirical results on Pong and Montezuma's Revenge show strong performance in low-data regimes, with the method generalizing to novel levels and enabling planning-based decision making and policy pretraining within the world model. The work emphasizes interpretability and controllability through explicit programs and hard constraints, advancing scalable symbolic reasoning for complex environments.

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

Figures (7)

• Figure 1: (a) World models predict the next state given a state-action history. We do this with a product of experts of many small programs. (b) The learner is given a short (<1 min) demonstration of gameplay as input, and uses it to 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.
• Figure 2: A "heat map" illustration of how simple Python programs are interpreted as distributions, and how they are combined into a single distribution over next-timestep object locations.
• Figure 3: A sequence of illustrations that demonstrates how our hierarchical planner works.
• 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 Revenge-Alt requires players to avoid moving skulls by jumping over them and climbing up ladders to reach and collect the key.
• Figure 5: Maximum score differences over random baseline achieved at different number of training steps (in log scale). The text "(demonstration)" annotates the length of demonstration used to initialize WorldCoder and PoE-World, which is not displayed for the Alt games since we use the world models initialized by the demonstrations of the corresponding base games.