Compositional World Models

Updated 3 July 2026

Compositional world models are generative models that explicitly factorize complex dynamics into modular, reusable, and interpretable components.
They employ diverse architectures such as sparse Bayesian networks, product-of-experts, and slot-based models to achieve zero-shot and few-shot generalization.
Modular training and inference pipelines enhance sample efficiency, robust planning, and effective adaptation in robotics, multi-agent systems, and more.

A compositional world model is a class of generative models in which the underlying dynamics, state, or transition structure is explicitly factorized into reusable, interpretable, or independently trainable modules, each describing some aspect, factor, or substructure of the environment. Compositionality is leveraged to enable scalability, adaptive structure learning, zero-shot or few-shot generalization to novel combinations, modular planning, and improved interpretability and sample efficiency—contrasting with monolithic models in which all world dynamics are encoded as a single unstructured neural network or factorization. Compositional world models have emerged in diverse forms, including sparse Bayesian networks, product-of-experts over neural or programmatic modules, slot/block/factor architectures in visual models, language-conditioned causal decomposition, neurosymbolic attention mechanisms, and hierarchical automata, all of which enforce (soft or hard) modularity and reusability in world representation and simulation.

1. Formal Architectures and Mathematical Principles

A compositional world model is typically defined by explicit modularity in its latent variables, dynamics, and/or generative structure. Canonical instantiations include:

Sparse Dynamic Bayes Nets: The model family proposed by Da Costa is a sparse class of dynamic Bayesian networks (DBNs) over observations $x_{1:T}$ and latent states $z_{1:T}$ , with parameter $\theta$ , and structure $G$ :

$p(x_{1:T}, z_{1:T} | \theta) = \prod_{t=1}^{T} p(x_t | z_t, \theta_{obs}) \cdot p(z_t | z_{pa(t)}, \theta_{trans})$

where $\text{pa}(t)$ denotes a (sparse) set of temporal parents inherited from the graphical structure $G$ (Costa, 2024).

Mixture- or Product-of-Experts: Multiple independently designed/synthesized modules, each encoding domain knowledge or subdynamics, are combined via explicit product (e.g., PoE-World (2505.10819)) or via gating mechanisms (e.g., MoE in PRISM-WM (Li et al., 9 Dec 2025)) to form the joint transition distribution:

$p(o_{t+1} | o_{1:t}, a_{1:t}) \propto \prod_{i=1}^K \left( p^{expert}_i(o_{t+1} | o_{1:t}, a_{1:t}) \right)^{\alpha_i} \cdot \mathbf{1}[\text{constraints}]$

Slot/Block/Factor Models: Video-world models that decompose latent state into object-centric "slots" and factor-specific "blocks," as in Dreamweaver (Baek et al., 24 Jan 2025), yield state representations $S_t \in \mathbb{R}^{N \times (M \cdot D)}$ , with $N$ slots and $z_{1:T}$ 0 blocks per slot, each block specialized for an attribute (e.g., color, shape, dynamics).
Causal/Language-Conditioned Partitioning: In WM3C, the latent state is explicitly factorized into $z_{1:T}$ 1 language-controlled blocks $z_{1:T}$ 2, each parameterized and trained independently to guarantee block-wise identification and intervention invariance (Wang et al., 13 May 2025).
Neurosymbolic Attribute/Routing: Cosmos employs neural vector slots, each augmented with a symbolic attribute vector from CLIP, and a learned attention mechanism to bind objects to interaction rules. This hybrid neurosymbolic routing realizes compositional generalization to new compositions (Sehgal et al., 2023).
Hierarchical Automata: The World Automata framework defines components (agents, fields, targets) as composable automata whose interactions yield modular, hierarchical world models for hybrid systems (Capiluppi et al., 2013).

These formal architectures are all characterized by explicit decomposition of world representation and modular or conditional composition at training and inference time, supporting local updates, selective recombination, and targeted adaptation.

2. Learning, Inference, and Structure Discovery

Compositional world models require both parameter and structure learning to discover and exploit the appropriate modular decomposition:

Bayesian Structure Learning: Sparse Bayes net models are equipped with priors over structure $z_{1:T}$ 3 (e.g., Bernoulli or spike-and-slab on edges) and parameters $z_{1:T}$ 4, yielding a structure posterior $z_{1:T}$ 5 (Costa, 2024). Inference involves variational message passing for states, conjugate updating for parameters, and coordinate ascent or MCMC over $z_{1:T}$ 6.
Modular Training Pipelines: In RISE (Yang et al., 11 Feb 2026), the world model is split into (i) a controllable, action-conditioned video diffusion dynamics model; (ii) a separately trained value regressor; (iii) an RL policy; the models are pre-trained/fine-tuned separately before being integrated for closed-loop imagined rollouts and policy improvement via advantage conditioning.
Program Synthesis and Pruning: PoE-World iteratively alternates LLM-driven program generation (for novel experts) and empirical estimation of expert mixture weights by maximum likelihood with sparsity regularization, aggressively pruning weak or redundant components (2505.10819).
Meta-Learning for Modular Fusion: In WorMI (Yoo et al., 4 Sep 2025), meta-training (via Reptile) learns the fusion/composer module $z_{1:T}$ 7 such that it generalizes to arbitrary subsets of domain-specific world models at test time. Test-time adaptation is implemented by prototype-based retrieval of relevant domain modules and dynamic composition (via hierarchical attention) with a frozen reasoning LLM.
Relational Attention and Equivariance: Slot/graph/object-based models employ learned attention or slot-binding matrices to dynamically associate action/transition components to latent slots, as in HOWM (Zhao et al., 2022).

Learning procedures generally exploit the independence, modularity, or compositionality of structure to (a) constrain credit assignment, (b) enable fast adaptation to novel compositions, and (c) guarantee tractable inference or sample efficiency.

3. Compositionality and Generalization Properties

The core advantage of compositional world models is in systematic generalization and data-efficient transfer:

Module Reuse and Zero-Shot Combinatorics: By structurally encoding the world as a collection of modules (object slots, causal blocks, symbolic rules, or neural/compositional experts), the model can handle novel combinations of previously seen factors without retraining the monolithic dynamics. For example, RoboDreamer can synthesize video plans corresponding to unseen instruction combinations by recombining text- or image-conditioned primitives via product-of-experts sampling (Zhou et al., 2024); Dreamweaver swaps or recombines concept blocks to synthesize OOD images (Baek et al., 24 Jan 2025); Cosmos fires the correct learned rule on unseen object compositions via symbolic attention (Sehgal et al., 2023).
Test-Time Modularity: WorMI allows seamless, instantaneous insertion and removal of domain world models at inference, yielding strong zero-shot cross-domain transfer (Yoo et al., 4 Sep 2025).
Hierarchy and Variable Factorization: Formal results in object-oriented models (e.g., HOWM (Zhao et al., 2022)) show that equivariance properties (e.g., permutation symmetry) can be preserved in the latent transition dynamics, guaranteeing soft or exact generalization to new object sets or configurations modulo object identity permutations.
Causal Identifiability Guarantees: WM3C proves that under certain smoothness + intervention constraints, block-wise decomposition of latent state is uniquely identified up to invertible transforms, so long as each block experiences sufficient language-controlled variation (Wang et al., 13 May 2025).
Sample Efficiency and Planning Fidelity: Empirical studies report dramatic improvements over monolithic or heuristic baselines, often with 30–45% absolute gain on contact-rich robot tasks (RISE (Yang et al., 11 Feb 2026)), perfect or near-perfect rollout tolerances in benchmark gridworlds (Cosmos, HOWM), and more robust, long-horizon planning in hybrid dynamics (PRISM-WM (Li et al., 9 Dec 2025)).

4. Compositional World Models in Robotics and Planning

Robotic and embodied AI scenarios have catalyzed the deployment of compositional world models:

Planning with Behavior Primitives: In jumpy world models, planning is performed not at the level of primitive actions but sequences of entire (pre-trained) policy modules; the composition of multi-horizon outcome models enables zero-shot assembly of new tasks from existing skills with greatly magnified planning horizons (Farebrother et al., 23 Feb 2026).
Multi-Agent Coordination: COMBO's compositional world model factorizes the joint action of multiple agents into agent-specific components and employs a compositional video diffusion model that scales to arbitrary agent multiplicities via score-composition and product-of-experts logic (Zhang et al., 2024).
Physics and Digital Twin Integration: DreMa constructs a compositional scene representation from Gaussian Splatting and per-object object-centric assets, pairing photoreal rendering and physics-simulated dynamics for robust policy imagination and data-efficient imitation learning (Barcellona et al., 2024).
Systematic Task Transfer: Language-guided modularity (WM3C) and policy/planning transplantation (WorMI, World Programs) provide mechanisms for fast task adaptation and cross-domain task assembly via compositional fusion, addressing the limitations of monolithic RL and model-free approaches (Wang et al., 13 May 2025, Yoo et al., 4 Sep 2025, Segler, 2019).

These architectures yield robust adaptation, systematic generalization, and efficient planning in real-world, multi-task, and highly stochastic domains.

5. Theoretical, Empirical, and Formal Guarantees

Rigorous analysis of compositional world models spans identifiability, symmetry, modularity, and compositionality guarantees:

Block-wise Identifiability and Causal Decomposition: WM3C provides theoretical conditions ensuring uniqueness of latent block decomposition conditioned on language tokens, and bounds on mutual information terms enforcing selective dependence (Wang et al., 13 May 2025).
Symmetry/Equivariance: Models such as HOWM encode permutation equivariance at the architecture and loss level, guaranteeing soft compositional generalization across arbitrary object combinations (Zhao et al., 2022).
Trace Inclusion and Substitutivity: World Automata preserve trace-inclusion relations under parallel composition and hierarchical "inplacement," supporting modular verification and hierarchical modeling (Capiluppi et al., 2013).
Sample and Data Efficiency: In PoE-World (2505.10819), a handful (few hundred to a few thousand) of transitions suffice to fit effective Atari dynamics; in Dreamweaver (Baek et al., 24 Jan 2025), DCI disentanglement and OOD rollout scores substantially exceed state-of-the-art; in RISE (Yang et al., 11 Feb 2026), simulated rollouts fully in the compositional model enable on-policy RL without risky physical trials.

6. Limitations, Open Challenges, and Future Directions

Documented limitations and prospects include:

Scalability of Inference: Models with many experts (PoE, MoE) incur inference-time cost, especially in settings like WorMI, where multi-domain retrieval and two-tiered attention add computational overhead (Yoo et al., 4 Sep 2025).
Complexity of Structure Search: Full Bayesian structure learning is deferred in Da Costa, as empirical/algorithmic tractability at scale (number of nodes/edges) remains open (Costa, 2024).
Robust Extraction of Modular Factors: Slot extraction, symbolic attribute mapping, and correct action binding are failure points in pixel-based and light supervision domains (HOWM, Dreamweaver, Cosmos).
Incomplete Physical Modeling: Block/slot models may not capture intricate object-object physical interactions (collisions, contact dynamics) unless enriched by relational or physics-based modules (Baek et al., 24 Jan 2025, Costa, 2024).
Integration with Model-Free and Hierarchical RL: Balancing modular model-based planning with model-free RL (e.g., for reward-rich settings) remains a subject of ongoing work.

Emerging research targets more expressive and scalable fusion (e.g., hierarchical or graph-structured composers (Yoo et al., 4 Sep 2025)), on-the-fly module synthesis (e.g., LLM-driven expert generation in PoE-World (2505.10819)), tighter integration of VLMs for compositional policy/value inference (Zhang et al., 2024), and the extension to naturalistic, high-diversity settings (web-scale videos, real-world 3D scenes).

Compositional world models thus constitute a foundational paradigm for interpretable, generalizable, and robust model-based agents across domains as varied as RL, robotics, multi-agent systems, and program synthesis, with rapidly expanding theoretical and empirical support from diverse communities (Costa, 2024, 2505.10819, Baek et al., 24 Jan 2025, Yang et al., 11 Feb 2026, Yoo et al., 4 Sep 2025, Li et al., 9 Dec 2025, Capiluppi et al., 2013, Zhao et al., 2022, Sehgal et al., 2023, Barcellona et al., 2024, Zhou et al., 2024, Hayashi et al., 13 Mar 2025).