Neurosymbolic Grounding for Compositional World Models
Authors: Atharva Sehgal
Arya Grayeli
Jennifer J. Sun
Swarat Chaudhuri
Venue: ICLR 2024 poster (2024)
Keywords: neurosymbolic learning
machine learning
world modeling
compositional generalization
Primary area: neurosymbolic & hybrid AI systems (physics-informed, logic & formal reasoning, etc.)
TL;DR: We study a new form of compositional generalization and develop a hybrid neurosymbolic world model for this form of compositional generalization..

Abstract: We introduce Cosmos, a framework for object-centric world modeling that is designed for compositional generalization (CompGen), i.e., high performance on unseen input scenes obtained through the composition of known visual "atoms." The central insight behind Cosmos is the use of a novel form of neurosymbolic grounding. Specifically, the framework introduces two new tools: (i) neurosymbolic scene encodings, which represent each entity in a scene using a real vector computed using a neural encoder, as well as a vector of composable symbols describing attributes of the entity, and (ii) a neurosymbolic attention mechanism that binds these entities to learned rules of interaction. Cosmos is end-to-end differentiable; also, unlike traditional neurosymbolic methods that require representations to be manually mapped to symbols, it computes an entity's symbolic attributes using vision-language foundation models. Through an evaluation that considers two different forms of CompGen on an established blocks-pushing domain, we show that the framework establishes a new state-of-the-art for CompGen in world modeling. Artifacts are available at: https://trishullab.github.io/cosmos-web/

Published as a conference paper at ICLR 2024
NEUROSYMBOLIC GROUNDING FOR COMPOSITIONAL WORLD MODELS
Atharva Sehgal UT Austin
Arya Grayeli
Jennifer J. Sun Caltech
Swarat Chaudhuri UT Austin
UT Austin
ABSTRACT
We introduce COSMOS, a framework for object-centric world modeling that is designed for compositional generalization (CompGen), i.e., high performance on unseen input scenes obtained through the composition of known visual "atoms." The central insight behind COSMOS is the use of a novel form of neurosym- bolic grounding. Specifically, the framework introduces two new tools: (i) neu- rosymbolic scene encodings, which represent each entity in a scene using a real vector computed using neural encoder, as well as vector of compos- able symbols describing attributes of the entity, and (ii) a neurosymbolic atten- tion mechanism that binds these entities to learned rules of interaction. Cos- MOS end-to-end differentiable; also, unlike traditional neurosymbolic meth- ods that require representations to be manually mapped to symbols, it putes an entity's symbolic attributes using vision-language foundation models. Through an evaluation that considers established blocks-pushing domain, we show that the framework establishes a new state-of-the-art for CompGen in world modeling. Artifacts are available bttps::/trishut ab. github. 1c/ coscoossmswbb
different forms of CompGen on an
INTRODUCTION
The discovery of world models - deep generative models that predict the outcome of an action in scene made of interacting entities - is central challenge in contemporary machine learning (Ha & Schmidhuber 2018 Hafner et al. 2023). As such models are naturally factorized by objects, methods for learning them (Zhao al. 2022 Goyal et al. 2021: Watters et al. 221 9a Kipf al.) 2019) commonly follow a modular, object-centric perspective. Given a scene represented as pixels, these methods first extract representations of the entities in the scene, then apply transition network to model interactions between the entities. The entity extractor and the transition model form an end-to-end differentiable pipeline.
Of particular interest in world modeling is the property of compositional generalization (Comp- Gen), i.e., test-time generalization to scenes that are novel compositions of known visual "atoms". Recently, Zhao et al. (2022) gave first approach to learning world models that compositionally generalize. Their method uses an action attention mechanism bind actions entities. The mech- anism is equivariant to the replacement of objects in a scene by other objects, enabling CompCien.
This paper continues the study of world models and compositional generalization. We that such generalization is hard for purely neural methods, as they cannot easily learn encodings that can be decomposed into well-defined parts. Our approach, COSMOS, uses a novel form neurosymbolic grounding to address this issue.
The centerpiece idea in COSMOS the notion of object-centric, neurosymbolic scene encodings. Like in prior modular approaches to world modeling, extract a discrete set of entity represen- tations from an input scene. However, each of these representations consists of: (i) a standard real vector representation constructed using neural encoder, and (ii) vector of symbolic attributes, capturing important properties for example, shape, color, and orientation - of the entity.
Like Goyal et al. (2021). model transitions in the world collection of neural modules, each capturing a "rule" for pairwise interaction between entities. However, in contrast to prior work, we bind these rules to the entities using novel neurosymbolic attention mechanism. In version of

---

Published as a conference paper at ICLR 2024
Evaluation Unseen Composition
Entity Composition
Generalization to compasitions
objects during training
Distribution of Ato
Sampled Composition
Image (799)
Ground Truth Next Image /(++3)
Prediction (Ours)
(Baseline)
Actions
Relational Composition
Scene Composition Objects with attributes have shared dynamics
Sampled Composition A)
Image (y(0)
Ground Truth Next Image z(++)
Prediction (Baseline)
Prediction (Ours)
Actions
Figure 1: Overview of compositional world modeling. We depict examples from 2D block pushing domain consisting of shapes that interact, where we can generate samples of different shapes and interactions. We aim to learn model that generalizes to compositions not seen during training, such as entity composition (top) and relational composition (bottom). Previous works (Goyal et al., 2021] focus on entity composition, and struggle generalize harder compositional environments. Our approach COSMOS leverages object-centric, neurosymbolie scene encodings to compositionally generalize across settings containing different types of compositions.
such attention, the keys symbolic and the queries neural. The symbolic keys matched with the symbolic components of the entity encodings, enabling decisions such "the i-th rule represents interactions between a black object and a circular object" (here, "black" and "circular" are symbolic attributes). The rules, the attention mechanism, and the entity extractor constitute an end-to-end differentiable pipeline.
The CompGen abilities of this model stems from the compositional nature of the symbolic attributes. The symbols naturally capture "parts" of scenes. The neurosymbolic attention mechanism fires rules based on (soft, neurally represented) logical conditions over the symbols and can cover new compositions of the parts.
A traditional issue with neurosymbolic methods for perception is that they need human-provided mapping between perceptual inputs and their symbolic attributes (Harnad) 1990 Tang & Ellis, 2023 However, COSMOS automatically obtains this mapping from vision vision-language foundation
models. Concretely, to compute the symbolic attributes of an object, we utilize CLIP (Radford et al.
2021 to assign each object values from a set of known attributes.
We compare COSMOS against a state-of- of-the-artt baseline and an ablation on an established domain of moving 2D shapes (Kipf et al. 2019). Our evaluation considers definitions of CompGen, illustrated in Figure
one of these, we want the model to generalize to new entity compositions, of entities that have been during training, but never together. The other definition, relational composition, is new to this paper: here, we additionally need to accommodate shared dynamics between objects with shared attributes (e.g., the color red). Our results show that COSMOS outperforms the competing approaches at next-state prediction (Figure I] visualizes a representative sample) and separability of the computed latent representations, well as accuracy in downstream planning task. Collectively, the results establish COSMOS as state- of-the-art for CompGen world modeling.
i.., input scenes comprising
To summarize our contributions, we offer:
CoSMOS, the first differentiable neurosymbolic approach based on a combination of neurosym- bolic scene encodings and neurosymbolic attention object-centric world modeling.
new way use foundation models to automate symbol construction neurosymbolic learning:
an evaluation that shows COSMOS to produce significant empirical gains over the state-of-the for compositionally generalizable world modeling.
2

---

Published as a conference paper at ICLR 2024
2 PROBLEM STATEMENT
We are interested in learning world models that compositionally generalize. World models arise naturally out of the formalism for Markov decision processes (MDPs). An MDP M is a tuple (S, T, R. 7) with states S, actions A, transition function T S x A + S x 22 reward function R x A R, and discount factor y [0, 1]. We make three additional constraints:
Object-Oriented state and action space: In our environments, the state space is realized as images. At each time step I, agent observes image I E Spizee and takes action A E A. However, learning Spizet directly is an intractable problem. Instead, we assume that the high dimensional state space can be decomposed into lower dimensional object-oriented state space S S X x Sk where k is the number of objects in the image. Now, factor S; E S, describes a single object in the image. Hence, the transition function has signature T (S1; A1 x-. x Ss; A&) - (S1x... x... x xS)) where each factor A; is factorized set of actions associated with object representation S; and (o: the concatenation operator. A pixel grounding function P4 : S + Spinet enables grounding a factored state into an image.
Symbolic object relations: We assume that objects in the state space share attributes. An at- For instance, in the 2D block push- ing domain (illustrated in Figure D. each object has a "color" attribute that can take on values Ceolor == (zed, green.... J. An object can be composed of many such attributes that can be re- trieved using an attribute projection function a+ : S, - A; == AA. where (CI,-..: Cp) is predefined, ordered list of attributes and A, selects the value in the Cp-th attribute that is most relevant to St. Note that a+ only depends on a single object and trivially generalizes to different compositions of objects (Keysers et al. 2020).
tribute is set of unique symbols Cp p..... C.
Compositional Generalization: We assume that each state in our MDP can be decomposed using two sets of elements: compounds and atoms. Compounds are sets of elements that can be decom- posed into smaller sets of elements, and atoms are sets of elements that cannot be decomposed further. For instance, in the block pushing domain (Figure we can designate each unique object as an atom and designate the co-occurrence of a set of atoms in a scene as a compound. We use A denote the atoms and C denote the compounds. The frequency distribution of distribu- tion o is denoted F. Given this, CompGen is expressed as a property of the train distribution Dirain and of the test distribution Drest undergoing a distributional shift of the compounds, while the distribution of atoms remains the same.
Given the following assumptions, for any experience buffer D == {f(r(), A (((T), S. Spizet X A) learning a world model boils down to learning the transition function T for the MDP using the sequences of observations collected by D. Specifically, for sfi.tt E ((tt)),AA))) A(10), r(r++1) E D, our objective function
T ΣΣ
IIP. (T (5((t); A((0)) ..... $6.00 , Al11)) - 10,+
L N(T-T)
==1 E=1
We study two kinds of compositions:
Figure shows an instance of entity centric CompGen in the block pushing domain. Here, there exist = 9 unique objects, but only A: = 3 are allowed to co-occur in any given realized scene. Each unique object represents an atom, and the co-occurrence of a set of & atoms in a scene represents the compound. Hence, there are total (7) possible compositions. The distribution of atoms in the train distribution and the test distribution does not change, while the distribution of compounds train time and test time are disjoint. So, FA(Deval) is equivalent to Fa(Dtrain)) while Fe(Devat) O Fe(D (Devain) = Ø.
Entity Composition:
Relational Composition: Figure shows an instance of relational composition in the block pushing domain. Here, there are n = 9, k = 3 unique objects. A composition of objects occurs when two objects share attributes (for instance, here, the shared attributes are color and adjacency). Objects with shared attributes share dynamics. As a result, if one object experiences force, others with the same attributes also undergo that force. Hence, assuming a single composition per scene, there are total (H) (*) possible compositions. As in the previous case, FA(eeval) is equivalent (Derain)) while Fe(Deeaat A Fe(Drrain) Ø.
3

---

Published as a conference paper at ICLR 2024
Prior Work Entity Extractor
Spotial Decoder
Module
Concat
Selector
Executor
Image
Actions
Ours
Entity Extractor
Decoder
Ardd
Arade Aquan
Module Selector
Sg
Ansi
Attention Concat
Audesh
Execunor
Image
Symbolic Labeling Module
A
Actions
Neurosymbolic Representations
Figure 2: Comparing world modeling frameworks between prior work (Goyal et al. 2021) and Cos- MOS. Both modules start with entity extraction, obtain neural object representations {S Sk} from the image (Section B.I.. While prior work uses this representation directly for the module se- lector, our work leverages symbolic labeling module, which outputs set of attributes A, to learn neurosymbolic representations (Section 3.2). We then perform action conditioning (Section 3.2 to keep track of corresponding actions, and update through transition model (Section 3.3
3 METHOD
We approach neurosymbolic world modeling through object-centric neurosymbolic representations (Figure 2 Our method consists of: extracting object-centric encodings (Section (3.I). representing each object using neural and symbolic attributes (Section 3.2). and learning transition model to update to next state (Section (.3). Similar to previous works, we use slot based autoencoders to ex- tract objects, but use symbolic grounding from foundation models alongside neural representations, which enables our model to achieve both entity and relational compositionality. The full algorithm is presented in Algorithm M and the model is visualized in Figure We will make of the example in Figure to motivate our algorithm.
Algorithm Neurosymbolic world model for k-factorized state space. The input is an image I of dimensions (C, H. W) = (3x 224 224) and a set of k-factorized actions dimension (A: x dact ion). The model is trained end-to-end except for the SLM modules, whose weights are held constant. There three global variables: T, {C1, -.. Cp} and (R, R.). T is the threshold of the number of repeated slot update steps, [CI,... Cp) denotes text encodings for the closed vocabulary and [ī,,.. } are learnable rule encodings.
function TRANSITIONIMMAGEU,, (A Ac)
2: (Si.... .k) < E-NTITYXXTRAATTRU))))
b. S, dim: (k, desot)
3: for in range(T)) do
4: [K.... K). (M,
M/) SPATIALDECODER((S.... Sz))
1 dim: (C, H. W)
[l,,... [.....4] < (M' F. Vi [],
6:
(Ai.... A.] < (SLM(1, (CI.... Cp)) [1,
b dim: (k.p)
7:
(A,, ......AA3 ACTIONATIN((AI. Ax), [A, Ak))
8:
P,c,r MODULARRULLNET Ar)
Aa), ( R())
9:
Sp Sp MLLLBAN(r/] (concatt (S,. S., R.))
10:
[K K.. (M/,... .MM + SPATIALDECODER((S..)) S.))
b 1 dim: (C, H, W)
11:
return Ei= -
3.1
SLOT-BASED AUTOENCODER
blot-based autoencoder transforms an image into a factorized hidden representation, with each factor representing single entity, and then reconstructs the image from this representation. Such
4

---

Published as a conference paper at ICLR 2024
autoencoder make two assumptions: (1) each factor captures specific property of the image (2) collectively, all factors describe the entire input image. This set-structured representation enables unsupervised disentanglement of objects into individual entities. Our slot based autoencoder has two components:
The entity extractor takes image and produces set- structured hidden representation describing each object. In our domains, slot attention and derivate works (Chang et al. 2022; Jia al. 2022) struggle to disentangle images into separate slots. To avoid the perception model becoming a bottleneck for studying dynamics learning, we propose a new entity extractor that uses Segment Anything Kirillov et al. (2023) with pretrained weights to produce set-structured segmentation masks to decompose the image into objects and Resnet- image encoder He al. (2016) to produce set-structured hidden representation for each object. This model allows us to perfectly match ground truth segmentations in our domains while preserving the assumptions of set structured hidden representations. More details are presented Section A.
ENTITYEXTRACTOR - ( ..... S
DECDER {S,,.... S + -/,.... ......I, K.J. {M' ..V........ MM): The slot decoder is a permuta- tion equivariant network that decodes the set of slots back to the input image. We follow previous works (Locatello et al.] 2020 Zhao et al. 2022, Goyal et al.) T02T) in using spatial decoder (Wat- ters et al. 2019b) that decodes a given set of vectors {Si, Sk individually into a set of image reconstructions { .............................. 1/) and a set of mask reconstructions (M/...., M,]. The final image is produced by taking the Hadamard product of each reconstruction and its corresponding mask and adding all the resulting images together. That is, I - M/
3.2 NEUROSY ENCODING
To achieve robust CompGen in our setting, our representation must be resilient to both object and attribute compositions. The slot-based encoding, by construction, generalizes to object replace- ment. For attribute compositions, however, it is essential to know the exact attributes to be targeted CompGen. Given these attributes, we propose describing each entity with a composition of symbol vectors. Each symbol vector is associated with a single entity and attribute, allowing it to trivially generalize to different attributes compositions. Moreover, we can ensure a canonical order- ing of the symbols, making downstream attention based computations invariant to permutations of attributes. We'll Il next detail our method for generating these symbol vectors.
The symbolic labelling module (SLM) pro- cesses an image and predefined list of attributes. Assuming this list comprehensive (though not exhaustive), the module employs pretrained C1.IP model to compute attention scores between the image features and each entity encoding. The resulting logits indicate the alignment between the image and each attribute. The attribute most aligned with the image is then identified using straight-through Gumbel softmax (Jang et al., 2016). The gumbel- softmax yields the index of the most likely value for each attribute one-hot vectors, which are concatenated together to form bit- vector. However, the discrete representation of such bit-vectors do not align well with downstream attention based modules so, instead of directly using one-hot vectors, the gumbel-softmax selects a value-specific encoding for each attribute. Thus, the resultant symbol vector is a composition of learnable latent encodings distinct each attribute value. In implementation. SLM utilizes another zero-shot symbolic module (a spatial-softmax operation (Watters et al.) (019b)) to extract positional attributes, such as the x and y position of the object, from the disentangled input vector.
SI. T, x (CI,...Cp)
A; == (A A
For instance, in Figure 3 SLM takes in the slot corresponding to the circle and a list of attributes (shape, color, etc.) and selects the most relevant attribute value ('circle', 'red', etc.). We subse- quently select trainable embeddings corresponding to each attribute value and concatenate them to construct the symbolic embedding.
Ag}: (A...., For accurate next state prediction, world models must condition the state on the corresponding action. Typically, this is done by con- catenating each action to its associated encoding if the slots have canonical ordering. However, the entities in a set structured representation do not have a fixed order. To find such a canonical ordering, we follow Zhao et al. (2022) in learning a permutation matrix between the actions and the slots and using the permutation matrix to reorder the slots and concatenate them with their respective actions. Attention, by construction, is equivariant with respect slot order, which avoids the need
ACTIONATTN: (AI.... A&3 {AI....AR} +
5

---

Published as a conference paper at ICLR 2024
each slot
Veston Foundation Model
Ceppo
8.75
M
A1
Ae
Ppppo
8.25
Ag
Cooloe
Blue
M
Gshagge
Symbolic Embeddings
MLP,
S1
Sj
MLL,s
.g
% Latent Embeddings
Sp
Keys
MLP,
(Ra..... Bi)
Image
Encoder
Neurasymbolic Attentton-based Module Selector
Decoder
MSE
Figure 3: A single update step of COSMOS. The image / is fed through slot-based autoencoder and a CLIP model generate the slot encodings [S 5 and the symbol vectors [A1:-.. .. A₂)- The actions and the symbolic encoding action attention module, which are used to select the update rule to be applied to the slots. This figure depicts single update step: in implementation, the sppdte-seleecttrrnsfon step repeated multiple times to model multi-object interactions.
aligned and concatenated using permutation equivariant
to enforce a canonical ordering. For instance, in Figure 3 if A₁ corresponds to the circle, a trained ACTIONATTN module reorders the actions so that A₁ corresponds to going downwards.
3.3 TRANSITION MODEL.
Monolithic transition models like GNNs and MLPs model every pairwise object interaction, leading to spurious correlations in domains with sparse interactions. Our modular transition model addresses this problem by selecting relevant pairwise interaction and updating the encodings for those objects locally (following Goyal et al.] (202I). We model this selection process by computing key-query neurosymbolic attention between ordered symbolic and neural rule encodings determine the most applicable rule-slot pair. As each entity is composed of shared symbol vectors, the dot-product activations between the symbolic encodings and the rule encodings remain consistent for objects with identical attributes. Next, we will discuss the mechanisms of module selection and transitions.
MODULESELECTOR [A1,...Ayy (R..... + (p. c, r): The goal of the selection process is select the primary slot, the contextual slot (which the primary slot will be conditioned on), and the update function to model the interaction. Query-key attention (Bahdanau et al. (2015) serves as the base mechanism the rule encoding and the action this selection will take (1+71) time, where ie is the number of slots and 1 the number of rules. In implementation, the selection of (p, c, r) can be reduced to runtime of Okl + k) by partial application of the query-key attention. This algorithm is presented in Section of the appendix.
will be using perform selections. We compute query-key attention between
symbolic encoding. The naive algorithm to compute
MLPBANK : + MLP,: Our modular transition function comprises a set of rules R....., .R,, with each rule defined R, (H², MLP,). ). Here, R, is lcarnable encoding, while MLP; Sp x S. S represents a submodule that facilitates pairwise interactions between objects. Intuitively, each submodule is learning how a primary state changes conditioned secondary state. In theory, each update function can be customized for different problems. However, in this study, we employ multi-layered perceptrons for all rules.
EXPERIMENTS
We demonstrate the effectiveness of COSMOS the 2D Block pushing domain (Kipf et al.) 2019; Zhao et al.) 2022] Goyal et al. 202T] Ke et al. 202T) with entity composition, and two instances of relational composition. We selected the 2D block pushing domain it is widely-adopted and

---

Published as a conference paper at ICLR 2024
objects
5 objects
Dataset
Model
MSE 4
AE-MSE 4
EqMRR
MSE +
AE-MSE
Eq f
RC
COSMOS
4.23F-00
4.90E-04
1.20E-01
4.15E-03
1.68E-03
3.67E-01
(Sticky)
ALIGNEDNPS
1.14E-02
7.72E-03
8.01E-02
6,07E-03
2.47E-03
3.62E-01
GNN
7.94F-03
5.11E-03
6.03E-04
6.21E-03
2.73E-03
5,30F-04
RC
COSMOS
4.60E-03
4.33E-04
1,04-001
5,53E-03
186E-03
2.86E-01
(Team)
ALIGNEDNPS
1.24E-02
8.36E-03
1.75E-01
9.64E-03
3.12E-03
2.93E-01
GNN
8.92E-03
3,82E-03
7.16E-04
7.01E-03
1.62E-03
5.46B-04
EC
COSMOS
7.66F.004
6.34E-055
2.99E-01
4.08F-04
2.92F-066
3.03F-01
ALIGNEDNPS
3.51E-03
2.69E-03
2.97E-01
2.45E-03
1.22E-03
3.19E-01
GNN
9.89F-03
1.03E-02
5.50E-01
1.20F-02
1.28F-02
5,25E-01
Table 1: Evaluation results on the 2D block pushing domain for entity composition (EC) and re- lational composition (RC) averaged across three seeds. We report next-state reconstruction er- ror (MSE), autoencoder reconstruction error (AE-MSE), and the equivariant mean reciprocal rank (Eq.MRR) for three transition models: our model (COSMOS), an improved version of Goyal et al. 2222 (ALIGNEDNPS), and a reimplementation of Zhao al. (2022) (GNN). Our model (COSMOs) achieves best next-state reconstructions for all datasets.
well-studied benchmark in the community. Notably, even in this synthetic domain, our instances of compositional generalization proved challenging to surpass for baseline models, underscoring the difficulty of the problem.
2D Block Pushing Domain: The 2D block pushing domain (Kipf et al. 2019 Ke et al. 2021 is a
two-dimensional environment that necessitates dynamic and perceptual reasoning. Figure presents an overview of this domain. All objects have four attributes: color (A (eCoiio) )) shape (A (isps)) ), x position Cr-eo)), and y position (ACr-e)). Objects can be pushed in one of the four cardinal directions -North-East-Soute West). Heavier objects can push lighter objects, but not the other way around. The weight of the object depends on the shape attribute. At each step, the agent observes an image of size 3 x 224 x 224 with k objects and an action pushing one of the objects. This action chosen from a uniform random distribution. The goal is for the agent to capture the dynamics of object movement. Furthermore, while there are n |Acoioz| x
AAsaspel unique objects in total,
only & < n objects are allowed co-occur in a realized scene.
We adapt the methodology from (Kipf et al. 2019 Z.hao et al. 2022 with minor changes. For entity compositions (EC), we construct training and testing datasets to have unique ob- ject combinations between them. In relational composition (RC), objects with matching attributes exhibit identical dynamics. Two specific cases are explored: Team Composition (RC-Team) where dynamics are shared based color, and Sticky Composition (RC-Sticky) where dynamics are shared based on color and adjacency. Further details are in the appendix (Section A..2). Our data gen- eration methodology ensures that the compound distribution is disjoint, while the atom distribution remains consistent across datasets, i.e. Fe Fe(Drrain) Fe (Deeatt = and Fa (Drrain)) = FA (D (Demat). The difficulty of the domain can be raised by increasing the number of objects. We sample datasets for 3 and 5 objects.
Dataset Setup:
We compare against past works in compositional world modeling with publically avail- able codebases at the time of writing. For the block pushing domain, we compare with homomor- phic world models (GNN) (Zhao al. 2022) and NPS (Goyal et 2021 with modifications to ensure that the actions are aligned with the slots PAS) The latter is equivalent to an ablation of COSMOS without the symbolic labelling module. We analyzed next-state predictions using mean squared error (MSE), current-state reconstructions using auto-encoder mean squared er- ror (AF-MSE), and latent state separation using (MRR) following previous work (Kipf al. 2019, Zhao et al. 2022; Goyal et al. 2021). Recognizing limitations in existing MRR calculation, introduce the Equivariant MRR (Eq.MRR) metric, which accounts for different slot orders when calculating the MRR score. More details can be found in the appendix Section A.4
Evaluation:
Results: Results are presented in Table 2 First, we find that COSMOS achieves the best next state prediction performance (MSE) all benchmarks. Moving from entity composition to relational composition datasets shows a drop in performance, underscoring the complexity of the task. Sur- prisingly, there is less than expected performance degradation moving from the three object five
7

---

Published as a conference paper at ICLR 2024
object domain. We attribute this to the nature of the block pushing domain and the choice of loss function. MSE loss measures the pixel error between the predicted and next image. As the density of objects in the image increases, the reconstruction target becomes more informative, encouraging better self-correction, and hence more efficient training.
Second, observe that, without neurosymbolic grounding, the slot autoencoder's ability encode and decode images degrades as the world model training progresses. This suggests that neurosym- bolic grounding guards against auto-encoder representation collapse.
Finally, we do not notice a consistent pattern in the Equivariant MRR scores between models. First, all models tend exhibit higher Eq.MRR score in the five object environments. How- ever, in many cases, models with high Eq.MRR score also have underperforming autoencoders. For instance, in the five object entity composition case, the GNN exhibits high Eq. score yet simultaneously has the worst autoencoder reconstruction error. We notice this happens when the model suffers partial representation collapse (overfitting to certain color, shapes combina- tion seen during training). This maps many slot encodings to the same neighborhood in latent space; making easier to retrieve similar encodings, boosting the MRR score. Given these observations, we conclude that MRR might not be optimal indicator of model's down- stream utility. For a comprehensive assessment of downstream utility, we turn to the methodol- ogy outlined by Veerapaneni et al. (2020), applying the models to a downstream planning task.
Downstream Utility: We evaluate the downstream utility all world mod- els using a simple planning task in 2D shapeworld environments. Our method,
Object Composition Object Composition Scene Composition (Sticky)
(Team)
Avengee torm tettoen Current Gool toee
30
AligndNPP
Hadd
Objects
GWN
20 Lo
inspired by Veerapaneni et al., (2020), in-
volves using a greedy planner that acts based on the Hungarian distance be-
tween the predicted and goal states. Due
to the compounding nature of actions over trajectory, there is an observed
20
Objects
divergence from the ground truth the deeper we get into the trajectory. We run these experiments on our test dataset and average the scores each trajectory depth. We showcase results in Figure Our model (CosMOs) shows the most consistency and least deviation from the goal state in all datasets, which suggests that neurosymbolic grounding helps im- prove the downstream efficacy of world models. More details can be found in appendix section A.5
20
10
20 10 20
20
10
20 20
Trajectory Length
Figure 4: Downstream utility of different world models using a greedy planner. The graph follows the average L1 error between the chosen next state and the ground truth next state as function of the number of steps the model takes. A lower LI error indicates better perfor- mance. COSMOS (in red) achieves the best performance.
5 RELATED WORK
Object-Oriented World Models: Object-oriented world models (Kipf et al. 2019 Zhao et al.,
2022 Van der Pol et al. 2020a) Goyal et al. 2021; Veerapaneni et al. (020) are constructed to learn structured latent spaces that facilitate efficient dynamics prediction environments that can be decomposed into unique objects. We highlight two primary themes in this domain.
Interaction Modeling and Modularity Representations: Kipf et al. (2099) builds contrastive world modeling technique using graph neural networks with an object-factored repreetttionn...... How- ever, GNNs may introduce compounding errors in sparse interaction domains. Addressing this, Goyal et al. (2021) introduces neural production systems (NPS) for modeling sparse interactions, akin to repeated dynamic instantiation GNN edge instantiation. CoSMOS is influenced by the NPS architecture, with distinctions highlighted in Figure 2 In parallel, Chang et al. (2023) studies hi- erarchical abstraction for world models, decomposing slots into a (dynamics relevant) state vector and (dynamies invariant) type vector, with dynamics prediction focusing the state vector. In
8

---

Published as a conference paper at ICLR 2024
COSMOS, instead of maintaining single "type vector", we maintain set of learnable symbol vec- tors and select a relevant subset for each entity. This allows COSMOS to naturally discover global and local invariances, utilizing them to route latent encodings (akin to "state vectors") to appropriate transition modules.
Compositional Generation and Equivariance in World Models: Equivariant neural networks har- ness group symmetries for efficient learning (Ravindran] 2004] |Cohen & Welling 2016a b Walters) 2004|CCohen
al. 2020, Park et al. 2022). In the context of CompGen, Van der Pol et al. (2020b) investigates constructing equivariant neural networks within the MDP state-action space. Zhao et al. (2022)
establishes a connection between homomorphic MDPs and compositional generalization, express- ing CompGen as a symmetry group of equivariance to object permutations and developing world model equivariant to object permutation (using action attention). We adopt this idea, but integrate a modular transition model that also respects permutation equivariance in slot order.
Vision-grounded Neurosymbolic Learning: Prior work in neurosymbolic learning has demon- strated that symbolic grounding helps facilitate domain generalization and data efficiency (Andreas) cet al.
2016a b; Sun et al. 2020; Zhan et al. [02T] [Shah et al. 2020; Mao et al., 2019 Hsu et a Yi et al.) 2018). The interplay between neural and symbolic encodings in these works can
2023
be abstracted into three categories: (1) scaffolding perceptual inputs with precomputed symbolic
information for downstream prediction (Mao et al. 2019 Ellis et al.) 2018 Andreas al. 2016b
Hsu et al., 2023 Valkov et al. 2018), (2) learning a symbolic abstraction from a perceptual input useful for downstream prediction Tang & Ellis (2023. , and (3) jointly learning neural and symbolic encodings leveraged for prediction Zhha et al. 1). Our approach aligns most closely with the third category. While Zhan et al. (2021) combine neural and symbolic encoders in a VAE setting, highlighting the regularization advantages of symbols for unsupervised clustering, they rely on a program synthesizer to search for symbolic transformations in a DDL--intoodun expressiveness issues. COSMOS also crafts a neurosymbolic encoding, but addresses scalability and expressiveness concerns of program synthesis by using a foundation model to generate the symbolic encodings.
scalability and
Foundation Models as Symbol Extractors: Many works employ foundation models to decompose complex visual reasoning tasks (Wang et al. 2023a Gupta & Kembhavi, 2023] Nayak et al. 2022) (Hsu et al. 2023; Wang et al. 2023b) decompose natural language instructions into executable pro-
grams using Code LLM for robotic manipulation and 3D understanding. Notably, ViperGPT (Suris ct al. 2023) uses Code LLMS to decomposes natural language queries to API calls in a library of pretrained models. Such approaches necessitate hand-engineering the API be expressive enough to generalize to all attributes in a domain. COSMOS builds upon the idea of using compositionality of symbols to execute parameterized modules but sidesteps the symbolic decomposition bottleneck by parsimoniously using the symbolic encodings only for selecting representative encodings. L.e., COSMOS does not need its symbols to learn perfect reconstructions. The symbolic encoding is only used for selecting modules, while the neural encoding can learn fine-grained dynamics- relevant at- tributes that may not be known ahead of time. Furthermore, to the best of our knowledge, CosMOS is the first work to leverage vision-language foundation models for compositional world modeling.
6 CONCLUSION
We have presented COSMOS, a new neurosymbolic approach for compositionally generalizable world modeling. Our two key findings are that annotating entities with symbolic attributes can help with CompGien and that it is possible to get these symbols "for free" from foundation models. We have considered two definitions of CompGen - one new to this paper - and show that: (i) CompGen world modeling still has a long way to go regarding performance, and (ii) neurosymbolic grounding helps enhance CompGen.
Our work here aimed to give an initial demonstration of how foundation models can help composi- tional world modeling. However, foundation models are advancing at a breathtaking pace. Future work could extend our framework with richer symbolic representations obtained from more power- ful vasion-language or vision-code models. Also, our neurosymbolic attention mechanism could be naturally expanded into a neurosymbolic transformer. Finally, the area of compositional world mod- cling needs more benchmarks and datasets. Future work should take on the design of such artifacts with the goal of developing generalizable and robust world models.
9

---

Published as a conference paper at ICLR 2024
7 REPRODUCIBILITY AND ETHICAL CONSIDERATIONS
7.1 REPRODUCIBILITY CONSIDERATIONS
We encourage reproducibility of our work through the following steps:
Providing Code and Setup Environments: Upon acceptance of our work, we will release the complete source code, pretrained models, and associated setup environments under the MIT license. This is to ensure researchers can faithfully replicate our results without any hindrance.
Improving Baseline Reproducibility: Our methodology extensively utilizes these public repositories (Zhao et al. 2022 Kipt et al. 2019 Goyal et al.] (021) for generating data and computing evaluation metrics. During development of our method, we made several improvements to these repositories: (1) We enhanced the performance and usability of the models, (2) We updated the algorithms to ac- commodate the latest versions of the machine learning libraries, and (3) we constructed reproducible anaconda environments for each package to work on Ubuntu 22.04. We intend to contribute back by sending pull requests to the repositories with our improvements, including other bug fixes, runtime enhancements, and updates.
Outlining Methodology Details: The details of our model are meticulously described in Section
A comprehensive overview is available in Algorithm and deeper dive into the transition function is in Appendix Section A.3 We mention training details in appendix section A.4
7.2 ETHICAL CONSIDERATIONS
Potential for Misuse: As with other ML techniques, world models can be harnessed by malicious actors to inflict societal harm. Our models are trained on synthetic block pushing datasets, which mitigates their potential for misuse.
Privacy Concerns: In the long term, as world models operate on unsupervised data, they enable learning behavioral profiles without the active knowledge explicit consent of the subjects. This raises privacy issues, especially when considering real-world, non-synthetic datasets. We did not collect retain any human-centric data during the course of this project.
Bias and Fairness: World models, generally, enable learning unbiased representations of data. How- ever, we leverage foundation models which could be trained on biased data and such biases can reflect in our world models.
8 ACKNOWLEDGEMEENS
This research was supported by the NSF National AI Institute for Foundations of Machine Learning (IFML), ARO award #W91 INF-21-1-001 and DARPA award #HR001 12320018.
REFERENCES
Jacob Andreas, Marcus Rohrbach, Trevor Darrell, and Dan Klein. Learning to compose neural networks for question answering. In Proceedings of the 2016 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies, pp- 1545-1554, San Diego, California, June 2016a.
Jacob Andreas, Marcus Rohrbach, Trevor Darrell, and Dan Klein. Neural module networks. In Pro- ceedings of the IEEE conference on computer vision and pattern recognition, pp. 39-49, 201
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. In ICLR, 2015.
Michael Chang, Tom Griffiths, and Sergey Levine. Object representations as fixed points: Training iterative refinement algorithms with implicit differentiation. Advances in Neural Information Processing Systems, 35:32694-3277 2022.
Michael Chang, Alyssa L. Dayan, Franziska Meier. Thomas L. Griffiths, Sergey Levine, and Amy Zhang. Hierarchical abstraction for combinatorial generalization in object rearrangement.
10

---

Published as a conference paper at ICLR 2024
In The Eleventh International Conference Learning Representations, ICLR 2023, Kigali, net
Rwanda, May 1-5, 2023. OpenReview.net, 2023. URL
net/pdf? pdf?
Id=EGG6vHp
Taco Cohen and Max Welling. Group equivariant convolutional networks. In International confer- ence on machine learming, pp. 2990-2999. PMLR, 2016a.
Taco S Cohen and Max Welling. Steerable cnns. arXiv preprint arXiv: 2016b.
Kevin Ellis, Daniel Ritchie, Armando Solar-L and Josh Tenenbaum. Learning to infer graph- ics programs from hand-drawn images. In Advances in neural information processing systems, pp. 6059-6068, 2018.
Anirudh Goyal, Aniket Didolkar, Nan Rosemary Ke, Charles Blundell, Philippe Beaudoin, Nicolas Heess, Michael Mozer, and Yoshua Bengio. Neural production systems. Neural Information Processing Systems, 2021.
Tanmay Gupta and Aniruddha Kembhavi. Visual programming: Compositional visual reasoning without training. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 14953-14962, 2023.
David Ha and Jürgen Schmidhuber. World models. arXiv preprint
2018.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301.04104 2023.
Stevan Harnad. The symbol grounding problem. Physica D: Nonlinear Phenomena, 42(1-3):335- 346, 1990.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Joy Hsu, Jiayuan Mao, and Jiajun Wu. Ns3d: Neuro-symbolic grounding of 3d objects and relations. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp- 2614-2623. 2023.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterrization with gumbel-softmax. arXiv preprint 1rXiv:1611. 144, 2016.
Baoxiong Jia, Yu Liu, and Siyuan Huang. Improving object-centric learning with query optimiza- tion. In The Eleventh International Conference on Learning Representations, 2022.
Nan Rosemary Ke Aniket Didolkar, Sarthak Mittal, Anirudh Goyal, Guillaume Lajoie, Stefan Bauer, Danilo Rezende, Yoshua Bengio, Michael Mozer, and Christopher Pal. Systematic evaluation of causal discovery visual model based reinforcement learning. arXiv preprint arXiv:2107.00848, 2021.
Daniel Keysers, Nathanael Schärli, Nathan Scales, Hylke Buisman, Daniel Furrer, Sergii Kashubin, Nikola Momchev, Danila Sinopalnikov, Lukasz Stafiniak, Tibor Tihon, Dmitry Tsarkov, Xiao Wang, Marc van Zee, and Olivier Bousquet. Measuring compositional generalization: com- prehensive method on realistic data. In International Conference on Learning Representations, 2020.
Thomas Kipf. Elise Van der Pol. and Max Welling, Contrastive learning of structured world models. arXiv preprint arXiv:1911.12244 2019.
Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander C. Berg, Wan-Yen Lo, Piotr Dollár, and Ross Girshick. Segment anything. arXiv:2304.02643 2023.
Francesco Locatello, Dirk Weissenborn. Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. Object-centric learning with slot atten- tion. Advances in Neural Information Processing Systems, 13:11525-11538, 2020.

---

Published as a conference paper at ICLR 2024
Jiayuan Mao, Xiuming Zhang, Yikai Li, William T Freeman, Joshua B Tenenbaum, and Jiajun Wu. brogram-guiden image manipulators. In Proceedings of the IEEE International Conference on Computer Vision, pp.
2019.
Jiayuan Mao, Tomas Lozano-Perez, Joshua B. Tenenbaum, and Leslie Pack Kaelbing. PDSketch: Integrated Domain Programming. Learning, and Planning. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
Nihal V Nayak, Peilin Yu, and Stephen H Bach. Learning to compose soft prompts for compositional zero-shot learning. arXiv preprint arXiv:2204. 03352444440374, 2022.
Jung Yeon Park, Ondrej Biza, Linfeng Zhao, Jan Willem van de Meent, and Robin Walters. Learning symmetrie embeddings for equivariant world models. arXiv preprint arXiv:2204. 13371, 2022.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pp- 8748-8763. PMLR, 2021.
Balaraman Ravindran. An algebraic approach to abstraction in reinforcement learning. University of Massachusetts Amherst, 2004.
Ameesh Shah, Eric Zhan, Jennifer Sun, Abhinav Verma, Yisong Yue, and Swarat Chaudhuri. Learning differentiable programs with admissible neural heuristics. In Advances in Neural Infor- mation Processing Systems, 2020.
Jennifer Sun, Ann Kennedy, Eric Zhan, Yisong Yue, and Pietro Perona. Task programming: Learning data efficient behavior representations. arXiv preprint arXiv:2011.1311,, pp. 2876- 2885, 2020.
Dídac Surís, Sachit Menon, and Carl Vondrick. Vipergpt: Visual inference via python execution for reasoning. arXiv preprint arXiv:2303.08/ /28, 2023.
Hao Tang and Kevin Ellis. From perception to programs: regularize, overparameterize, and amor- tize. In International Conference on Machine Learning, pp. 33616-33631. PMLR, 2023.
Lazar Valkov, Dipak Chaudhari, Akash Srivastava, Charles Sutton, and Swarat Chaudhuri. Houdini: Lifelong learning as program synthesis. Advances in Neural Information Processing Systems, pp. 8701-8712, 2018.
Elise Van der Pol, Thomas Kipf, Frans A Oliehoek, and Max Welling. Plannable approximations to mdp homomorphisms: Equivariance under actions. arXiv preprint arXiv:2002. 2020a.
Elisc Van der Pol, Daniel Worrall, Herke van Hoof, Frans Olichock, and Max Welling. Mdp homo- morphic networks: Group symmetries in reinforcement learning. Advances in Neural Information Processing Systems, 33:419 4210, 2020b.
Rishi Veerapaneni, John D. Co-Reyes, Michael Chang, Michael Janner, Chelsea Finn, Jiajun Wu, Joshua Tenenbaum, and Sergey Levine. Entity abstraction in visual model-based reinforcement learning. In Leslie Pack Kaelbling, Danica Kragic. and Komei Sugiura (eds.). Proceedings of the Conference on Robot Learning, volume 100 of Proceedings of Machine Learning Research, pp. 1439-145555 1456. PMLR, 30 Oct-01 Nov 2020. UtU/attss.//prrceeddi.. html
.ml. press/
w100/veerapaneni222
Robin Walters, Jinxi Li, and Rose Yu. Trajectory prediction using equivariant continuous convolu- tion. arXiv preprint arXiv:2010.11344 2020.
Guanzhi Wang, Yugi Xie, Yunfan Jiang, Ajay Mandlekar, Chaowei Xiao, Yuke Zhu, Linxi Fan, and Anima Anandkumar. Voyager: An open-ended embodied agent with large language models. arXiv preprint axXiv:2005.16297 2023a.
Renhao Wang, Jiayuan Mao, Joy Hsu, Hang Zhao, Jiajun Wu, and Yang Gao. Programmatically grounded, compositionally generalizable robotic manipulation. arXiv preprint 23X4::333.13826 2023b.
12

---

Published as a conference paper at ICLR 2024
Nicholas Watters, Loic Matthey, Matko Bosnjak, Christopher P Burgess, and Alexander Lerchner. Cobra: Data-efficient model-based rl through unsupervised object discovery and curiosity-driven exploration. arXiv preprint arXiv:1905.09275, 2019a.
Nicholas Watters, Loic Matthey, Christopher P Burgess, and Alexander Lerchner. Spatial broadcast decoder: A simple architecture for learning disentangled representations in vaes. arXiv preprint 1iXiv:1901.00017, 2019b.
Ziyi Wu, Nikita Dvomik, Klaus Greff, Thomas Kipf, and Animesh Garg. Slotformer: Unsupervised visual dynamics simulation with object-centric models, 2023.
Kexin Yi, Jiajun Wu, Chuang Gan, Antonio Torralba, Pushmeet Kohli, and Josh Tenenbaum. Neural- symbolic vqa: Disentangling reasoning from vision and language understanding. Advances in neural information processing systems, 31, 2018.
Eric Zhan, Jennifer J. Sun, Ann Kennedy, Yisong Yue, and Swarat Chaudhuri. Unsupervised leam- ing of neurosymbolic encoders, CoRR, abs/2107.13132, 2021. URL https://arxiv.orr fbs/2107.13132
leam-
Linfeng Zhao, Lingzhi Kong, Robin Walters, and Lawson LS Wong. Toward compositional gener- alization in object-oriented world modeling. In International Conference on Machine Learning, pp. 26841-26864, PMLR, 2022.
13

---

Published as a conference paper at ICLR 2024
Entity Composition
Relational Composition
Generalication compoosiion shapes not togecher during trannig
Scone Compostion Objects with shared attriuttes have shared dynamics.
Nktribution Atoms L
Sampled Compesitio 10 .4)
Sampled Compostoon [o Al
"Sticky"
Shared Adjaeency
Color
Image (1
Ground Truth
Image y(++1)
[(t+
2(++3)
Actions
Figure 5: Overview of types of compositions studied. Entity composition (left) necessitates learning a world model that is equivariant to object replacement. Relational compositions (right) necessitates learning the properties of entity composition well as additional constraints where objects with shared attributes also have shared dynamiics. We study two instantiations of shared attributes sets: "Sticky" and "Team", Details on these instantiations are given Appendix A.2
A APPENDIX
The Appendix is divided into eight sections. Section A.1 erate a set-structured representation, Section A.2 surveys the types of compositions we study in detail, Section A.3 introduces faster algorithm used in implementation for module selection, Sec- tion A.4 outlines experiment details and, notably, presents justification for the Equivariant MRR metric employed study encoding separation, [
explains how leverage SAM to gen-
details of how the downstream planning experiments were conducted, Section [A.6] goes over our reasoning. for selecting relevant benchmarks. Section A.7 details ablation study with "fully-symbolic" model, and Section A.8 showcases qualitative results on a randomly sampled subset of five-object state-action pairs.
A.1 GENERATING SET-STRU REPRESENTATION. WITH SEGMENT ANYTHING
We prompt SAM (Kirillov al. 2023) with the image (/) and 64 x 3 potential masks (as there are three channels). To ensure a set structured representation, we ensure that (1) each mask captures specific property of the image, (2) collectively, all masks describe the entire image. We ensure (1) by removing duplicate masks and (2) by evaluating all combinations of remaining slots and selecting the tuple (where ie is the number of slots) that, when summed, most closely matches the image. The resulting masks {M,.... Mt are pose multiplied with the image yield (I,.... Ig]. Each masked image passed through finetuned Resnet to yield (S,.....S. Sz
8 grid of points. This yields
A.2 TYPES OF COMPOSITIONS
Entity Composition: Entity composition (Figure necessitates learning world model that is equivariant to object replacement. The dynamics of the environment depends on which objects are present in the scene.
Relational Composition: Relational composition (Figure necessitates learning all the properties present entity composition. Additionally, in relational composition, the composition is determined by constraints placed on observable attributes of individual objects. For instance, in Sticky block pushing (Fig 5). the scene constrained so that two objects start out with the same color adjacent to each other; and action on one object moves all objects of the same color with it. This gives the appearance of two objects being stuck to each other. At test time, the objects stuck together change. Sticky block pushing demonstrates compositionality constraints based on two attributes: position and color. In the team block pushing (Figure 55), we relax the adjacency constraint in the sticky block pushing domain. An action any object also moves other objects of the color. This
14

---

Published as a conference paper at ICLR 2024
allows to study whether the adjacency constraint places larger burden dynamics learning than the color constraint.
A.3 TRANSITION ALGORITHM
NPS (Goyal et al. 2021 necessitates selecting a primary slot (p) to be modified, a contextual slot (c), and a rule (r) modify the primary slot in the presence of the contextual slot. The naive algorithm to compute this tuple has a runtime of O(k-l) where is the number of slots and / the number of rules. However, in implementation, the selection of (r, p. c) be reduced to runtime of O(kl + k) by partial application of the query-key attention. This is achieved by selecting the primary slot p and MLP,, partially transforming Sp using partial transition module (L.(E. selecting the contextual slot c, and performing final transformation of (L.(e. (Sp) with S. using (i..ight)) Algorithm 2 presents this faster algorithm.
Algorithm 2 Faster Transition Algorithm. This algorithm has faster runtime than the one presented in the manuscript. The main difference is that the transition step is bifurcated into two parts, reducing the runtime of the selection from O(22t) to O(kl k) where & is the number of slots and / is the number of rules.
1: function TRANSITION(Ke><(AA.... T&J. (R Rd), (Sl.e=(S)...... Sk)) As), query= (R
2: A +- GumnbessftmaaxxXKXtteetooonnnn
3: p, <- argmax(A", axis = 'all')
4:
S* MLPr,1eIC (concat Sy. R,))
5:
A GumbelSoftmax (actiimmaxxxAAttt (QAAttention(key((A A,). query=(S" S*)
6:
C, argmax(A)) axis *al1')
7:
S MLPr,right (conghttccooactccc (concat(S., S)))
8:
return S2
A.4 EVALUATION PROCEDURE
Dataset Generation: To generate each dataset, we first create a scene configuration file that = 0
prescribes the permissible shapes and col- ors for objects within a given dataset, The scene configuration ensures that FD(Dtrain) n Fe(Devat)
Eq.MRR accounts for all permutation and correctly calculates this as closest embedding.
D. Next, we sample Dirain and Devel from the permissible scenes. Each dataset is a trajectory of state-action pairs, where state the image of the shape 3 x 224 224 and the action is vector. factor- ized by object ID (Object x North-East-Soutte West). Overall, of length 32 each where the actions are sam- pled from a random uniform distribution. Our domain is equivalent to the observed weighted shapes setup studied in Ke al. (202T) with compositionality constraint where the weight of the objects depend only on the shapes.
MRR incorrectly calculates this the closest embedding
(s
generate 3000 trajectories
2
Figure 6: Intuition for shortcomings of MRR when number slots = 2 and dslot = 1. The MRR metric incorrectly finds point (š, 9) that is + units away from (x. y) while Equivari- ant MRR considers all possible permutations and finds a point (ý, Z) that is units away from (y. x) and, in turn, closer (x, y) than (f. 9).
We compare against past works in compositional world modeling with publi- cally available codebases at the time of writ- ing. For the block pushing domain, we compare against homomorphic world models (GNN) (Zhao et al. (022) and an ablation of our model without symbols (ALIGNEDNPS), GNN uses a slot autoencoder, an action attention module and graph neural network for modeling transitions. requires two-step training process: first
Baselines:
15

---

Published as a conference paper at ICLR 2024
Dutnsct
Model
MSSE
AE-MEE
Eq,MKK 1
MSB L
AE-MISE 1
EQMMRR T
RC
Cosstos
4.238 03 +/ 1.415-04
4.WIE-M ni- 1833-04
1.3RE-01
2.55-02
4.15E-83 e/- 5.21 413
LASE-03
1.88B-13
3.6T-001
7:738-02
(Sickk)
ALIGNENNPS
N0F-04
1.736-01
12215488
5.01E-02
6:076-01
8.330-4444
2.47E-01
3,606-04
1.315-00
GNN
7,94E-03
5.78-008
5.11E:03 4.44443
6.00--04 1.00-044
6.218-03
1.26E-03
2.758-03
1.276-03
5.30E-04
+d- 5.155-05
COSMIS
4.008
2.328-3
4.33E-04
L.NE-04
1.4E 01 3.19
5,53E-05
1959
LAKE
L.1E
2,862 01 4.328
(Team)
ALKOMEONSS
1.24E
4.11
5.30
6.79E 04
1.75E-01
2ANE-02
9,48 013 1.55E
3.12E (13 6,07E
2955-001
2.022-
ONN
K998-00
3.836-03
7.16E-04
1.008.04
7.01B-03 4/- 9.811-04
LA22-40 a/- L.44--3
5.465-04 +/- 1.37B-04
PC
CORMOS
7.66R-44 +/- 4.88-844
6.348-68 +/- 2.011-055
3.90F-01
+/ 8.5--2
4.861-84-2 -/- 4.688-466
2.928-406 +/ 6.348-40
5639500 -/- 1:558-002
ALRINENPS
3.51E 05
6.30 04
2.09E 03 17. 5.89 05
2.97E 01 1/ 7.99E 02
2,4SE: 03 e/ 3,47E 04
1.22E 03 9.00E 04
3.199 D1 m/ LIOIE: 01
ONN
0.806-03
5.776-03
1.016-02
5.44508
5,505-01
5.18.00--01
1:306-02
1.136-02
1.385-02
1.085-02
5.255-01
2.678-0
Table 2: Evaluation results on the 2D block pushing domain for entity composition (EC) and rela- tional composition (RC) averaged across three seeds. This table includes standard deviation numbers as well. Our model (CosMOS) achieves best next-state reconstructions for all datasets.
warm starting the slot-autoencoder and then training the action attention model and GNN with an equivariant contrastive loss (Hungarian matching loss). ALIGNEDNPS uses a slot autoencoder for modeling perceptions and NPS module (Goyal al. (02T) for modeling transitions. The pipeline is trained end to end with contrastive loss. We use action attention with NPS as well. For both mod- els, we weren't able reproduce the results using the provided codebases due to issues in training robust perception models for large images (3 x 224 x 224).. To ensure a fair comparison, and since both these methods are agnostic to the perception model, we opt to reimplement the core ideas for both these models and use the same fine-tuned perception model for all models.
Evaluation Procedure:
We evaluate all models on single 48 GB NVIDIA A40 GPU with a (maximum possible) batch size of 64 for 50 epochs for three random seeds. Contrastive learning necessitates large batch size to ensure a diverse negative sampling set. As a result, the small batch size made contrastive learning challenging in our domain. To ensure fair comparison, we report results for all models trained using reconstruction loss. We first train the slot autoencoder (ENTITYEXTRACTI)
and SPATIALDECODER until the model shows no training improvement for 5 epochs. This is sufficient to learn slot autoencoders with near-pent sear-perfect state reconstructions. All transition models are initialized with the same slot-autoencoder and are optimized to minimize mixture of the autoencoder reconstruction loss and the next-state reconstruction loss. For composi- tional world modeling, we are interested in two aspects of model performance: next-state predictions and separation between latent states. We evaluate next-state predictions on all models using the mean squared error (MSE) between the predicted next image and the ground truth next image in the ex- perience buffer. We also measure the performance of the autoencoder on reconstructing the current state by calculating the aot-auuoencooer mean squared error (AE-MSE). Generally, training world models improves the perception model's ability to reconstruct states as well. We also evaluate the separability of the learned latent encodings. This is done by measuring the L.2 distance between the predicted next slot encodings and the ground-truth next slot encodings obtained from the encoder and using information theoretic measures such as mean reciprocal rank (MRR) to measure simi- larity. Notably, the MRR computation in previous work does not to account for the non-canonical order of slots, causing higher I.2 distance and, consecutively, higher MRR scores when the target and predicted slots have different orderings. The core issue here that MRR computation, as used in previous works, fixed the order of the slots before calculating L2 distance. This ignored k! - possible orderings where a closer target encoding could be found. To rectify this, we propose new metric, Equivariant MRR (ER.MRR), which uses the minimum L2 distance among all permu- tations of slot encodings to calculate mean reciprocal rank. This metric ensures that the latent slot encodings are not penalized for having different slot orders. Figure presents an illustration of the shortcomings of MRR on simple example. This limitation is characteristic of algorithms which do not align the slots to canonical slot ordering. In practice, we observe that the Equivariant MRR is always lower than or equal to the MRR.
A.5 DOWNSTREAM EVALUATION SETUP
Following Veerapaneni et al. 2020 we use greedy planner that chooses the action that minimizes the Hungarian distance between the current and the goal state. These actions are applied t times over trajectory of length with the output from the world model at the (d 1)-th step becoming the state for step d in the trajectory. Due to this compounding nature, we see an increased divergence from the ground truth get deeper into the trajectory. At each step d in the trajectory,
16

---

Published as a conference paper at ICLR 2024
Dataset
Only Symbols (MSE
Only Neural AlignedNP'S (MSE )) COSMOS (MSE J)
Object RC- Sticky
1.36E-02
1. 1.14E-02
4.23E.-03
Objecc RC- Team
1.39F-02
1.24F-02
4.60E:033
3 Object EC
1.21E-02
3.51F-03
7.66E:004
Table 3: Evaluation results on the 2D block pushing domain for ablations of COSMOS.
the accuracy of the world model is evaluated as the LI error of the difference between the current ground truth and predicted states in the form of their zy-coordinates. These zyycoordinatess initialized for each object to the origin and updated with every action taken by the corresponding rule. For example, after one step, if the ground truth moves object the east, but the planner chooses to move the same object to the west, then the distance between the two states would be 2. We these experiments for the 500 trajectories of length /. 32 in test dataset and average the scores at each trajectory depth. We showcase results in Figure [4] Our model (COSMOS) shows the most consistency and least deviation from the goal state in all datasets, which suggests that neurosymbolic grounding helps improve the downstream efficacy of world models.
A.6 DATASET COMPARISON
The focus our paper is to demonstrate the first neuro-symbolic framework leveraging foundation models for compositional object-oriente world modelling, and evaluate on the same benchmarks as existing work (Kipf et al. 2019 Goyal et al. 2021 Zhao et al. 2022). We chose our evaluation domain based four properties: (1) object-oriented state and action space, (2) history-invariant dy- namics, (3) action conditioned (plannable) dynamics, and (4) ease of generating new configurations (to evaluate entity and relational composition). We curate the following list of domains from related
work to explain what properties are missing for each dataset.
Dataset and Relevant Works
Object Oriented State and Action Space
History Independent
Action Conditioned (Plannable)
Configurable
2D Block Pushing
Kipf et al.
2019
Goyal al. 2021
(7hao al. 2022
Ke clal 2021
Pnysion WD
Phyre Wu cl a.
2014
Clevrer WW at
3DObj Wi
21A3
MineRL (Famar er 2023
Atari (Polt//pacee wtt.eesss/oooeaaaeeoooo Goyal et 2021
Minigrid/Baby-A1 Mao et al. 2022
A.7 SYMBOLIC ABLATION OF COSMOS
ALIGNEDNPS serves as a "fully neural" ablation to demonstrate the effectiveness our model. In this section, we detail another "fully symbolic" ablation of our model to demonstrate the need for a neurosymbolic approach. Specifically, we maintain the algorithm presented in [ but modify the transition model to the symbolic embedding to predict the next state. Specifically, line changes to:
S MLPBANK///(conea((/ A.t R.))
The results, detailed in Table indicate that the "symbols-only" model significantly underperforms compared to COSMOS. We believe this is because the symbolic embedding is constructed by con- catenating symbolic attributes, and the rule module not aware of this structure. This causes the MLP to overfit to the attribute compositions seen at train time. COSMOS sidesteps this issue by using the symbolic embedding in the key-query attention module select the relevant rule module, while allowing the real vector to learn local features useful for modeling action conditionnd transitions.
A.8 QUALITATIVE RESULTS
17

---

Published as a conference paper at ICLR 2024
Higure 7: Qualitative outputs on randomly chosen state-action pairs for all baselines. We show two samples for each experiment and dataset type with objects. Color is shown for illustrative purposes only; in implementation. the action conditioning does not carry any information about color.
Predicted Next Image
Ground Truth Next Image
Image
Predicted Next Image
Relational Composition (Sticky Shapeworld)
Ground Truth Next Image
Image
Entity Composition
Als
AlignedNPS
COSMOS
COSMOS
GNN
GNN
18


Reviews:
[{"review_id": "XI5kSWc2v4", "reviewer": "Reviewer_5cti", "summary": "This paper proposed COSMOS framework for compositional generalization in world modeling, that involves a blend of neural and symbolic representations to understand scene entities and their interactions.", "questions": "point 2 from weakness section", "limitations": "", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1.     This paper proposes an end-to-end differentiable framework with a novel neuro-symbolic attention\n2.\tThis is a well written and well-structured paper.\n3.\tMost of the traditional neuro-symbolic methods map the representations to symbol manually while this paper does it without any manual effort", "weaknesses": "1.\tThe effectiveness of the framework might be constrained dealing with larger and more complex real-world scenarios. Since scalability of neuro-symbolic methods are required to handle more diverse environments.\n2.\tWhen the model will encounter with noisy or incomplete input, how will the model perform?\n3.\tCombining neural and symbolic inputs might be computationally heavy, but there is no significant discussion about the computational complexity.\n4.\tAlthough the model showcases strong performance in the 2D block pushing domain with MSE but that’s not the case for MRR and Eq.MRR always. More experimental results are required to establish this as the new state-of-the-art.\n\n5. There are typos in introduction section.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OM51e16qWb", "reviewer": "Reviewer_CQ89", "summary": "This paper presents a neurosymbolic approach to learning a transition model in the pixel space. The proposed model, namely Cosmos, takes in the current state as an image, the action (encoded as one-hot vectors), and predicts the next state. The model is trained on a dataset of state-action-state transition tuples and evaluated on unseen state-action combinations. The model focuses on two types of compositional generalization tests: scene composition and object composition.", "questions": "I don't have particular questions. Please address the missing related works; and consider reframing and better illustrating the object compositional generalization.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "This paper tackles an important problem that is interesting to the ICLR community, specifically learning world models from pixels. The overall presentation of the paper is good. The organization of the method section clearly illustrates the number of modules in the system and how they are connected with each other, which is obviously helpful for readers. The description of experimental setups is clear, and the authors have done a sufficient number of ablation studies and comparisons with baselines.", "weaknesses": "There are two main weaknesses of the paper.\n\nFirst, the problem setting of object composition seems very contrived to me, for two reasons.\n- In the physical world, it is unclear what's a concrete example where such kind of metarules would apply. In particular, the authors are training the model on seeing two red blocks moving together and two green blocks moving together, and hope that the model would generalize to predict two blue blocks would also move together. Such kind of \"attribute-relationship\" based generalization doesn't seem natural to me. Arguably, this kind of generalization can be dangerous: two blocks can be stacked together; two cylinders can be stacked together, but not two spheres, in the physical world.\n- There is some serious machine learning identifiability issue with this setting. If the model does not have inductive biases in training, there is no way that it can generalize.\n--- Based on these two concerns, the arguments around object compositional generalization is weak.\n\nSecond, the model is only trained and evaluated on a fairly toy environment, and the downstream application to planning is only shown in a very simple setting. It is unclear how this approach can be generalized to more complicated scenarios.\n\nSlightly minor is that the paper missed some important related work along the direction of learning neuro-symbolic transition models. For example,\n- PDSketch: Integrated Planning Domain Programming and Learning https://arxiv.org/abs/2303.05501\n- Learning Object-Oriented Dynamics for Planning from Text https://openreview.net/forum?id=B6EIcyp-Rb7\n\nThey are not exactly the same setting but a lot of the high-level ideas are definitely the same, including learning lifted transition rules, using factorized embeddings (colors, shapes, etc.) to represent objects.\n\nWhile I overall like the presentation of the paper---it's well-organized and overall good. I found the description of some details very unclear. In particular:\n- Page 3, the object composition part. I have to check the appendix and read through paragraphs/figures several times in order to understand what this object composition means. I think the name is not very descriptive. The authors should consider change it to a better name that describes such kind of \"metarules\" (e.g., two objects have relations if they share the same color) and present concrete examples in the main text.\n- Parge 4-5: the authors should keep the \"...\" in the sets. Otherwise it's very confusing to look at \"{c1, c2, cp}\"\n- The writing of the method section could be further improved by having a running example (and referring back to this example in 3.1, 3.2, and 3.3).\n\nFinally, the paper does not have a limitation discussion section.\n\nMinor notes on the writing: I think using CG as an acronym for \"compositional generalization\" is a bit uncommon. The term is easily confused with other concepts like \"computer graphics.\"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "B7RJbq2egb", "reviewer": "Reviewer_QCEm", "summary": "The authors proposed a framework for object-centric world modeling in compositional generalization settings. They propose a neuro-symbolic scene encoding, consisting of real vectors and vectors of symbols describing attributes, as well as a neuro-symbolic attention mechanism, which binds entities to rules of interaction. They use foundation models to extract each entity’s symbolic attributes, and show their method’s performance on a block-pushing domain.", "questions": "Q1. Can you clarify how SAM is used with ResNet to produce a set-structured hidden representation for each object? \n\nQ2. Is there a way to interpret the learned rules and qualitatively see how well it aligns with the ground truth rules?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "I am appreciative of the neural-symbolic attention mechanism, computed between ordered symbolic and neural rule encodings to determine the most applicable rule-slot pair. I also appreciate that this is end-to-end differentiable as a permutation equivariant action attention module. Using a frozen foundation model to capture the symbolic attributes in the scene encodings is also a simple but elegant way of decomposing entities into attributes.", "weaknesses": "W1. I am not convinced that the proposed neuro-symbolic scene encoding is the optimal formulation. I would have appreciated ablations in this paper, including ones that explored whether or not having both real vectors and a vector of symbolic attributes is actually helpful. What should the real vector capture that the symbolic attribute is not capturing, especially given that the rules of the evaluation domain seem to only rely on these attributes?\n\nW2. An additional note is that I would have appreciated more clarity of what information is given to the model, and whether other methods have access to the same information. From my understanding, a strong assumption is that the symbolic labeling module requires a predefined list of attributes that is important for the downstream task. I think in many downstream tasks this may not be reasonable to know a priori. Potentially, you can run experiments showing that with a superset of attributes, directly predicted by some foundation model, that you can still learn the correct correspondence to rules given this noise. \n\nW3. Similarly, are there assumptions made on how many rules there are in the evaluation domain? I understand that the rule is a learnable encoding, but is the amount of rules learned as well? One can imagine that the method discovers rules that are correct, but not optimal (e.g., decompose rules into many smaller rules that overfit to the train set). \n\nW4. I would have appreciated evaluation on a different domain, such as maybe Physion, and learn more complex and less obvious rules such as rigid and soft-body collisions, stable multi-object configurations, etc. In the block-pushing domain, it seems like the rules are tied to clear attributes such as shape and color, without having to further learn whether non-uniform combinations of these would lead to certain downstream effects. Also related to the known attributes assumption in W2.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)