Neurosymbolic Grounding for Compositional World Models

• Each entity combines a neural real-vector representation with symbolic attributes such as shape, color, and orientation.Source: §1 · p. 1
• Neurosymbolic attention matches symbolic keys to entity attributes while using neural queries to bind learned interaction rules to entities.Source: §1 · pp. 1–2
• Compositional symbolic attributes let the attention mechanism apply rules to new compositions of scene parts, supporting both entity and relational compositionality.Source: §1 · p. 2
• COSMOS uses CLIP to assign symbolic attribute values, avoiding the manually specified perceptual-to-symbolic mapping required by traditional neurosymbolic methods.Source: §1 · p. 2
• COSMOS is a differentiable neurosymbolic approach to object-centric world modeling that combines neurosymbolic scene encodings with neurosymbolic attention.Source: §1 · p. 2
• The evaluation reports significant gains over the state of the art for compositional world modeling on an established 2D moving-shapes domain.Source: §1 · p. 2

• World modeling learns an action-conditioned transition function over object-factored states rather than directly modeling the intractable pixel state space.Source: §2 · p. 3

Symbolic object relations

• Objects are assigned symbolic attributes through projections over a predefined ordered attribute list, allowing attributes to be represented separately from each object.Source: §2 · p. 3

• CompGen requires a distribution shift in compounds while preserving the distribution of atoms between training and test data.Source: §2 · p. 3

Entity Composition

• Entity composition holds atom frequencies fixed but makes training and test compound distributions disjoint, so familiar objects appear in unseen co-occurrences.Source: §2 · p. 3

Relational Composition

• Relational composition introduces shared dynamics between objects with common attributes, such as color or adjacency, while keeping compounds disjoint across splits.Source: §2 · p. 3

Object-centric encoding

• COSMOS extracts one slot-based representation per entity and reconstructs the image from the set of slots.The entity extractor uses Segment Anything with pretrained weights to produce set-structured object representations, while a permutation-equivariant spatial decoder reconstructs the image from masked slot reconstructions.Source: §3.1 · pp. 4–5

Neurosymbolic encoding

• Each entity combines a neural vector with composable symbol vectors for attributes such as shape, color, orientation, and position.The symbolic labelling module uses CLIP and straight-through Gumbel softmax to select value-specific trainable embeddings, which are concatenated into each symbolic representation.Source: §3.2 · p. 5

Action conditioning

• Action attention learns a permutation alignment between actions and unordered entity slots before transition modeling.This reordering makes each action correspond to its associated entity, such as aligning an action with the slot representing a circle.Source: §3.2 · pp. 5–6

End-to-end pipeline

• The complete model repeats update, selection, and transformation steps to represent multi-object interactions while remaining end-to-end differentiable.Figure 3 describes one step; the implementation repeats the process multiple times, and the rule, attention, and entity-extraction components are trained as a differentiable pipeline.Source: §3 · p. 4, Figure 3 · p. 6

Modular transition model

• Neurosymbolic attention selects a primary entity, contextual entity, and update rule by matching symbolic entity encodings with neural rule queries.The modular transition model updates only the selected pair, avoiding the spurious interactions produced when monolithic models score every object pair.Source: §3.3 · p. 6

Results

• COSMOS achieves the best next-state reconstruction performance across entity-composition and relational-composition benchmarks.The evaluation reports next-state reconstruction error using MSE, averaged across three seeds.Source: Table 1 · p. 7, §4 · p. 7

Experimental Setup

• The benchmark uses images of moving 2D shapes whose objects vary by color, shape, position, and push dynamics, with training and testing compositions separated.Entity composition holds out object combinations, while relational composition shares dynamics among objects with matching attributes.Source: §4 · p. 7

Results

• Relational composition reduces performance relative to entity composition, underscoring the greater complexity of shared-dynamics generalization.The evaluation includes team and sticky relational-composition settings in the block-pushing domain.Source: §4 · p. 7
• Without neurosymbolic grounding, slot autoencoder encoding and decoding degrade during training, suggesting protection against representation collapse.The reported collapse is an observed degradation rather than a separate causal intervention analysis.Source: §4 · p. 8
• Equivariant MRR lacks a consistent cross-model pattern and can be high when autoencoder reconstruction is poor, limiting its value as a downstream-utility indicator.The paper motivates downstream planning evaluation because partial representation collapse can artificially improve retrieval-based MRR.Source: §4 · p. 8

Downstream Utility

• COSMOS shows the most consistency and least deviation from goal states in greedy-planner evaluations across all datasets.Downstream utility is measured by average L1 error between chosen next states and ground-truth next states over trajectory depth; lower error is better.Source: §4 · p. 8, Figure 4 · p. 8

Interaction Modeling and Modularity in Representations

• Prior world-modeling work uses object-factored representations, graph neural networks, or neural production systems to model structured dynamics and sparse interactions.COSMOS is influenced by the neural production-system architecture while introducing symbolic routing.Source: §5 · p. 8
• COSMOS extends modular interaction modeling by selecting learnable symbol subsets per entity instead of maintaining a single type vector.The selected symbols route neural state encodings to appropriate transition modules and support global and local invariances.Source: §5 · pp. 8–9

Compositional Generalization and Equivariance in World Models

• Equivariant and homomorphic approaches connect symmetry structure with compositional generalization in neural world models.The related work frames COSMOS within research on symmetry-aware state-action representations and homomorphic MDPs.Source: §5 · p. 9

Vision-grounded Neurosymbolic Learning

• Vision-grounded neurosymbolic learning uses symbolic information to support generalization and data efficiency, while differing in how symbolic and neural encodings interact.COSMOS belongs to this broader family but obtains attributes from vision-language foundation models for module selection.Source: §5 · p. 9

Foundation Models as Symbol Extractors

• Unlike approaches that require symbols to reconstruct every perceptual detail, COSMOS uses symbols primarily to select transition modules while neural encodings capture fine-grained dynamics.This design avoids requiring symbolic attributes to represent all dynamics-relevant information.Source: §5 · p. 9

Conclusion

• COSMOS demonstrates that annotating entities with symbolic attributes can enhance compositional generalization in world models.The conclusion covers both entity composition and the paper’s new relational-composition definition.Source: §6 · p. 9
• Foundation models can provide the symbolic attributes used by COSMOS without requiring the paper’s traditional human-provided perceptual-to-symbol mapping.The conclusion characterizes this as obtaining symbols “for free” from foundation models.Source: §6 · p. 9

Future Work

• The study is an initial demonstration, with future work targeting richer symbolic representations, neurosymbolic transformers, and more compositional-world-modeling benchmarks and datasets.These directions are presented as responses to rapidly advancing foundation models and limited available evaluation artifacts.Source: §6 · p. 9

7 REPRODUCIBILITY AND ETHICAL CONSIDERATIONS

• The authors plan to release source code, pretrained models, and setup environments under the MIT license to support replication.They also describe improvements to public repositories used for data generation and evaluation.Source: §7.1 · p. 10
• Methodological reproducibility is supported through detailed model descriptions, Algorithm 1, transition-function documentation, and appendix training details.Source: §7.1 · p. 10
• World models may be misused for societal harm, although training on synthetic block-pushing datasets mitigates this potential.Source: §7.2 · p. 10
• The paper notes privacy risks from behavioral profiling on unsupervised real-world data and bias risks from foundation models trained on biased data.No human-centric data were collected or retained in this project.Source: §7.2 · p. 10

A APPENDIX

• Section A.1 explains how SAM generates set-structured representations, while A.2 surveys the studied composition types.Source: A APPENDIX · p. 14
• Section A.3 presents a faster module-selection algorithm, and A.4 covers experiment details and the Equivariant MRR metric.Source: A APPENDIX · p. 14
• The appendix also includes planning-experiment details and additional implementation or training material.Source: A APPENDIX · p. 14

A.1 GENERATING SET-STRUCTURED REPRESENTATIONS WITH SEGMENT ANYTHING

• SAM receives the image and an 8 × 8 point grid, producing 64 × 3 potential masks across three channels.Source: §A.1 · p. 14
• Duplicate masks are removed, and combinations of remaining slots are evaluated to select a k-tuple that collectively describes the image.The selection enforces both property-specific masks and complete image coverage.Source: §A.1 · p. 14
• Selected masks are point-wise multiplied with the image, and each resulting masked image is passed through a finetuned ResNet to produce an entity representation.Source: §A.1 · p. 14

A.2 TYPES OF COMPOSITIONS

• Entity composition requires a world model equivariant to object replacement because environmental dynamics depend on which objects are present.Source: §A.2 · p. 14
• Relational composition includes entity-composition requirements and additionally imposes constraints on observable object attributes.Source: §A.2 · p. 14
• In sticky block pushing, same-colored adjacent objects move together when one is acted upon, enabling comparison of adjacency and color constraints.This setup tests whether adjacency places a greater burden on dynamics learning than color.Source: §A.2 · pp. 14–15

A.3 TRANSITION ALGORITHM

• O(k^2l) becomes O(kl + k) by partially applying query-key attention during transition selection.The faster algorithm separates selecting the primary slot and rule from selecting the contextual slot.Source: §A.3 · p. 15
• The first attention stage jointly selects the primary slot and rule before applying the selected rule module.The procedure uses Gumbel-Softmax attention and then computes the updated primary representation with the chosen rule.Source: §A.3 · p. 15
• A second attention stage selects the contextual slot using the updated primary representation as the query.This bifurcation enables the reduced runtime while retaining rule-conditioned pairwise interaction selection.Source: §A.3 · p. 15

Dataset generation

• The datasets contain 3,000 random-action trajectories of length 32, with disjoint training and evaluation composition sets.States are 3 × 224 × 224 images, and actions are factorized by object ID and cardinal direction.Source: §A.4 · p. 15

Baselines and procedure

• The comparison includes GNN homomorphic world models and ALIGNEDNPS, a symbol-free NPS ablation, using reconstruction loss for a fair comparison.Training uses three random seeds, 50 epochs, and a maximum batch size of 64 on one 48 GB NVIDIA A40 GPU.Source: §A.4 · pp. 15–16

Evaluation results

• COSMOS achieves the best next-state reconstructions across the entity-composition and relational-composition datasets.Table 2 averages results across three random seeds and reports standard deviations.Source: Table 2 · p. 16

Downstream evaluation

• Downstream planning evaluates L1 xy-coordinate error over 500 test trajectories of length 32 and averages error at each trajectory depth.Because predictions are recursively fed into later steps, divergence from ground truth increases with trajectory depth; COSMOS shows the least deviation across datasets.Source: §A.5 · pp. 16–17

Evaluation setting

• The evaluation domain is selected for object-oriented, history-invariant, action-conditioned dynamics and easy generation of new configurations.The domain is the established 2D block-pushing benchmark used in prior work.Source: §A.6 · p. 17

Symbolic ablation

• The symbols-only ablation significantly underperforms COSMOS because concatenated symbolic attributes encourage overfitting to training-time attribute compositions.COSMOS instead uses symbols to select rules while retaining real vectors for local action-conditioned transition features.Source: §A.7 · p. 17