Compositionality and Systematic Generalization

Imagine you have trained a world model for a household robot on a stipulated dataset containing thousands of hours of kitchen activity. The model predicts, with impressive accuracy, how a mug slides across a countertop when pushed, and how a drawer swings open when pulled. Now you place the mug on a different counter and ask the robot to slide it while a drawer stands open nearby. Nothing about the physics is new under the example's autonomy assumptions: the sliding mechanism is the same, the drawer mechanism is the same, and neither depends on the other. A human predicts the outcome by putting the two known pieces together. A monolithic model trained on the observed distribution of scenes can fail, because the exact combination "this mug, that counter, drawer open" never appeared in its training data and no compositional constraint ties it to the observed scenes.

The failure here is subtle and worth dwelling on before we fix it. The stipulated dataset is large in duration, but duration alone does not establish coverage of the combinations that matter. What this particular model lacked was the right inductive structure. Every ingredient of the novel prediction was demonstrated in training, but the model was never forced to represent the ingredients separately, so its output on their new combination is not determined by those demonstrations. It learned answers, not a grammar for producing answers. Without another useful inductive bias, it may fail outright or produce a confident guess with no grounding in the missing combination.

This chapter is about designing and evaluating world models so that this kind of prediction is possible. The core idea is compositionality: represent the world as a set of reusable parts (objects, relations, mechanisms, skills) whose behaviors are learned separately and combined by known rules. A model can predict unobserved combinations after seeing a vanishing fraction of the configuration space only when the reusable components are adequately identified on the state-action support the new combinations require and the composition rule remains valid there. This conditional capacity is called systematic generalization, and it is an important property for world models intended to support planning in novel situations.

The stakes are especially high for world models specifically, as opposed to predictive models in general, because of how agents use them. A planner queries a world model at the states reached by candidate action sequences, and unconstrained candidate rollouts can leave the support of the data-collection policy. Generalization is therefore a core requirement whenever planning induces this distribution shift. Compositional structure is one important response when the assumed mechanisms and composition rule match the world; support constraints, uncertainty-aware objectives, robust planning, and new data address complementary failure modes.

We will build the idea in four steps: why compositionality matters for prediction and control, how factorized mechanisms formalize it, how to recombine objects, relations, and skills, and how to measure whether a model generalizes systematically. Along the way we will implement a small experiment where a factorized model and a monolithic model see identical data but behave completely differently out of distribution.

Recall from Part I that a world model earns its keep by supporting decisions, not just by fitting observed transitions. This distinction becomes sharp when we think about generalization. A model can have excellent one-step prediction accuracy on data drawn from the training distribution and still fail in closed-loop use when a policy reaches states it has not seen. Compositionality is one structural response to this problem when the relevant mechanisms can be reused, and it is worth understanding precisely what that response does and does not guarantee.

The mechanism of the failure is a feedback loop between the agent and the model. During training, the states in the dataset follow the trajectories of whatever policy or process collected them. When a planner searches over action sequences, it may favor states where the predicted outcome is good, including states near or beyond the regions the training policy visited. Without support constraints or adequate uncertainty controls, out-of-distribution behavior can therefore become a central operating regime and the planner can exploit optimistic model errors. Constrained or well-covered planning need not create this shift.

The natural world is combinatorial. A kitchen is not one monolithic scene; it is a particular mug, on a particular counter, with a particular drawer state, combined from a vocabulary of objects, surfaces, and mechanisms. If your vocabulary has $K$ factors with $q$ options each, the number of distinct worlds is $q^K$. Under any feasible data budget, you will observe only a tiny subset of these combinations. A model that treats each observed world as an independent learning problem has no principled prediction for an unseen combination. A model that represents the world as a composition of separately learned parts can determine that prediction from the behaviors of the parts, provided those parts were sufficiently observed and the composition rule is valid.

The exponential arithmetic deserves a moment of reflection, because it is the entire motivation for the chapter. With even a modest vocabulary, say twenty factors with five options each, the space of possible worlds contains $5^{20}$, roughly ninety-five trillion, distinct configurations. That finite space could be enumerated in principle, but exhaustive coverage is computationally and statistically impractical: even a passive sampler with full support would need an enormous coupon-collector budget to see every configuration. The practical route is to learn at the level of the factors themselves, so that each piece of data can update reusable components rather than only one point in the grid. This is an engineering response to exponential configuration growth, not a mathematical impossibility theorem about finite datasets.

The figure below makes the counting argument concrete: it plots the size of the configuration space against the fraction any realistic dataset can cover.

It is worth separating two things that are easy to conflate. Prediction quality asks whether the model's simulated rollouts match reality. Decision usefulness asks whether an agent planning against those rollouts chooses good actions. Valid compositional structure can improve prediction on novel combinations when the component estimates have relevant support and the composition rule is correct. It does not by itself make unsupported search trustworthy: calibration, uncertainty handling, and robust decision procedures remain separate requirements. We will keep both lenses in view throughout.

It is also worth pulling apart two further notions that are often collapsed into "the model is good." Representation quality asks whether the model's internal variables correspond to the real factors of variation in the world, so that the pieces it can recombine are the pieces the world recombines. Uncertainty quality asks whether the model knows what it does not know, so that its confident predictions in novel combinations are earned rather than reflexive. A model can predict well without having the right representations (a lookup table predicts perfectly on seen worlds), and it can have the right representations without calibrated uncertainty (a factorized model can be confidently wrong about a mechanism it saw in narrow conditions). Compositionality, taken seriously, addresses prediction and representation together, but it does not by itself deliver calibrated uncertainty; that requires the machinery from Stochasticity and Uncertainty. Keeping these four notions separate, predictive fidelity, representation quality, uncertainty quality, and decision usefulness, prevents the most common evaluation mistake, which is reading a low test loss as evidence of all four at once.

The formal backbone of compositionality in world models is the idea that the world's dynamics decompose into modules that can vary independently. This section develops that idea formally and explains what it buys you.

Building on Causal Models and Interventions, the central assumption is this: the process that generates the world is not one big function but a set of mechanisms, each responsible for a distinct aspect of the dynamics, and changing one mechanism does not change the others.

Before formalizing, it helps to see the intuition behind the assumption. In the kitchen example, changing the countertop material need not change the drawer hinges, and replacing a hinge need not change countertop friction. When such autonomy holds, an intervention can isolate one mechanism while the others remain stable. The formal principle below treats that locality as a modeling hypothesis rather than a universal law of nature, and it also tells us what would falsify the hypothesis: a change in one proposed mechanism that systematically alters another.

The phrase "replacing it with a different mechanism of the same kind" is doing real work. It says that mechanism families are shared across worlds: there is a family of possible friction behaviors, and worlds differ in which member of that family they instantiate. This is what makes cross-world comparison possible at all, and it is what the mechanism-indexed formulation below will make explicit. In Haavelmo's econometric program, autonomy is a modeling requirement to test over specified perturbations: a proposed structural relation should remain stable under the interventions for which it is claimed to be structural, rather than being assumed invariant under every possible change.

Concretely, suppose the state decomposes or factors into components, and the transition can be written as a product of conditional mechanisms:

where:

• $s_t$: the current state at time $t$, with $a_t$ the action taken at that time
• $s_{t+1}^{(k)}$: the $k$-th component of the next state, the slice of the world that mechanism $k$ is responsible for updating
• $p_k$: the conditional distribution that mechanism $k$ assigns to its own component
• $\mathrm{pa}(k)$: the (hopefully small) set of variables that mechanism $k$ depends on, drawn from the current state, the action, and possibly earlier components of $s_{t+1}$ in an acyclic same-slice graph
• $K$: the number of mechanisms (and corresponding state components) into which the dynamics are factorized
• $\prod_{k=1}^{K}$: the product over all $K$ mechanisms, which reassembles the full next-state distribution from the per-mechanism pieces

Read the product form carefully, because every term in it carries a commitment. As written, it defines a dynamic Bayesian network only when any dependencies among next-state components form a directed acyclic graph; the factors can then be arranged in a topological order, and the product is a normalized joint distribution with the corresponding local conditional independences. A next-state component may therefore depend on an earlier next-state component in that order. Genuinely simultaneous or bidirectional constraints, such as a coupled collision equilibrium, require different semantics: for example, structural equations with a well-defined equilibrium or normalized potentials with a partition function. Multiplying arbitrary cyclic conditionals is not enough to define a joint distribution.

Two properties make this useful:

• Sparsity: each mechanism depends on few variables, which reduces input dimension and can lower estimation burden; the local function may still be complex or poorly identified.
• Modularity: each mechanism can be reparameterized or intervened on without touching the others, which is exactly what lets you recombine them.

The two properties offer complementary benefits. Sparsity can reduce the dimensionality and estimation burden of each mechanism. Modularity lets the same estimated mechanism be reused across worlds. A sparse but world-specific mechanism still requires relearning per world; a dense mechanism can nevertheless remain compositional if it is autonomous, shared, and combined by a valid rule. Sparsity helps statistical and computational efficiency, but it is not logically required for recombination.

In a world with multiple environments, we can go further and index each mechanism by a setting. Write a world as a tuple $w = (m_1, m_2, \ldots, m_K)$, where each $m_k \in \mathcal{M}_k$ is a value of the $k$-th mechanism (friction is low or high, the object is a mug or a pan). The transition kernel is then:

where:

• $p_w$: the transition kernel belonging to world $w$
• $s'$: the next state
• $s, a$: the current state and the action taken
• $m_1, \ldots, m_K$: the mechanism settings, held fixed for the given world; conditioning on them selects one member of each mechanism family, so each world has its own transition kernel

The set of possible worlds is the combinatorial product $\mathcal{M}_1 \times \cdots \times \mathcal{M}_K$. This is the formal statement of the kitchen intuition: worlds are combinations of mechanisms, and the number of worlds grows exponentially in the number of factors.

Notice what this formulation does and does not assume. It assumes worlds differ only in mechanism settings, not in the structure of the factorization itself; the same $K$ families exist in every world. When a new world contains a genuinely new kind of mechanism, no element of $\mathcal{M}_1 \times \cdots \times \mathcal{M}_K$ describes it, and we land in the "novel structure" category of the taxonomy later in this chapter. The formulation also quietly assumes that the mechanism settings are stable within a world, which fits controlled transfer settings but not worlds whose mechanisms change over time; nonstationarity requires the machinery of later chapters.

The payoff is statistical. Compare two learning strategies in a world with $K$ binary mechanisms, giving $2^K$ possible worlds, where you collect $N$ trajectories per world.

A monolithic model treats each world as a separate learning problem. It must estimate $2^K$ distinct transition functions, seeing $N$ trajectories each, for a total of $N \cdot 2^K$ trajectories before it can predict everywhere. Worse, nothing in its hypothesis space tells it what to do with a world it has never seen. Even if the model is a flexible neural network conditioned on a world identifier, the identifier is an arbitrary label: knowing the transition for world $0$ and world $1$ constrains nothing about world $2$, because the label carries no information about which mechanisms differ. The monolithic model's knowledge is a set of isolated islands, with no bridges between them.

A factorized model can update a mechanism value from every trajectory in which it is active. In this full-factorial thought experiment, mechanism $k$ with a given value appears in $N \cdot 2^{K-1}$ trajectories (all worlds sharing that value), which is $2^{K-1}$ times the raw trajectory count available to one isolated world model. With $K = 10$ binary mechanisms, that is 512 times as many pooled trajectories containing the value. This is not a 512-fold sample-efficiency theorem: the full factorial collection itself contains $N \cdot 2^K$ trajectories, and under a fixed balanced budget $B$, a binary value appears in roughly $B/2$ trajectories regardless of $K$. The structural saving is instead that the model stores a library of $K \cdot q$ value-specific components plus a composition rule rather than a separate component for every world.

It is worth seeing where the pooling number comes from, because the same counting discipline governs experimental design. Fix a value, say $m_k = 0$, and count the worlds that contain it: exactly half of all $2^K$ worlds, that is $2^{K-1}$ of them. Each of those worlds contributes $N$ trajectories, so the value occurs in $N \cdot 2^{K-1}$ trajectories. Whether those trajectories provide proportionally more independent information depends on excitation, correlation, and model specification. Under a fixed total budget, the useful advantage is that each trajectory informs several reusable components and that the design can balance coverage of mechanism values without enumerating every world.

The amplification factor is easiest to absorb as a curve: for each number of binary mechanisms, compare the trajectories available per mechanism value under the two strategies.

This construction echoes a recurring idea in system identification: choose or learn a dynamical model structure, then estimate its behavior from input-output data. Textbook examples often fit a few interpretable parameters such as stiffness or damping, while the broader field also includes structure selection and flexible nonlinear models. Factorized world models add an explicit reusable decomposition: component functions may be learned flexibly, but the architecture still encodes which parts are separable and how they combine.

The spring example captures the useful analogy. A spring constant identified on one test rig is intended to describe the spring rather than the rig, so it can transfer when the relevant operating assumptions continue to hold. A learned factorized component aims for the same scene-independent semantics, but that semantics is a hypothesis to validate rather than a consequence of its name. Modern methods may learn both the component function and parts of its parameterization from data.

There is an important caveat that separates textbook treatments from practice. Everything above assumed the model receives the mechanism identity $(m_1, \ldots, m_K)$ as an input. In real problems, the factors are not labeled. The agent observes pixels and forces, and must infer which latent mechanisms are present and how they combine.

This is the domain of causal representation learning, and it is hard. As discussed in Latent Variable Representations and Object-Centric and Relational Worlds, learning a factorized representation from passive observations alone is impossible in general without additional inductive biases, as formalized by Locatello et al. (2019): many different latent factorizations can produce identical observations. Interventions are one powerful way to break this symmetry by changing one thing while holding others fixed. Temporal or auxiliary structure, grouped observations, supervision, and architectural assumptions can also provide identifying information under their own conditions.

The impossibility claim is worth stating in intuitive terms so that it does not read as pessimism about a solvable problem. A passive dataset is a set of samples from one joint distribution, and without additional assumptions that distribution can support many latent descriptions: an invertible transformation may leave observed data unchanged while scrambling which latent controls what. If friction and wetness always vary together, passive samples alone do not distinguish them. An intervention that varies friction while holding wetness fixed can cut that tie, while other auxiliary information can do so in different problem settings. The hierarchy below grades problems by how much identifying structure is available.

The practical hierarchy to keep in mind:

• Oracle factors: mechanism identities are given (simulators, labeled domains, controlled experiments). This removes the factor-discovery problem, but architecture, state-action support, excitation, finite-sample estimation, optimization, misspecification, and validation remain real concerns.
• Inferred factors from known variables: the factors exist in the observation space (object positions, temperatures) and the model must learn to route them. Methods from Object-Centric and Relational Worlds live here. The hard part shifts from discovering the factors to deciding which observed variables each mechanism reads and writes, and mistakes in routing show up as silent composition failures.
• Fully latent factors: the model must discover both the factors and the mechanisms from raw interaction. This remains an active research problem. Weak architectural biases such as permutation invariance or locality and deliberately structured data collection are two sources of identifying information, but neither combination is a general recipe and each method must state its assumptions.

The hierarchy is useful because it separates statistical, architectural, and representational difficulties. Moving down it generally requires more identifying assumptions or auxiliary information; interventions are one important source, but not the only one. A team that knows where it stands can budget data collection and validation accordingly.

The same factorization logic applies one level up. In the options framework of Sutton, Precup, and Singh (1999), an option is a temporally extended closed-loop policy defined by an initiation set, an internal policy, and a termination condition. An option model or skill-effect model separately summarizes the option's state-transition and reward consequences for planning. If you learn such a model for each skill, for example how "push object" changes position or how "open container" changes container state, then planning can compose those abstract consequences. We will develop the hierarchical machinery for this in Part V's chapter on hierarchical and hybrid models and in Part VII on planning.

The parallel is useful but not an identity. A low-level mechanism maps state and action variables to next-state variables; an option model maps a starting state through a temporally extended policy to a distribution over termination states, durations, and rewards. The option's initiation set determines where it may start, while its termination rule determines when it hands control back. Both settings replace "learn what happens in each full situation" with reusable pieces plus a combination rule, but the temporal abstraction adds policy and termination semantics that a one-step mechanism does not have.

Factorized mechanisms tell us that the world decomposes. Object-centric and relational representations, which we developed in Part III, tell us what the pieces usually are. This section connects the two.

Physical environments are naturally described as a set of objects, each with its own state, its own intrinsic dynamics, and interactions governed by relations. When a world model represents state as a set of object states rather than a flat vector, it inherits a powerful form of combinatorial structure:

• Permutation sharing: all objects of the same type share the same dynamics. A model that learns "how mugs slide" once can apply it to every mug, which is data amplification exactly like the factorized case. Every mug in every training scene becomes a training example for the shared mug-dynamics module, so the effective dataset for that module is the total number of mug-steps across all scenes, not the number of scenes.
• Permutation invariance: predictions should not depend on the arbitrary order of objects in the input. Architectures that process sets or graphs enforce this by construction. Without architectural invariance or adequate order augmentation, a model trained with mugs always listed before pans can fail when the order reverses; a non-invariant architecture may still learn the symmetry from sufficiently varied data.
• Variable cardinality: a model that computes per-object and pairwise quantities can, in principle, run on a scene with more objects than it ever saw during training, because nothing in its computation depends on the count. The per-object encoder applies to each object regardless of how many there are, and the aggregation combines any number of pairwise messages into a fixed-size output.

These three properties make the intended sharing explicit in object-centric models. A flat-vector model can also implement equivariance, parameter sharing, locality, or algebraic factorization, but a flat model without equivalent structure has no architectural guarantee of those symmetries. An object-centric model predicts object-wise futures and relation changes, then assembles the next state subject to local couplings such as collisions rather than treating object futures as independent.

Relations are the second half of the vocabulary. Contact, support, containment, and collision apply to pairs of objects. Implementing a relation as a shared function, as message-passing architectures do, creates a route for transfer across pairs. A learned "supported-by" relation can transfer from mugs and shelves to books and tables only when its inputs capture the relevant physics and the new pair lies within the relation model's supported regime.

The caveat in that last clause deserves emphasis. A message-passing network can implement "shared relation" at different scopes. A relation function based on physical pair quantities such as relative position, velocity, and contact geometry can potentially transfer broadly. A module conditioned on type embeddings may still share parameters, but only within the transfer scope supported by those types. If it memorizes individual object pairs, it cannot reuse the relation on a new pair. Whether the learned sharing matches the intended mechanism is an empirical question, and the swap test later in this chapter is designed to probe it.

This is where the relational inductive bias becomes concrete. A model with per-object encoders, pairwise relation functions, and permutation-invariant aggregation encodes sharing directly, while a fixed-length flat network without equivalent sharing leaves recombination unconstrained. The component inventory can scale with objects and relation types rather than scene combinations, but empirical coverage still depends on support and estimation quality. We will see graph simulators and slot-based models in Part V; the point here is that architecture can make a composition rule explicit.

A distinct generalization axis deserves its own name. Productivity is the ability to handle structures longer, deeper, or larger than those seen in training: a chain of ten dominoes after training on three, a sentence with a nested clause after training on short ones, a robot assembly with twelve parts after training on five. Size-agnostic shared computation can make such inputs architecturally executable without adding a new module for every length or object count. That does not guarantee accurate extrapolation, and a fixed-size or otherwise monolithic model can also be designed for variable length. Productivity is where the limits of learned mechanisms show up: a contact model that is approximately right per interaction can be badly wrong about how errors accumulate over a long chain.

The domino example makes both sides of this vivid. Applying the same local "one domino topples the next" mechanism ten times is trivial for a compositional model; it is literally the same computation repeated, and no architectural change is needed to go from three dominoes to ten. But ten applications also means ten chances for the local mechanism's error to compound, and if each toppling slightly overestimates the transferred momentum, the tenth domino's behavior may be badly wrong even though every two-domino prediction was excellent. Productivity therefore splits into two questions that must be evaluated separately: can the model run on larger structures (an architectural question, where compositional models win by construction), and are its predictions still accurate on larger structures (a statistical question about error accumulation, where no architecture saves you automatically). The evaluation methodology later in the chapter, particularly horizon-resolved rollout error, is exactly the instrument needed for the second question.

We now have the vocabulary to be precise about what "out of distribution" means for world models, because not all distribution shift is the same kind of problem. Lumping them together leads to bad diagnosis and bad remediation: a team that treats a cheap recombination failure as if it needed new primitive learning will waste data, and a team that expects recombination to solve a novel structure will ship a silently broken model. The taxonomy below is the diagnostic instrument.

Consider a model trained on a set of worlds drawn from $\mathcal{M}_1 \times \cdots \times \mathcal{M}_K$. The following situations require different responses:

• Novel recombination: each mechanism value was seen during training, but a particular combination was not (the mug on the new counter). No world-specific learning is needed once the relevant component functions are adequately identified on the state-action support required by the new pairing and the composition rule is valid. If those conditions hold, failure suggests a factorization or composition problem; without them, insufficient support, excitation, estimation, or optimization remain alternative explanations.
• Novel instance within a known class: a new object that is recognizably a mug but with a different handle. Here generalization must come from the representation abstracting away instance details, which is the concern of Part III's chapters on sufficient state and object-centric representations. The relevant question is whether the representation has discarded the details that do not matter (decorative variations) while retaining the ones that do (mass distribution, handle geometry that affects grasping), a judgment that only a sufficient-state analysis can make.
• Novel mechanism parameters: a continuous mechanism (friction, mass, elasticity) takes a value outside the training range. A physically meaningful parameter supplies an ordered extrapolation axis that an opaque categorical embedding lacks, but predictive accuracy still depends on the functional form, regularity, and training support. Conditioning on a friction coefficient of $0.9$ gives the model a notion of direction beyond a training range of $0.1$ through $0.5$; it does not guarantee a correct prediction there.
• Novel structure: a new relation or a new way for mechanisms to interact that cannot be expressed with the current vocabulary and composition rule. A new object class is structurally novel only if its relevant properties or mechanisms cannot be represented by existing primitives. Recombination alone cannot cover a genuine vocabulary gap; the model needs richer primitives or structural revision. Uncertainty estimates and mechanism-level diagnostics can help flag such cases, but no generic method guarantees detection of every novel structure, so deployment needs explicit abstention and data-collection policies.

The practical lesson is that "OOD" is not one problem. Diagnosing which kind of shift you face tells you which test or remedy is appropriate. A useful habit is to write the shift in mechanism notation before reacting to it: if the new situation is a point in $\mathcal{M}_1 \times \cdots \times \mathcal{M}_K$ with familiar coordinates, recombination is plausible only after checking support, identification, and the composition rule; if it is outside that space, existing primitives alone do not determine the answer.

The warning example is worth unpacking, because interactions between mechanisms are common in the physical world and easy to miss in a factorial design. If the model factors dynamics into "object mechanism" and "surface mechanism," and the surface mechanism was learned only with light objects, then composing the heavy-object mechanism with the polished-surface mechanism may predict a confident glide where reality delivers stiction or deformation. Nothing in the training objective flagged the interaction, because the training combinations never separated heavy objects from rough surfaces. This is why the factorial design of training data matters as much as the factorial design of the model: combinations in the training data are the only place where mechanism interactions could have been observed.

It is worth being precise about why an unconstrained function approximator may fail at novel combinations. Empirical loss directly constrains a flexible network only on sampled support. Without a valid inductive bias or additional assumptions, multiple functions can agree on every training transition and disagree on a missing cell, so off-support behavior is not identified. Some architectures and target functions extrapolate successfully; the training loss alone does not guarantee which extension is learned.

The missing cell has no direct empirical-loss term: no observed transition evaluates the prediction at that world. Gradients from observed cells still update shared weights and thereby change its output indirectly, according to the architecture, regularization, initialization, and optimization path. Those inductive couplings may be useful, but the data do not identify which off-cell extension is correct. A missing cell can even sit between seen inputs, where neighboring constraints may disagree.

A compositional architecture changes this by constraining the hypothesis space. If the model must compute its prediction as $f_1(x; m_1) + f_2(x, a; m_2)$, then the value at the missing combination $(m_1, m_2)$ is determined by the learned $f_1(\cdot; m_1)$ and $f_2(\cdot; m_2)$, each of which was trained on abundant data from other combinations. The generalization is not luck; it is a consequence of restricting the model to functions that factor. This is inductive bias in the strict sense: you give up the ability to represent non-factorizing worlds in exchange for automatic coverage of all combinations.

The grid picture makes the implemented contrast visible: the categorical world-ID baseline has no sharing rule that constrains a missing cell, whereas the additive compositional model assembles that cell from its learned row and column factors. Other flat or monolithic architectures may encode equivalent sharing through different inductive biases.

When the factorization matches the world, the coverage trade can be favorable. When it is materially wrong and the model lacks corrective interaction terms or calibrated uncertainty, it can be confidently wrong in a structured way.

The additive example also shows how composition rules trade expressiveness for coverage, and why the choice of rule is a real design decision. An additive rule composes contributions that do not interact; a multiplicative or gating rule allows interaction but needs data to estimate the interaction terms. The true dynamics of the worked example below are additive, so the additive bias matches reality exactly; had the true dynamics contained a term like $m_1 m_2 x$, the additive model would be systematically biased at every combination where the interaction term is large, including seen ones. Matching the composition rule to the world's actual combination structure is as important as factorizing at all.

Two themes from earlier parts become sharper here. First, compositional correctness is about one-step mechanisms, but agents act over horizons. Even small per-step errors in a novel combination compound over an autoregressive rollout, so a model that is "right" about the composition can still produce useless long-horizon predictions. Evaluation must therefore measure fixed-action autoregressive error as a function of horizon and, separately, closed-loop decision performance when a policy or planner selects actions, a distinction that connects to the uncertainty machinery of Stochasticity and Uncertainty.

Second, in model-based reinforcement learning (Part II's chapter on model-based RL), the planner searches for action sequences where predicted outcomes are good. Out-of-distribution and novel combinations can contain large optimistic errors that such a planner exploits, especially without support constraints or uncertainty penalties. They need not be the worst-error regions or coincide with every composition test. Compositional structure can reduce some unconstrained variation, but it does not eliminate unsupported component outputs or planning exploitation.

The two error-growth regimes are worth seeing side by side, because they call for different fixes.

Factorization ties a novel world's prediction to its component outputs, but those outputs can still be arbitrary outside their learned state-action support. Their composition may therefore assign an optimistically low cost to an unsupported combination. Distribution shift and model exploitation are treated as first-class reliability concerns in Part XII.

If we care about compositional generalization, we must evaluate for it. In factorial datasets with multiple observations per world, an observation-level i.i.d. split will usually put the same world combinations in both training and test, so it does not isolate novel composition. The chance of overlap depends on the number of observations per cell and the split fraction, rather than following from i.i.d. sampling alone. A model can score well on the repeated combinations while remaining unconstrained on combinations withheld at the world level. Everything in this section exists to close that gap.

The canonical recipe has three requirements.

• Primitive coverage: every primitive (mechanism value, object type, relation) should appear in multiple training combinations. Marginal appearance is necessary but not sufficient: verify state-action support, excitation, sample adequacy, and component estimation quality in the conditions relevant to the held-out pairing. Only after those controls does failure specifically implicate factorization or composition rather than primitive learning.
• Held-out combinations: test worlds are combinations absent from training. A clean way to guarantee this is to hold out a structured subset, such as an entire anti-diagonal of a factorial design, so that every factor value remains well represented in training. The anti-diagonal has a pleasing property: removing it removes every combination of the form $(m, c-m)$ while leaving every row and every column of the grid populated, so both mechanisms keep full coverage of their value sets.

The following figure shows the split used in the worked example: removing the anti-diagonal leaves every row and column populated.

• Multiple seeds and splits: with few held-out cells, results are noisy. Repeat the split with different held-out cells and report the spread. A model that passes the anti-diagonal split might fail when a different diagonal is held out, and that sensitivity is itself a finding: it suggests the model's composition is only partially factorized, with some mechanisms entangled in ways other splits do not expose.

In the compositional generalization literature for language and reasoning, splits are often scored with divergence measures that quantify how different train and test compositions are. The same discipline applies here: design the split so memorizing observed combinations is insufficient. For each held-out world, check not only whether every mechanism value appears in training, but whether the data cover and excite the component behavior at relevant states and actions. These controls make composition the intended novelty; they cannot prove it is the only source of difficulty.

A single test loss hides the distinctions that matter. At minimum, report:

• One-step error on held-out combinations: this measures whether the composed predictor works on sampled true states without rollout compounding. Good one-step error with poor autoregressive performance points toward compounding or rollout-distribution shift, but additional state-coverage checks are needed before declaring the composition globally correct.
• Fixed-action autoregressive rollout error as a function of horizon: replay the same exogenous actions while feeding model predictions back as states. This isolates multi-step predictive compounding; it is not closed-loop control. The curve's shape also matters: roughly linear growth suggests persistent bias, while rapid growth suggests instability, subject to checking scale, dynamics, and noise.
• Planning success with the model held fixed: train the model on the compositional split, then let a planner use it in held-out worlds. This is the decision-usefulness number, and it can decouple from prediction accuracy when the model's errors are systematically biased in directions the planner exploits. Zero-mean errors under a fixed sampling distribution do not by themselves make a model serviceable for planning: a maximizing planner can select optimistic tails, especially when errors are heteroscedastic or correlated across candidates. Lower average error can likewise coexist with a harmful directional bias, so decision usefulness must be measured directly.
• Sample efficiency for adaptation: given a small budget of trajectories in a new world, how quickly does fine-tuning recover accuracy? Selective updates and fast repair are evidence consistent with localized mechanisms, but generic feature reuse or meta-learning can produce the same pattern. Track which parameters change and control support, optimizer, and adaptation budget before attributing the result to factorization.

No single number among these substitutes for the others, and reporting all four is cheap relative to the cost of deploying a model whose evaluation summarized away the failure mode that matters.

Several established benchmarks probe related abilities, but their split semantics differ. SCAN tests compositional command generalization through structured command splits. CLEVR-CoGenT tests visual question answering under held-out combinations of object colors and shapes. COGS uses structured generalization splits for semantic parsing. Procgen instead evaluates transfer to held-out procedurally generated levels, while manipulation suites can be designed around held-out object combinations. The shared lesson is to define the intended shift explicitly; only some of these benchmarks implement the exact fixed-primitives, held-out-combinations design used in this chapter.

Many controlled domains admit an approximate factorial construction once candidate mechanism families are explicit. Where interventions are feasible, enumerate or sample combinations, verify primitive support in training, hold out a structured subset, and run the metric suite above. Continuous, coupled, irreversible, or ethically nonmanipulable mechanisms may require partial factorial designs or different stress tests. A domain-specific split is useful when its assumptions match the system.

A particularly informative diagnostic is the swap test: take environments that differ in a candidate mechanism, swap that mechanism, and check whether predictions change accordingly. Passing is evidence consistent with the proposed factorization, while the categorical-ID baseline below has no sharing rule that requires it to pass. Coincidence, support differences, and unequal task difficulty remain alternative explanations, so use matched controls and repeated swaps.

The swap test deserves a precise statement. Let worlds $w = (m_1, m_2)$ and $w' = (m_1', m_2')$ share no mechanism values. Check both cross-combinations $(m_1, m_2')$ and $(m_1', m_2)$. An asymmetric failure can suggest which candidate mechanism or routing path is problematic, but support, noise, and unequal difficulty must be ruled out before causal localization. Simulation permits exact swaps; real systems may permit controlled physical recombination.

Let's make all of this concrete with the smallest possible factorial world. The state is a single scalar $x$. The world is described by two independent mechanisms:

• Mechanism A (intrinsic dynamics), with three variants indexed by $m \in \{0, 1, 2\}$, each with a decay $\alpha_m$ and a curvature $\kappa_m$.
• Mechanism B (action response), with three variants indexed by $n \in \{0, 1, 2\}$, each with an action gain $\beta_n$ and a state-action coupling $\gamma_n$.

The transition is:

where:

• $x_t$: the scalar state at time $t$
• $a_t$: the action taken at time $t$
• $\alpha_m, \kappa_m$: the decay and curvature parameters of mechanism A variant $m$
• $\tanh(x_t)$: the hyperbolic tangent of the state, which bounds the curvature contribution to magnitude at most $|\kappa_m|$; the full intrinsic term can still grow linearly through $\alpha_m x_t$
• $\beta_n, \gamma_n$: the action gain and the state-action coupling of mechanism B variant $n$
• $\varepsilon_t$: small process noise added to the state at each transition, with standard deviation $0.02$ in the implementation below

The structural form of this equation is the entire point of the exercise. The right-hand side is a literal sum of a mechanism-A term and a mechanism-B term, with no interaction between them, so the true dynamics satisfy the additive factorization exactly. This is the ideal case for a factorized model: its inductive bias is not an approximation of the truth but a restatement of it. The two mechanisms are also semantically distinct, one governing how the state evolves on its own and one governing how actions perturb it, so the factorization matches a natural physical reading of the world rather than a mere algebraic convenience.

There are $3 \times 3 = 9$ worlds. We will train on six of them and hold out the anti-diagonal: $(m{=}2, n{=}0)$, $(m{=}1, n{=}1)$, and $(m{=}0, n{=}2)$.

To see the composition rule in action, fix $x_t = 2$ and $a_t = -1$, and ignore noise. Mechanism A with $m=1$ contributes $0.85 \times 2 - 0.3 \times \tanh(2) = 1.70 - 0.29 \approx 1.41$, using $\tanh(2) \approx 0.96$. Mechanism B with $n=0$ contributes $0.8 \times (-1) = -0.8$, giving $x_{t+1} \approx 1.41 - 0.80 = 0.61$. Mechanism B with $n=1$ instead contributes $1.6 \times (-1) + 0.2 \times (-1) \times 2 = -1.6 - 0.4 = -2.0$, giving $x_{t+1} \approx 1.41 - 2.0 = -0.59$. The two worlds share the intrinsic term exactly and differ only in the action response. If a model has adequately identified both contributions on the state-action support used by the held-out rollout, the valid additive rule determines their combination; learning a separate categorical world embedding supplies no corresponding constraint on $(1,1)$.

The arithmetic makes the data-sharing argument tangible. Every training world with $m=1$ in our split, namely $(1,0)$ and $(1,2)$, provides observations of the same intrinsic term $0.85 x - 0.3\tanh(x)$, and every world with $n=1$, namely $(0,1)$ and $(2,1)$, provides observations of the same action-response term. The held-out world $(1,1)$ asks the model to add two familiar contributions. In the true simulator the pairing introduces no interaction term, but finite data, limited excitation, optimization, and off-support component error can still leave the learned contributions inaccurate.

The composition rule can be drawn directly: each mechanism contributes a curve over states, and the held-out prediction is simply their sum.

We will now build both models in PyTorch, train them on identical data from the six training worlds, and compare fixed-action autoregressive rollouts on all nine. The monolithic model receives the world's identity as a single index and uses one shared prediction head. The factorized model receives the two mechanism indices separately and computes its prediction as the sum of two mechanism-specific heads. The implemented models do not have matched parameter counts: the monolithic model has 5,001 trainable parameters and the factorized model has 9,906. The experiment therefore illustrates the two constructions under matched data and optimization, with model capacity left as a confound; it is not a capacity-controlled causal estimate of structure alone.

Before reading the code, it helps to know exactly what the comparison is and is not. It is a single-seed, single-split pedagogical illustration under matched data and optimization but unmatched parameter count. It does not satisfy the minimum serious-evaluation protocol stated above: it has no repeated seeds or splits, no planner that selects actions, and no adaptation study. It shows how this particular untrained world-index representation behaves on cells the data never touched; capacity-matched controls and the full evaluation suite would be required to isolate structure's causal contribution or make a deployment claim.

First, import the remaining libraries and define the ground-truth world. The mechanism tables store the parameters for each variant, and a single rollout function composes them exactly as the equation above specifies.

The key parameters for the factorized and monolithic comparison are:

• d: Embedding dimension for the mechanism and world embeddings. Both models use 8, but equal embedding width does not equalize total capacity.
• hidden: Width of the hidden layers in each MLP. The monolithic model has one width-64 MLP, while the factorized model has two. Including embeddings, their trainable parameter counts are 5,001 and 9,906 respectively.
• lr: Adam learning rate, set to 0.001 for both models so neither receives a more favorable optimization setting.
• steps: Number of minibatch gradient steps (4000), the full training budget each model receives on the pooled training transitions.
• batch: Minibatch size (512) of transitions drawn at each step from the pooled training data.

The data, optimizer, learning rate, minibatch size, and training steps are matched, but the parameter budget is not. Consequently, this run supports a mechanism demonstration (unseen world embeddings remain untrained while shared mechanism embeddings receive updates), not the causal claim that structure alone produced the full performance difference. A capacity-controlled experiment would add a matched-width factorized model or a larger monolithic control and rerun the comparison under the same seeds.

The honest reading of this chapter is that compositional generalization is purchased with assumptions, and each assumption can fail. The ICM principle says mechanisms are separable, but real mechanisms interact: a gripper's contact dynamics depend on both the object's surface and the gripper's material in ways that are not cleanly additive, and in fluids or soft bodies the very notion of a separable mechanism can be approximate at best. The factor labels in our experiment were handed to the model. From raw observations alone, those labels are underdetermined in general; identifying them requires additional assumptions or information, which may come from interventions, temporal or auxiliary structure, grouped observations, supervision, or architectural constraints. A model can also factorize along the wrong axes, such as visual appearance rather than physical mechanism. These failures need not trigger high uncertainty, which is why swap tests, mechanism-level diagnostics, and uncertainty quantification remain necessary even for structured models.

The wrong-axis failure deserves emphasis. A model that factors by appearance can pass splits dominated by visual variation yet fail physical swaps, such as visually similar surfaces with different friction. The structural factorization alone need not expose that confusion, although a calibrated uncertainty method may. Design diagnostics around physical interventions rather than visual variation, and test for invariance to appearance together with sensitivity to dynamics.

There is also a standing tension with scale. Strong structural priors can improve sample efficiency and compositional coverage when their commitments are correct, while diverse data can expose a less structured model to many useful recombinations. These approaches are complements rather than guaranteed substitutes: architecture constrains how components can combine, and data determines which behavior has empirical support. The data-side implications, including interaction, curation, and synthetic data, are taken up in Part VI.

An implicit factorization provides no architectural certificate for which recombinations are grounded, yet it remains empirically diagnosable through held-out-combination tests, swaps, probes, and interventions. Explicit structure makes its commitments easier to inspect and test, at the cost of committing to a factorization that may be wrong. In safety-relevant systems, that inspectability can be useful only when paired with evidence that the structure is accurate, calibrated uncertainty, and robust decision procedures; inspectability alone is not a safety guarantee.

The compositional viewpoint has produced useful tools in concrete systems, but the scope of each result matters. Interaction Networks, for example, reported transfer to some synthetic physical systems with object counts beyond training; that result does not establish arbitrary arrangement or real-world transfer. A correctly specified mechanism family can pool informative interactions across related systems when the data excite the relevant parameters. Option models can let a planner reason over temporally extended actions rather than only primitive actions, but whether this improves planning depends on model accuracy, initiation and termination semantics, and search. Primitive-covered splits, horizon-resolved rollout error, and planning-success measures test those claims in the situations where an agent will act.

• Compositionality is the design commitment that a world model should represent the world as a combination of reusable parts (objects, relations, mechanisms, skills) rather than as a monolithic function from histories to predictions.
• Factorized mechanisms formalize this: the transition decomposes into independently varying modules, and worlds are points in the product of mechanism spaces. For $K$ families with $q$ values each, the representation stores $Kq$ value-specific components rather than a separate component for every one of the $q^K$ worlds. Balanced data let each trajectory update several reusable components, but the resulting sample-efficiency gain depends on the fixed data budget, coverage design, excitation, and dependence among trajectories.
• Recombination of objects, relations, and skills can give zero-shot coverage of novel pairings when the architecture enforces a valid composition rule and the components are identified on the relevant state-action support.
• Out-of-distribution composition comes in distinct kinds (novel recombination, novel instances, novel parameters, novel structure) with distinct remedies; conflating them hides both failures and cheap wins.
• Systematic generalization must be measured, with primitive-covered splits, held-out combinations, horizon-resolved rollout error, swap-consistency tests, and planning success rather than i.i.d. one-step loss alone.
• The factorization is a testable hypothesis. A wrong factorization can fail without warning, so pair structured models with uncertainty quantification and mechanism-level diagnostics.

One closing habit to carry forward from this chapter: whenever you evaluate or design a world model, ask what constrains its prediction for a novel combination. Learned parts and a valid composition rule provide one answer; other models may encode equivalent sharing or equivariance without explicit object modules. If no valid inductive bias or data support constrains the prediction, it is an unsupported off-data extension that planning may exploit—not necessarily interpolation, and not an inevitable property of every monolithic model.

The next part of the book turns from what world models should represent to how concrete architectures realize it. We begin with the recurrent state-space model, a recurrent latent-dynamics architecture that couples deterministic and stochastic state evolution. Keep this chapter's question in mind as we get there: nothing in the RSSM objective forces its latent transition to factorize, so whether these models compose is a property of their design and data, and a property you now know how to test.

Ready to test your understanding? Take this quick quiz to reinforce what you've learned about compositionality and systematic generalization in world models.