Learning Dynamic Attribute-factored World Models for Efficient Multi-object Reinforcement Learning
Authors: Fan Feng
Sara Magliacane
Venue: NeurIPS 2023 poster (2023)
Keywords: multi-object RL
compositional generalization
factored representations
TL;DR: Learning factored representations at the attribute level, including for interactions across objects, helps to learn a policy that generalizes across unseen objects with varying attributes and latent parameters.

Abstract: In many reinforcement learning tasks, the agent has to learn to interact with many objects of different types and generalize to unseen combinations and numbers of objects. Often a task is a composition of previously learned tasks (e.g. block stacking).
These are examples of compositional generalization, in which we compose object-centric representations to solve complex tasks. Recent works have shown the benefits of object-factored representations and hierarchical abstractions for improving sample efficiency in these settings. On the other hand, these methods do not fully exploit the benefits of factorization in terms of object attributes. In this paper, we address this opportunity and introduce the Dynamic Attribute FacTored RL (DAFT-RL) framework. 
In DAFT-RL, we leverage object-centric representation learning to extract objects from visual inputs. We learn to classify them into classes and infer their latent parameters. For each class of object, we learn a class template graph that describes how the dynamics and reward of an object of this class factorize according to its attributes. We also learn an interaction pattern graph that describes how objects of different classes interact with each other at the attribute level. Through these graphs and a dynamic interaction graph that models the interactions between objects, we can learn a policy that can then be directly applied in a new environment by estimating the interactions and latent parameters.
We evaluate DAFT-RL in three benchmark datasets and show our framework outperforms the state-of-the-art in generalizing across unseen objects with varying attributes and latent parameters, as well as in the composition of previously learned tasks.

Learning Dynamic Attribute-factrre World Models for Efficient Multi-object Reinforcement Learning
Fan Feng City University of Hong Kong @eengl017@gmanl com
Sara Magliacane University of Amsterdam MIT-IBM Watson Al Lab sara . mag iacaeeallllaeeeaaaaaaam com
Abstract
In many reinforcement learning tasks, the agent has to learn to interact with many objects of different types and generalize to unseen combinations and numbers of objects. Often a task is composition of previously learned tasks (e.g. block stacking). These are examples compositional generalization, in which we compose object-centric representations solve complex tasks. Recent works have shown the benefits of object-factored representations and hierarchical abstractions for improving sample efficiency in these settings. On the other hand, these methods do not fully exploit the benefits of factorization in terms of object attributes. In this paper, we address this opportunity and introduce the Dynamic Attribute FacTored RL (DAFT-RL) framework. In DAFT-RL, we leverage object-centric representation learning extract objects from visual inputs. We learn to classify them into classes and infer their latent parameters. For each class of object, we learn class template graph that describes how the dynamics and reward of an object of this class factorize according to its attributes. We also learn an interaction pattern graph that describes how objects of different classes interact with each other at the attribute level. Through these graphs and dynamic interaction graph that models the interactions between objects, we can learn a policy that can then be directly applied in a new environment by estimating the interactions and latent parameters. We evaluate DAFT-RL in three benchmark datasets and show our framework outperforms the wate-of-the-an in generalizing across unseen objects with varying attributes and latent parameters, as well as in the composition of previously learned tasks.
Introduction
Model-based reinforcement learning (MBRL) and world models [1-3] have demonstrated improved performance in many RI. tasks by providing better sample efficiency. However, most world models focus only on modeling a single object or holistic modeling over the environment, while in real- world tasks, we often have environments with multiple objects that interact, and we are interested in generalizing to unseen combinations and numbers objects. In recent years, there have been several studies exploring and learning object-oriented environment models policy models [4-17] and tackling the problem of compositional or combinatorial generalization, which we consider combining the modeling of multiple objects or tasks to solve new task.
Although these methods have effectively leveraged object-centric and object-factored representations in RI. and thus improved the sample efficiency in multi-object settings, they have not fully exploited the benefits of factorization in terms of object attributes. Often an object's transition and reward functions are influenced only by sparse subset of attributes, e.g. object's position and reward are affected by its previous position, but not by its appearance or activation state once the object of interest is identified (e.g.. "the red block"). In these environments, interactions between objects are
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

often sparse, both in time and in terms of which attributes are affected, e.g. the position of a box affected by the position of another box at the timestep in which they collide, but not directly by the other object's friction coefficient. Additionally, objects of the same type share similar factored dynamies, modulated by an object-specific latent parameter, while objects of different types might instead have different attributes, dynamics, and rewards.
In this paper, we propose Dynamic Attribute Fore RL (DAFT-RL), a framework that learns fine-grained attribute-factored representation across objects, including a dynamic graph for modeling interactions between objects. As part of this framework, we propose a model, DAFT-MDP, that builds on Factored (PO)MDPs [1 [18-23], Relational MDPs [24-26] and especially Object- jijiet--riented (PO)MDPs [27-29], but focuses on a more fine-grained factorization at the attribute level and dynamic graphs. We implement our framework as a structured and sequential generative model by combining it with state-of-the-an object-centric representation learning [30, 31] for extracting objects and attributes from visual inputs, factored adaptation approaches inspired by the causality literature [32, 33] for estimating the factored dynamics and reward, soft attention networks [34] for action binding [35] and (dynamic) Neural Relational Inference [36, 37] for modeling interactions.
Our framework allows us learn a policy on set of source environments that can successfully generalize to new environments with unseen combinations of objects with different latent parameters (possibly with unseen values) and types, as well as to combinations of previously learned tasks on different objects, without any further policy learning. We show the benefits of DAFT-RL in three benchmark datasets for compositional generalization, where outperforms the baselines.
2 Dynamic Attribute-FacToree MDPs (DAFT-MDP)
We formalize our assumptions by introducing our DAFT-MDP model, which is an extension with class template graphs, interaction patterns, and interaction graphs of factored (PO)MDPs [18-23]. This extension takes inspiration from Relational MDPs [24-26] and their literature, especially Object- Oriented (PO)MDPs [27-29], but we propose more fine fine-grained factorization of the transition and reward at the object attribute level, based on estimating template and dynamic interaction graphs.
Intuitively, we will consider learning policy that can be generalized across different environments that vary in the number and characteristics of their objects. We will assume that each environment is composed of multiple objects, each of a specific type, or class. Each object has several observable attributes (e.g. position, velocity) and some latent constant parameters (e.g. an object-specific friction coefficient). Objects of the same class will have the same set of attributes, the same transition and reward functions, but can differ in the values of the attributes (c.g. they are at different positions) and in the value of the latent parameters (e.g. they have different friction coefficients). We will assume that the transition and reward functions can be factored in terms of attributes and that for given attribute only a sparse subset of other attributes influences these functions.
In our setting, objects can interact with each other, which might inffuence their dynamics. We will assume that these interactions are sparse, both in time and in terms of the effect on the attributes of each object, and that all objects in a class have the same interaction pattern in terms of how the attributes interact with objects of another class. In each environment, we will assume that an action has only an effect on one object at a time.
We formalize these assumptions in the following. We start by defining our class system, and then we describe three types of graphs (class template graphs, interaction patterns, and dynamic interaction graphs in Fig. 1) that describe how dynamics and reward factorize based on the classes, attributes, and interactions, and finally provide formal definition of Dynamic Attribute-Fach
MDP.
Class system, attributes, and objects. We assume a known set of classes C = (CI,.... C of objects, where each class C, describes a set of observable attributes [Cy.s1..... Cy.s,,, which we assume for simplicity are the same number in each class. We assume that each class has a set of latent constant parameters C,.0, that represent physical properties of the object that can vary across different objects of the same type. For example, in one of the benchmarks [15], we consider two types of objects C [BOX, sw). boxes and switches, with class attributes (BOx.cole .ollor, pox. Box,vel, BOX.0) representing the object color, position, velocity, and friction coefficient of box, and [sw.color, sw. ww.p swaactive), representing the color, position, and ac- tivation state of switch. The class system specifies template for set of environments that can
2

---

A. Class Template Graphs
Interaction Pattern Graph
Interaction Graph
GLT
:
Z
S
C
C,.9
A
Figure I: The graphical representation of DAFT-MDP. The colors denote the attributes for an object or class, the red dashed lines denote edges that can be switched on or off at different timesteps.
vary in the number and characteristies of the objects. Each environment has a fixed set of objects
O = foi..... 0m ). where each object is an instance of a class in C. which we denote as o₂ E C, for i E fl....,m} m).j E {1....,k) k}. We denote the class of object og with C(t). For example, using the class system from the previous example, we can represent source environment as O' = {0 02) where 01 E BOX and 02 E sw, and a target environment
O" {or, 02, 03), where we add 03 E BOX to the original objects. For each object oR E Cj and timestep 1 = {1...., T,). we denote its attributes at time
of (o..st...... 04-5%,,, ). which (C..S1,..., Cj.snh, and its constant parameters as o,.0. D. In our example, for box or the attributes 01.84 0g.sh and 0 are its position and velocity at time and its friction coefficient, while for the switch 0g the attributes 02.s and 02.sh are its position and activation at time t.
instantiations of the class attributes
States, actions, transitions, and rewards. We define the state at time t as the collection of all the object states, i.e., s² = {of.... .o....... I with domain S. We collect all object-specific latent parameters in global parameter 0 = (o.. 0m.0) with domain θ. We define the space of actions A and use a' A to denote the action time t. We denote the transition probability as p(s sst a', 0). The reward at time t is denoted as r' and the reward probability is denoted as p(r'ls* a', O). The transition and reward probabilities are factorized according to the class template graphs, interaction pattern graphs, and interaction graphs that we can learn from data, and that we will introduce below.
Class template graphs (Fig, 1A). We assume that all objects within the same class share the same factorization in terms of how their own attributes, latent parameters, and actions influence their dynamics and rewards. For example, all boxes will have a similar relationship between their position, velocity, and latent friction parameters, since they follow the same physical laws. Similarly, all switches will share similar dynamics, which will be different from the boxes. We describe these relationships as Dynamic Bayesian Networks (DBNs) [38], which are graphical models that describe the template for the relations between two contiguous timesteps, t and t + 1, and assume this structure is time-invariant. In particular, for each class C,. we learn a DBN Gc, over the nodes .Cy.si....,C, Cy.sf+1 ..... Cy. Cj 0, a², r'r" that represents a template for the instance graph between the attributes for each object of this class G5 for o, C.. In particular, the edges between the action a' and the attributes timestep + 1 can be switched on or switched off at different timesteps, since as we will see in the description of the interaction graphs, they represent the interaction pattern of the agent with any specific object.
We show an example of two DBNs for two classes of objects (boxes and switches) in Fig. la. In this example, Cj.sf influences C,-st+1 and Ci.s but it's not influenced by Cy.st The reward rd is only influenced by Ci.sl and the action will only have an effect on Cl.s1 C2.sh influences only C2.s and y.h
f+1 and Ci.a5+1 Moreover,
Interaction pattern graphs (Fig. IB). When two objects interact, it often happens that only some of the attributes of each object affect the other. For example, when two boxes collide, their positions and velocities might change, but their masses will not. Therefore, we will assume that from attribute perspective, the interplay between objects during interaction is also factored
3

---

and sparse. Additionally, we assume that the interactions between two objects follow patterns based on the classes of the objects. For example, the attribute of box will always interact in the same way with the attributes of a switch, regardless of the specific object. In particular, for any pair of classes C, and Cj (possibly also with = =}), we learn a DBN Gc,c, over the nodes [C J= C,O0, R- 1 [C (ft.s++1a Ji Cj.03. We show an example of DBN describing how boxes interact with other boxes in Fig. lb. In this case, the interaction is limited to c1.s from one object influencing the cg.s and C1.52 f+1 from the other object, and the latent parameters 0 from each object influencing C1.s² 1+1 from the other object. While the interaction patterns are invariann,,, we assume that for each pair of objects, the interactions are only switched on at some points in time, we will now describe through the interaction graphs.
Dynamic interaction graph (Fig. IC). The class template and interaction pattern graphs that we just described can model the general behavior of the classes of objects in static, time-invariant way. On the other hand, in many multi-object environments object interactions occur sparsely, the pairs of interacted objects not fixed, and the action has effect only objects at any given time (e.g. we will assume only one for simplicity). We, therefore, propose to model these interactions between objects as a dynamic graph Ginter = (Gimer) y at the object level,
limited number of
which is a sequence of graphs Pinter with edges from a subset of {oi..... on a²} to a subset of {oft1 ...... , ". Each edge of + -0 ofth represents an interaction between an object o, and an object 0z. This interaction is instantiated in an instance interaction graph Co..o, following the interaction pattern graph Gc,cc, for each pair of objects E C; and 0j Cj-
We also learn an object selector a' E fl,....,m m) that selects the single object on which the action a' has an effect at timestep I, which we represent with an edge from a' the selected object of, in Pinter In particular, this selection (or action binding) activates the edges from a' to the object attributes described in the class template graph Gc, in the instantiated version for object o, E Cj, which we denote Co The graph Gitter is dynamic because the edges of each graph in the sequence can change at each timestep t I...., T.
We show an example of interactions between three objects different timesteps in Fig. 1C. this example, at timestep / there is an interaction between and 03, which will follow the interaction pattern presented in Fig. IB, i.c. 01.s} + {og.s ++1 .3..5 ++1] and 01 o,.0º - 03.s ++1 and viceversa for 03 towards 0₁. Moreover, af has an effect on 02, and specifically 02.8/ 1-1 following the class template graph presented in Fig. 1A. Instead, at timestep t + 1, there are no interactions between the objects and +++1 has effect 01, specifically on 01.8/+2 and o1 01 -4/2 For completeness, in App. B we provide an example of this graph combined with the instantiated class template graphs for each object Go: g and Gos for three timesteps t, f + 1, 2, well as the instantiated interaction pattern graph GoL.os that is switched on at timestep t.
Our modeling assumptions. Now that we have introduced all of the graphical structures that we will need, we can describe our assumptions as a Dynamic Attribute- FacToored Markov Decision Process (DAFT-MDP). We will assume that the class system is fixed and that the objects vary in each environment, as can their interaction graphs Under this assumption, a DAFT-MDP defines a family of MDPs, that parametrized in the objects O and their dynamic interaction graph Ginter-
Definition 1 (DAFT-MDP). A Dynamic Attribute-Fac is a tuple (C, A, R. where is the set of classes, is the set of objects, θ is the space of the constant latent parameters, A the action space, P, is the transition distribution, R is the reward function and G is a set of graphs that contains the collection of class template graphs for each class {9c, Te,ee, the collection of interaction pattern graphs for each pair of classes (9c,,,, ,,,,,, and the dynamic interaction graph Ginten, as defined previously. These graphs define the factorization of the transition distribution per object and per attribute, as follows:
Markov Decision Process (DAFT-MDP)
2=(8*+118',0,00 = IIII P. (or.s)* (o..s(+2), (pag: Soioo =1
(o (o0 -o.EE
where s² is the collection of all attributes of all objects at time t, O E O is the collection of all latent constant parameters for all objects, a' E A is the action. In the right-hand term, or.s] t+1 is altribute 81 of object or at time + 1, while pag: (0n.8/+1) ) are the parents of the altribute 1 for

---

object og based on the class template graph Gatey where (()) is the class of and where the action binding a' activates any potential connections from a'. In the second term of the conditioning, we iterate over the objects Oj that are interacting with o; at time in the dynamic interaction graph Og - O, E Tutter For each of these objects oß we collect the attributes that interact with 0g.81 in the instance interaction pattern So,.o based on interaction pattern graph (C()) (C()) for the respective classes C(i) and Ckk). Similarly, we define the factorization of the reward function per object and per attribute Ris', a(s".8r a², 0) = RR(ppg,, (r*) (r')),,E)) where for each object og we collect all the attributes that have edge the reward in the instantiation of the class template graph.
In the following, we assume that the classes C are known and fixed across environments, while the objects O can vary, as can the latent parameters 0. In the training phase, we will learn how to classify objects, the transition and reward functions based on the class template graphs {9c, Jc, Ec and the interaction patterns {9c,.c CCC,,,,,,CCCc In the testing phase, will infer the class and latent parameters of each object, as well as the interactions between the objects in the dynamic interaction graph Ginter which specify the transition and reward functions in the new environment.
3 The DAFT-RL Framework
In the previous section, introduced model, DAFT-MDP (Fig. 2). In this section we provide a framework for estimating DAFT-MDPs, leveraging them for policy learning in a set source environments, and adapting the policy new target environment with different objects, without any additional policy learning. Our framework is divided in four steps: (i) offline class learning in single-object environments, (ii) offline interaction learning and latent parameter inference in multi- object environments, (iii) policy learning and imagination in multi-object environments, and finally (iv) adaptation new multi-object environment. We present each step and its implementation in the following.
Step 1: Class learning single-object environments
Stage 3: Policy Learning
(1 of Paa
10 dg P
(00,
Object-centric model
Object-canting model
(dr, ofy, Ma
1f
[0), d', o', of PM)
Object class Inference
C.C.G
imagined trajectories
((6 di. it, 6. PD1
Learning cless templete graph
Ee, Sey &c,
Learned policy
s"taal x., Bend
Stage 2: Interaction Learning and Latents Inference
Stage Adaptation
10.. Objoct-contric contric model
160. PI
Object-cantre model
0%) of, d'rr'))
[o], ol, 0/, of, 1/)
Learn the object solector
a' E inter
Interred constant latent paramoters
0
Learn object interaction graph
object inertaction graph Apply pollcy
inter
Learn interaction pattern graphs
8 cc-c, (1.23). (1.23)
(r $.0. me))
Learn the Istent constant parameterss
0,.0.0 0,.00.0..0.. 0.
Figure 2: The learning pipelines of the DAFT-RL framework.
In all steps and all environments, if the input is an image, we extract the objects and their attributes {of, o², ... onn) from sequences of images x² with pre-trained object-centric methods, e.g. SA [30] and AIR [31]. For symbolic inputs, we directly access the attributes of all objects s' = {of of.... of J. For each set of objects O, we learn to classify the objects into their classes C with supervised learning, which we describe in detail in App. C.
5

---

3.1 Step 1: Class Learning in Single-Object Environments
In this step we consider data from m single-object envirouments with different objects and no agent interaction. In particular, for each class Cj we collect the transitions for several objects o, Cj as {or, af. ri in environments in which there is only object o,, and a random policy is used to generate actions. We denote these data as Dsingle- = [fo], a r/] |Ee,, veji,,jjil,....... k,i m
We initiate the class template Gc; for each class C randomly, and then use Deingle to learn it, except the contribution of the latent parameters, which is learned in the next phase. In particular, we learn the class template graph by maximizing the log-likelihood of log PA (of , of. a², gC)), where A = {^g, } } are the parameters of dynamics and reward models. For the implementation, we use Gated Recurrent Units (GRU) [39] to learn the dynamics and reward models jointly with the class template graphs. At time step for each object with class Cj. the inputs to the GRU are {page, (o) )) and {page, (r:)) and the GRU outputs of and rf. The learning objective of this step given below, where we maximize the log-likelihoods of dynamics and reward models and regularize the graph to be sparse:
T ΣΣΣ
(logpx(o.s(-1.
(-11, .r pageo (o,.s(+1) ( pagcroo (r9)
- lI
argmax,,, X. (G
where and ls indicate the number of single-object environments and object type classes, respectively, and pagcos denotes the parents of a variable in the template class graph for the class C(i) of object
og- After this step, fix the learned (gc: ge,..__..... ...., (Ck) } with the exception of the edges from the latent parameters 0, which here we assume are disabled and we will learn in the next step. In later stages, we will reuse the learned reward model Ap and the class template graphs.
3.2 Step 2: Interaction Learning and Latent Parameter Inference in Multi-object Environments
In this step, consider data Dimalti from N multi-object environments with different object con- figurations and in which the objects can have varying latent parameters. Formally, we define Dimuti {fof, al, ri, o6, ab, r². -... ofn a r 17" NN m-= In each of these environments sume the agent can interact only with one object at a time. On these data, we again extract the objects and their attributes from sequence of images with pretrained object-centric methods and classify the objects using the object classifier. We use these data to learn the interaction pattern graphs Gc.,c, for each pair of classes C, and Cj and the dynamic interaction graph Ginter by exploiting the previously learned class template graphs. In particular, we first learn the action binding a', and at second stage, we jointly learn the rest of the dynamic interaction graph Ginters the interaction patterns Gc, for each pair of classes, the sppect-pe latent parameters o,0 and their edges to the other attributes in Gc,- We describe these two stages in detail in the following.
as-
3.2.1 Step 2.1: Learning the Action Binding
Motivated by [35], we learn the dynamic action binding = {a a²... ..) using soft attention networks. which are modified by the single-head self-attention module in the Transformer model [34]. Specifically, we perform non-linear transformations on the states and actions using multi-layer perceptrons (MLPs) derive the key k' (f& (of). fe(ot)...... (o query q' = f, (a'), and
value v (of). (o)( ..... (o )) respectively. We then compute the attention weights af = softmax ((k) q", (k))* .... ) )')). We use the learned attention weights a as the action binding selector, as it provides an estimation of the binding affinity from the action to the objects at each time step. The soft attention mechanism assigns weights from the action to the objects by multiplying the value vector of with the attention weights a' and then embeds the weighted actions into the dynamics of each object. We maintain a fixed structure for the class template graphs and focus on learning the action binding selector by updating fk. fq. and f..
3.2.2 Step 2.2: Learning Dynamic Interaction Graphs
As noted in Section 2, the interaction among objects may change over time and usually occurs in a sparse manner. To learn this dynamic graph, we leverage sequential latent variable model to infer

---

the object interaction graph. Following the neural relational inference (NRI) works [36, 37], we use an encoder to generate the latent variables and subsequently sample the interaction graph based on these variables. Specifically, we use graph neural networks (GNN) as the encoder module, where the nodes represent the states of each object, and the predicted edges denote the temporal interactions between the objects. In line with the dynamic NRI (dNRI) [37], we use recurrent units to model the temporal relations of the interaction graph. We outline the key components of the model in the following, and provide the detailed description in App. C.
Encoder and prior During training, at each time step t, we use GNN layers to generate hidden embeddings h' = GNN (oj 02 ................................ ). ). For each object pair o; and 0g, we then obtain h(.)). For the encoder, we use a Gated Recurrent Unit (GRU) [39] to model the temporal dependency of the interaction graphs. The inputs to the GRU include the future embeddings of the states, which then generate the updated embeddings:
,au GRU (h (.)) h ). For the prior, we also use (i.j):
(i.)).
a GRU, but in this context, we append ((g)) with the output of the GRU from the previous step h-1 (rii. as the input. In specific terms, we have h (i.j.prrior = GRUp (h h (a.).prior) For both the encoder and the prior, we feed the output embedding from the GRUs to an MLP layer to derive the distribution of the latent variables Is (z' | where s2:2² {or", ;.... Om ) ](=1 and prior distribution pa(2" sli =11t-1), respectively. We assume that the encoder outputs Bernoulli distribution for each edge and the graph Fimer is sampled using the Gumbel-
trick [40].
Decoder We perform one-step prediction to generate 6, for each object o; predicting the state dynamics with the learned graphs, including Gitter, as well as the interaction pattern graph 9c,, for two objects with class C, and Cy. At the same time, our goal is also learn the interaction pattern graph with sparsity regularization. We also learn the latent parameters, O at this stage. Specifically, we also incorporate 0 and the graph from Gc, 0.o; for each object o with class C₁ into the dynamies model and use the sparsity regularization for GC,00...
Therefore, the learning objectives include maximizing the likelihood of the dynamics model, min- imizing the KL divergence between Pe and qs to estimate 0, and encouraging the sparsity of the interaction pattern graph Gc..c, and the subgraph from latent parameters to states Ic.
T
m
n
argmax,
lo
( (0,.s) 1+1 | paay(or.s) ( {pac; (0 ( 1o
1=1
k |9 0. 111
T KL (go (z' ( lps (z" sltt 211, 1))-1)))
|90.0, 11 - ΣΣ
Σ 1=2
i=
i=1==1
where 9 is the dynamic interaction graph Ginter, interaction pattern graphs (Go,o, 1 li {1, .......... k},j E e {1,2, m).j E (1.2...., k)}. Similarly, to Definition 1, pag; (o..s)*) ( indicates the par- ents of the attribute 8 for object o based on the class template graph Go(s). where C(i) is the class of o4, and where the action binding a' activates deactivates any potential connections from a'. In the second term of the conditioning, we iterate over the objects o% that are interacting with o, at time in the dynamic interaction graph 0 - 0 T For each of these objects og we collect the attributes that interact with 0,.8t in the instance interaction pattern G5 based on interaction pattern graph (C()).C()) for the respective classes C(i) and Clk). A₂ and ø indicate the dynamics model and encoder parameters, while v, k q are the parameters of MLPs for learning the attention models. After this, we have learned all the graphs, dynamics, and reward models that estimate DAFT-MDP.
k}} and the subgraph from latent parameters to states {9c, 0.o.
3.3 Step 3: Policy Learning and Imagination in Multi-object Environments
In the first two phases, we have learned the template for our world model, which we can now finetune to new multi-object domains by inferring the environment-spectin latent parameters 0 and the interaction graph Ginter- We again consider several multi-object environments with different object configurations and in which the objects can have varying latent parameters. For each environment, we then use the finetuned environment--pecife world model to create set of imagined trajectories. Finally, we can learn a policy x (aa|ss, f, Ginter across different environments, based on the real and imagined trajectories. We can apply policy learning or planning methods using any RL algorithms. To
7

---

take full advantage of the estimated models, we use RL or planning methods such as model predictive control (MPC) [41] and proximal policy optimization (PPO) [42] to learn x². Detailed information about the domain parameters are provided in App. C.
3.4 Step 4: Adaptation to the New multi-Object Environment
In new environment, we apply the policy x (a'|s* 0, Iinter) by inferring latent parameters 0 and dynamic interaction graphs Ginter based on a few trajectories, without any policy learning.
4 Related Work
We shortly summarize the related work in this section and provide more detailed discussion of related work, including the discussion of each method and a comparison based on the method features in App. A. Recent work in object-oriented and relational RL has incorporated various inductive biases for modeling object relations into both model-based based and model-free RL frameworks. Zhou et al. [15] investigate the benefits of deep learning architectures, such as MLP. self-attention, and deep-sets in goal-conditioned RL with factorized object states. Likewise, Yoon et al. [13] provide a comprehensive empirical evaluation of pre-trained object-oriented models, such as SA [30] and SLATE [43], for model and policy learning in multi-object RL. Mambelli et al. [7] use linear relational networks to model object interactions and learn the policy. Another line of research focuses on learning structured representations among objects or their interactions [9, 5, 6, 8, 12, 17, 16, 44-49]. Most of these approaches aim to learn an object-wise factorization model, either with structured symbolic input or digh-ddimensoonal raw pixels as input. NCS [12] and STEDIE [17] go further by disentangling netion/control-relevaae or irrelevant features for each object. Unlike these works, we propose a more fine-grained factored world model that considers the structure among all attributes of the objects as well as the dynamic interaction among all objects.
Experimental Evaluation
We consider diverse set of RL benchmarks and setups, including modified OpenAl Fetch envi- ronments [15, 50, 51] (symbolic inputs), Spriteworld [9] (visual inputs), and the Block-stacking benchmark [52](visual inputs). We compare our approach with several baseline models. These include methods using Deep sets, GNN, and Self-Attenton as the inductive bias [15, 4]. such as SRICS [6]. STOVE [8]. SMORL [5], NCS [12], LRN [7]. and COBRA [9]. To ensure fair com- parison, we modify these baselines, so that methods that originally only support symbolic input can also be adapted to handle image- based input by using visual encoders obtain the symbolic states, and add an imagination component. We provide the details of these modifications in App. D.1. We provide descriptions of all setups, training and testing domains for each benchmark in App. D.2. Here we describe the most representative results, but we provide the complete results in App. D.3.
Symbolic benchmark: OpenAl Fetch Push and Switch. Following [15], we modify the OpenAl Gym Fetch environment [50] to create the N-push and N-switch benchmarks, where N denotes the number of objects. In the N-push task, the agent is trained to push all cubes to their corresponding target positions. Similarly, for the N-switch task, the agent is required to flip all switches in the environment. In this benchmark, all inputs are symbolic. As sanity check, we show in App. D.3 that our method has comparable s with the baselines in the single-task mode, in which we train the model estimation and policy learning individually on the 2-Push, 3-Push, 2-Switch, and 3-Switch tasks. To evaluate compositional generulization, we train on the set of tasks (1-Push, 1-Switch, 2-Push, 2-Switch); while during the test phase, we test the model in different settings as shown in Table 1. We consider combinations of skills (denoted by S), e.g. 2-Push+2-Switch (S), which combines the training tasks 2-Push, 2-Switch. We also consider changes in the number of objects (denoted by O) (S+O), which combines the training tasks but also varies the number of objects. We also test whether our model can achieve efficient transfer during testing when the objects' latent parameters differ from the training samples (denoted by L). Generally, during training, we consider the objects with masses and friction coefficients uniformly sampled from a set of values, and during testing, we test the model with two different sets of values. For example 2-Switch (L) considers this case. Finally, we consider the challenging setting in which we combine all of these changes. As seen in Table 1, in most cases DAFT-RL outperforms the baselines, with bigger gain
and skills, e.g. 3- 3-Push+3-Swinn
8

---

Experimcet
Slethod
Scttiags
DAFF-BLSSynhdii)
Sea
Selfoatzation
SRICS
GINN
STONE
SMORL
NCS
2PPub 2 25
0,551 0.005
0.004
0.367 =
0,833 (.127
0,902 =0,050
=0.01
0.912
0.022
1-Snchh (S+0)
0,885 0.324
0:334 1.0.012
0.085 1 (.13)
0.683
00/19
0.39: 1 [0023
(1873 1 10.031
0,09:
00.07
0.751 1 0.123
0.715
00.01
2-hoh
1.96% 8,8%6
B.80 0.02
0.941 1 1:0.99
0.022 0092
0,955 10.007
0.911
1 10817
0994
10.02
00964
1 0.039
08933
00.08
2-S4ic3 (t)
0,933 = 0.865
B591 + 03043
0902
= 0.006
0.99 00022
03733
= +0.005
0.834
02892
+0.003
0.864
= 0.002
01907
-0.349
0,921 0.837
0.088
0.007 6.105
0.088
0,822 0.003
0.91
0.09
0.000 0.036
0.838
(L.0)
2/Pub 2.Swihh (L+5)
08783
0.826
03844 0.077
0.351 4-010
0.476 0.062
0.505 0.043
0.067
0.501 0.051
5/huh 3.Sench (1.005)
0.783 - 0.825
0,500 6 0.045
0.030
0.029
0.00 0,58
0,002 + 0.005
0.50 6 0.042
0.380 + 0.001
0.509 + 0.027
01,11 0.001
Table I: Average success rate over random seeds for Push & Switch compositional generalization in terms of combination of skills (S), changing number of objects (O), and changing latent parameters (L) with respect to training. The numbers in bold highlight the top-performing method.
Expertmment
Mechod
Settings
DAFT RL (SA)
DAFTRL (AIR)
SMORL
SRICS
GNN
STOVE
COBRA
NCS
LRN
Object Giaal
0.997 0.8.1
0.991 0100
0.710 a: 0.093
0.330 a U. H01
01616
(1R8) 0.01
U.47
U.759
Objest Inseraction
(U070
0.035
0.063
0.742
0.104
0,520
0.036
0.m6
0.078
0.827
0.002
0.749
Object Compation
0J067
0.051
8,000
0.775
0.072
0,209
0.121
0.724
0.797 0.057
0.766
0.054
Propenty Compariuu
Table 2: Average success rate over random seeds for Spriteworld with unseen object numbers, color and shape combinations. The numbers in bold highlight the top-performing method.
in the difficult settings. In this setting, we do not compare with COBRA [9]. since it expects pixel inputs. As a representative example, we show the smoothed learning curve for 2-Push+2-Switeh (L++S) in Fig. 3A, in which we display the top three methods in terms of success rate.
Image benchmark: Spriteworld We consider four tasks in Spriteworld [9]: object goal, interaction, comparison, and property comparison. Given that the observations are in pixel space, we use pre- trained SA [30] or AIR [31] as preitrained encoders to obtain the object factored states. Following [13], we generate the datasets with varying numbers of objects placed randomly for pre- pre-training the object- centric model. We consider cases in which we vary the number of objects, have unseen combinations of colors for the objects or unseen combinations of shapes in the target task, and provide the results in Table 2, showing that DAFT-RL with Slot Attention outperforms all baselines, followed closely by DAFT-RL with AIR. In this setting, we do not show the results from DeepSets or Self-Attention, because of bad performance, but we provide the results in App. D.3. We show the smoothed learning curve for the object comparison task in Fig. 3B, in which we only display the top three methods.
Block-stacking In the stacksst benchmark [52], the task is to stack all the blocks into the target position. We use the pre-trained encoders to obtain the object attributes and MPC [41] to learn the optimal policy. We follow the experimental setup in the multi-step planning task configuration in [17], which is specified in App. D.2. As sanity check, we show in App. D.3 that DAFT-RL provides comparable performances to the baselines the single-task mode, We train on varying number of objects and objects with different masses and then use the trained model to test on the domains where there exist an unseen number of objects and unseen combinations of masses. We show the results in Table 3 and Fig. 3C for varying numbers of blocks. Our approaches consistently outperform baselines, with bigger gains for more objects.
Image benchmark:
Ablation studies To evaluate the effectiveness of each compo- nent in DAFT, we conduct the following ablation studies: (I) DAFT w/o latents, (II) DAFT w/o class template graphs, (III) DAFT w/o dynamic interaction graph, (IV) DAFT w/o interaction pattern graphs; (V) Using single class learning (i.e., w/o multi-class la- bels); (VI) DAFT w/o action binding graphs; (VII) Using hard attention networks to learn the action binding. As a representative result, we show the results of 3-Push + 3- 3-Switch (L+O+S) in Fig. 4. The results further indicate the contribution of each part to the effectiveness of final policy learning. We also investigate the effect of using end-to-end learning instead of multi-step learning, the im- portance of imagination, the choice of object-centric models, and the effects of the learned latents and graphs. All results are given in App. D.4-D.7. Additionally, we also provide visualizations of
022
Bate
0.20
0000
20
chee ma TT. vil Brandd aneattec cools
113 waing
Figure 4: Ablation study on the 3-Push+3-Switch task.

---

Experiment
Method
Settings
DAFT-RL(SA)
DAFT-RL (AIR)
SMORL
SRICS
GNN
STOVE
NCS
LRN
2 Blocks
0,809 1: 0.019
0,838 0.030
0.658 1 0.028
0.704 1: 0.016
0.549 1: 0.016
0.728 1: 0.041
0.799 4: 0.085
0.649 0.028
4 Blocks
0.735 1: 0,032
0,608 0.022
0.635
0.033
0.591 +: 0.0:9
0.525
0:011
0.408 0.013
0.571 4: 0.025
0.461 0.026
6 Blocks
0.664 0.017
0.536
0.461
0,88
0.475
0.602
Blocks
0.500 1: 0.083
0.571
0.039
0.386
0.062
0.420 1: 0.061
0.331 4: 0.017
0.278 1. 0.088
0.397 1: 0.052
0. 0.483 1: 0.077
Table Average success rate over random seeds for Block-stacking with unseen object numbers and mass combinations. The numbers in bold highlight the top-performing method.
0.8
NCS
0.8
Shooos Bate
0.6
Surcess Rave
0.6
0.4
0.4
Stocoss
0.6
GINT
SLOVE
0.2
0.2
DAFT STO9E COBRA
0.4
LBS
RUS
0.00
0.5
3.0
2.0
Time Steps (107)
Time Steps (10%)
Number of Objects
The smoothed learning curve for 2-Push 2-Switch (L+S) with different friction coefficients for each object (for clarity, show only the top three methods in terms of the success rate); B. The smoothed learning curve for the object comparison task in Spriteworld with unseen object numbers, combinations of colors and shapes (for clarity, we show only the top three methods in terms the success rate); C. Success rate versus number of blocks the stacking task, where each block has distinct mass.
Figure A.
the learned graphs in App. D.3, which show that our model is capable of learning the true causal of a single object in symbolic cases. For pixel inputs, we do not have ground -ruth causal graphs
Conclusions, Limitations, and Future Work
We proposed Dynamic Attribute FacToored RL (DAFT-RL), a framework that leverages learned Autribute-factored representations with dynamic graphs. For each class object, we learned a class template graph that describes how the dynamics and reward of an object of this class factorize according to its attributes, as well as an interaction pattern graph that describes how it interacts with objects of different classes at the attribute level. We also learned about interactions between objects and with the agent with a dynamic graph. Through this template world model, we learned policy that then be directly applied in new environment by estimating the interactions and latent parameters. We showed that DAFT-RL outperforms the state-of-the-art in three compositional generalization benchmarks.
One limitation of our method is that during the adaptation phase, we only consider object classes that the agent has encountered before. This means that model might encounter challenges when introduced to unseen object classes, One potential area for future exploration is the integration of open- set recognition [53] to identify and adapt unseen object classes. Furthermore, the causal representations [54] in multi-entity environments that could benefit the downstream compositional generalization tasks. Another exciting future direction is to explore the real-world object-centric manipulation or interactive tasks [55-58] or even complex embodied Al tasks [59, 60].
could integrate
Acknowledgments
FF would like to acknowledge the CityU Pirf-Peerfrmance Computing (HPC) resources in Hong Kong SAR and LISA HPC from the SURFnl. SM was supported by the MIT-IBM Watson AI Lab and the Air Force Office of Scientific Research under award number FA8655-22-1- 1-55-22-1-7155.. Any opinions, findings, and conclusions or recommendations expressed this material are those of the author(s) and do not necessarily reflect the views of the United States Air Force.
10

---

References
[1] Danijar Hafner, Timothy Lillicrap, lan Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International Conference on Machine Learning (ICML), pages 2555-2565 PMLR, 2019.
[2] Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. International Conference on Learning Representa- tions (ICLR), 2020. URL https: http///ppeeeeeeeeeeeeeeeeeeeeeee net/ forum? www.ummiiisslil d=$1100C4tDs.
[3] Vincent Micheli. Eloi Alonso, and François Fleuret.
Transformers sample-cfficient world models. In The Eleventh International Conference on Learning Representations, 2023. URL https: / /pwp//pppeeeeeeeeeeeeeee net/ forum? ? idd AAbb0xb
[4] Richard Li, Allan Jabri, Trevor Darrell, and Pulkit Agrawal. Towards practical multi-objeet manipulation using relational reinforcement learning. In 2020 JEEE international conference on robotics and automation (ICRA), pages 4051-4058 IEEE, 2020.
[5] Andrii Zadaianchuk, Maximilian Seitzer, and Georg Martius. Self-supervised visual re- inforcement learning with object-centric representations. In International Conference Learning Representations (ICIR), 2021. URL https: /owp/oopeevveeeeeee net. forum? ? iid= xppImXCbOwl.
[6] Andrii Zadaianchuk, Georg Martius, and Fanny Yang. Self-supervised reinforcement learning with independently controllable subgoals. In Conference on Robot Learning (CoRI.), pages 384- 394. PMLR, 2022. URL https: httpss///ppoeee /ppereeviwww...//////p net./ forum? ded-eddwwennr..
[7] Davide Mambelli, Frederik Träuble, Stefan Bauer, Bernhard Schölkopf. and Francesco Locatello. Compositional multi-object reinforcement learning with linear relation networks. In ICLR2022 Workshop on the Elements of Reasoning: Objects, Structure and Causality, 2022. URL https://opat http///pppeevvveeeeeeeeeeeeeeeee net/ f 1dd-FFOXPr__Seee
[8] Jannik Kossen, Karl Stelzner, Marcel Hussing, Claas Voelcker, and Kristian Kersting. Structured object-aware physics prediction for video modeling and planning. In International Conference on Learning Representations (ICLR), 2020. URL https: /pwpreeeeeeeeeeeeeee ew net. forum?
id=Ble-kxSKDB
[9] Nicholas Watters, Loie Matthey, Matko Bosnjak, Christopher Burgess, and Alexander Lerch- ner. Cobra: Data-efficient model-based rl through unsupervised object discovery and curiosity- driven exploration. arXiv preprint arXiv: 0ttl:/9905000000 9905009775 2019.
[10] Anirudh Goyal, Alex Lamb, Jordan Hoffmann, Shagun Sodhani, Sergey Levine, Yoshua Bengio, and Bernhard Schölkopf. Recurrent independent mechanisms. In International Conference on Learning Representations (ICLR), 2021. URL https://p/// /wwpereeeiveeeeeeee eew net/ / forum? mLcmd EUxy-
[11] Jongwook Choi, Sungtae Lee, Xinyu Wang, Sungryull Sohn, and Honglak Lee. Unsupervised object interaction learning with counterfactual dynamics models. In Workshop on Reincarnating Reinforcement Learning at ICLR 2023, 2023. URL https: /openrevi net forum?
id=dYj JH8Nv81K..
[12] Michael Chang, Alyssa L. Dayan, Franziska Meier, Thomas I.. Griffiths, Sergey Levine, and Amy Zhang. Hierarchical abstraction for combinatorial generalization in object rearrangement. In The Fleventh International Conference on Learning Representations (ICLR), 2023. URL
ht https: wpwprevien.eeeeeeeeeeeeeeeeeeeee net/ forum? ? 2id wwd.nddvvdd-ddp......mmm HGGGvipp3W9.
[13] Jaesik Yoon, Yi-Fu Wu, Heechul Bae, and Sungjin Ahn. An investigation into pre-training object-centric representations for reinforcement learning. In International Conference on Machine Learning (ICML), pages 40147-40174, PMLR, 2023.
[14] Linfeng Zhao, Lingzhi Kong, Robin Walters, and Lawson LS Wong. Toward compositional generalization in object-oriented world modeling. In International Conference on Machine Learning (ICML), pages 26841-26864, PMLR, 2022.

---

[15] Allan Zhou, Vikash Kumar, Chelsea Finn, and Aravind Rajeswaran. Policy architectures for compositional generalization in control. In ICML Workshop on Spurious Correlations, Invariance and Stability, 2022.
[16] Ondrej Biza, Thomas Kipf, David Klee, Robert Platt, Jan-Willem van de Meent, and Lawson LS Wong. Factored world models for zero-shot generalization robotic manipulation. arXiv preprini 2rXXi::2022.5333,, 2022.
[17] Akihiro Nakano, Masahiro Suzuki, and Yutaka Matsuo. interaction-based disentanglement of entities for object-centric world models. The Eleventh International Conference on Learning Representations, 2023. URL https: /ttp:///ppppeeie/eeee/e wwwp://wwwwwwweeeeeeeeeeeeeeeeee
htte://eooooummmmde/eeeeeeoeeeee 1d=
[18) Brian Sallans and Geoffrey E Hinton. Reinforcement learning with factored states and actions. The Journal of Machine Learning Research, 5:1063-1088, 2004.
[19] Michael Kearns and Daphne Koller. Efficient reinforcement learning in factored mdps. In Proceedings of the 16th international joint conference on Artificial Intelligence (ICCAI), pages 740-747, 1999.
[20] Craig Boutilier, Richard Dearden, and Moisés Goldszmidt. Stochastic dynamic programming with factored representations. Artificial intelligence, 121(1-2):49-107, 2000.
[21] Carlos Guestrin, Daphne Koller, Ronald Parr, and Shobha Venkataraman. Efficient solution algorithms for factored mdps. Journal of Artificial Intelligence Research, 19:399-468 2003.
[22] Assaf Hallak, François Schnitzler, Timothy Mann, and Shie Mannor. Off-policy nodel-based learning under unknown factored dynamics. In International Conference on Machine Learning (ICML), pages 711-719. PMLR, 2015.
[23] Sammie Katt, Frans A. Oliehoek, and Christopher Amato. Bayesian reinforcement learning in factored pomdps. In Proceedings of the 18th International Conference on Autonomous Agents and MultiAgent Systems (AAMAS), page 7-15, 2019.
[24] Natalia Gardiol and Leslie Kaelbling. Envelope-based planning in relational mdps. Advances in Neural Information Processing Systems (NIPS), 16, 2003.
[25] Carlos Guestrin, Daphne Koller, Chris Gearhart, and Neal Kanodia. Generalizing plans to new environments in relational mdps. In Proceedings of the 18th International Joint Conference on Artificial Intelligence (IJCAI), pages 1003-1010, 2003.
[26] Martijn Van Otterlo. A survey of reinforcement learning in relational domains. Centre for Telematics and Information Technology (CTIT) University of Twente, Tech. Rep, 2005.
[27] Carlos Diuk, Andre Cohen, and Michael L Littman. An object-oriented representation for efficient reinforcement learning. In International Conference Machine Learning (ICML), pages 240-247, 2008
[28] Jonathan Scholz, Martin Levihn, Charles Isbell, and David Wingate. A Ahysics-baseee bassd model prior for object-oriented mdps. In International Conference Machine Learning (ICML), pages 1089-1097. PMLR, 2014.
(29) Arthur Wandzel, Yoonseon Oh, Michael Fishman, Nishanth Kumar, Lawson L.S. Wong, and Stefanie Tellex. Multi-object search using object-oriented pomdps. In 2019 International Conference on Robotics and Automation (ICRA), pages 7194-7200, 2019. doi: 10.1109/C
2019.879388..
[30] Francesco Locatello, Dirk Weissenborn, Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. Object-centric learning with slot attention. Advances Neural Information Processing Systems (NeurIPS), 33:11525-11553 2020.
[31] SM Eslami, Nicolas Heess, Theophane Weber, Yuval Tassa, David Szepesvari, Geoffrey E Hinton, et al. Attend, infer, repeat: Fast scene understanding with generative models. Advances in neural information processing systems, 29, 2016.
12

---

[32] Biwei Huang, Fan Feng, Chaochao Lu, Sara Magliacane, and Kun Zhang. AdaRL: What, where, and how adapt in transfer reinforcement learning. In International Conference on Learning Representations (ICLR), 2022. URL https: /ttp:///pp//vvveeeee ew. net forum? f 8H5bpVwvt.5.
[33] Fan Feng, Biwei Huang, Kun Zhang, and Sara Magliacane. Factored adaptation for non- stationary reinforcement learning. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
[34] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems (NIPS), 30, 2017.
[35] Ondrej Biza, Robert Platt, aan-Willem van de Meent, Lawson L.S. Wong, and Thomas Kipf. Binding actions to objects in world models. In ICLR2022 Workshop on the Elements of Reasoning: Objects, Structure and Causality, 2022. URL https: http/////eeeeeeeeeeeeeeeeeeeeeem net/ forum? Imm BuUccl..
[36] Thomas Kipf. Ethan Fetaya, Kuan-Chich Wang, Max Welling, and Richard Zemel. Neural relational inference for interacting systems. In International Conference on Machine Learning (ICML), pages 2688 2697. PMLR, 2018.
[37] Colin Graber and Alexander G Schwing. Dynamic neural relational inference. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 8513-8522 2020.
[38] Kevin Murphy. Dynamic bayesian networks: Representation, inference and learning. UC Berkeley, Computer Science Division, 2002.
[39] Junyoung Chung, Caglar Gulcehre, KyungHyun Cho, and Yoshua Bengio. Empirical evaluation of gated recurrent neural networks on sequence modeling. NIPS 2014 Deep Learning and Representation Learning Workshop, 2014.
[40] Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterizztion with gumbel-softmax. In International Conference on Learning Representations (ICLR), 2017. URL https:// openreview. /
[41] Eduardo F Camacho and Carlos Bordons Alba. Model predictive control. Springer science business media, 2013.
[42] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347 2017.
[43] Gautam Singh, Fei Deng, and Sungjin Ahn. Illiterate DALL-e learns to compose. In Interna- tional Conference on Learning Representations (ICMI.), 2022. URL http:::tpp/p/ottt htttp::/opeerevieet net/ / forum? htw.ummmm.ooooo00eee3eo....
[44] Thomas Kipf, Elise van der Pol, and Max Welling. Contrastive learning of structured world models. In International Conference Learning Representations (ICLR), 2020. https:
//openreview.. net. wtwpmmmddddlglgllttttt
[45] Caleb Chuck, Kevin Black, Aditya Arjun, Yuke Zhu, and Scott Niekum. Granger-causal hierarchical skill discovery. arXiv preprini arXiv:2306.0095 2023.
[46] Jiaheng Hu, Zizhao Wang, Peter Stone, and Roberto Martin-Martin Elden: Exploration via local dependencies. arXiv preprint arXiv:2310.0877 2023.
[47] Rim Assouel, Pau Rodriguez, Perouz Taslakian, David Vazquez, and Yoshua Bengio. Object- centric compositional imagination for visual abstract reasoning. In ICIR2022 Workshop on the Elements of Reasoning: Objects, Structure and Causality, 2022.
[48] Jason Hartford, Devon Graham, Kevin Leyton Brown, and Siamak Ravanbakhsh. Deep models of interactions across sets. In International Conference on Machine Learning, pages 1909-1918. PMLR, 2018.
13

---

[49] Somjit Nath, Gopeshh Subbaraj, Khimya Khetarpal, and Samira Ebrahimi Kahou. Discovering object-centric generalized value functions from pixels. International Conference on Machine Learning (ICML), pages 35755-25768 PMLR, 2023.
[50] Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojeiech Zaremba. Openai gym. arXiv preprint 1r06vv1606.015555 2016.
[51] John Banister Lanier. Curiosity-driven multi-criteria hindsight experience replay. University of California, Irvine, 2019.
[52] Matthias Plappert, Marcin Andrychowicz, Alex Ray, Bob McGrew, Bowen Baker, Glenn Powell, Jonas Schneider, Josh Tobin, Maciek Chociej, Peter Welinder, et al. Multi-goal reinforcement learning: Challenging robotics environments and request for research. arXiv 2018.
preprint
[53] Chuanxing Geng, Sheng-jun Huang, and Songcan Chen. Recent advances in open set recog- nition: survey. IEEE transactions on pattern analysis and machine intelligence, 43(10): 3614-3631, 2020.
[54] Phillip Lippe. Sara Magliacane, Sindy Löwe, Yuki M Asano, Taco Cohen, and Efstratios Gavves. BISCUIT: Causal representation learning from binary interactions. In Proceedings of the Thirty-Ninth Conference on Uncertainty in Artificial Intelligence (UAI), volume 216 of Proceedings of Machine Learning Research, pages 1263-1273. PMLR, 2023.
[55] Anurag Ajay, Seungwook Han, Yilun Du, Shaung Li, Abhi Gupta, Tommi Jaakkola, Josh Tenenbaum, Leslie Kaclbling, Akash Srivastava, and Pulkit Agrawal. Compositional foundation models for hierarchical planning. arXiv preprint arXiv:2309.08581 2023.
(56) Mengjiao Yang, Yilun Du, Kamyar Ghasemipour, Jonathan Tompson, Dale Schuurmans, and Pieter Abbeel. Learning interactive real-world simulators. arXiv preprint 10000:210006114, 2023.
[57] Yifeng Zhu, Zhenyu Jiang, Peter Stone, and Yuke Zhu. Learning generalizable manipulation policies with object-centric 3d representations. In 7th Annual Conference on Robot Learning (CoRL), 2023.
[58] Jiayuan Mao, Tomás Lozano-Péres Joshua B Tenenbaum, and Leslie Pack Kaelbling. Learning reusable manipulation strategies. In 7th Annual Conference on Robot Learning (CoRL), 2023.
[59] Abhishek Padalkar, Acorn Pooley, Ajinkya Jain, Alex Bewley, Alex Herzog, Alex Irpan, Alexander Khazatsky, Anant Rai, Anikait Singh, Anthony Brohan, et Open x-embodiment: Robotic learning datasets and rt-x models. arXiv preprint wrXiv:2310.08864 2023.
[60] Chengshu Li, Ruohan Zhang, Josiah Wong, Cem Gokmen, Sanjana Srivastava, Roberto Martín- Martín, Chen Wang, Gabrael Levine, Michael Lingelbach, Jiankai Sun, et al. Behavior- lk: A benchmark for embodied ai with 1,000 everyday activities and realistic simulation. In Conference on Robot Learning (CoRL), pages 80-93. PMLR, 2023.
[61] 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 Conference on Robot Learning (CoRL), pages 1439-1456 PMLR, 2020.
14

---

Appendix for "Learning Dynamic Attribute-Factored World Models for Efficient Multi-object Reinforcement Learning"
A Detailed Discussion and Comparison with Related Work
In this section we discuss the related work. We first shortly discuss related work in terms of our modelling assumptions and then provide a comparison with other object-centric RI methods for compositional generalization.
A.1 Factored, Relational and Object-oriented MDPs
Our modelling assumptions, formalized Dynamic Attribute FacTorre Markov Decision Process (DAFT-MDP) in Definition 1, are related the literature on factored and object-oriented (PO)MDPs. In particular DAFT-MDPss are first order extension with class template graphs, interaction patterns and interaction graphs of factored (PO)MDPs [18-23]. In particular they define a family of factored POMDPs, in which the objects O vary across environments in numbers, types and latent parameters. This extension takes inspiration from Relational MDPs [24-26] and their literature, especially Object-Oriented (PO)MDPs [27-29], which also define family of MDPs that vary in terms of objects and types. Most of these methods focus on discrete states and often define relations states as first-order predicates, with the exception of [28], which propose physics- based approach that circumvents these issues, allowing for more realistic robotic applications.
As opposed to these works, we also consider continuous state spaces and propose more fine-grained factorization of the transition and reward the object attribute level, based on estimating class template graphs and interaction pattern graphs. Additionally, parameters that modulate the dynamics cach object and use dynamic graphs to account for the sparsity of interactions between objects and between the agent and object.
consider object-specific latent
A.2 Compositional Generalization with Object-centric RL Approaches
Table Al provides a comparison of the object-centric RL methods discussed in Section The criteria for comparison include the extent of factorization in the dynamics, the interaction modeling, and the model's ability to adapt to changes in the environment with changing latent factors.
Methods
Dynamies
Factored
Interaction
Modeling
Generalize to Changing Lafeut Factors
GNN [1]
Object
Fully conneted
Self-ottemion [15]
Object
felly---mmettl
Derp Sets [15]
Objod
Fully-eemmetal
LRN [7]
Object
Fully connected
COBRA [9]
Object
x Fully colletedmeeted
FWM
Object
SMORL[1 [S]
Object
x Fully connected
STOVE [9]
Object & determined atrributes
SRICS (6]
Object
Sparse & dynamic (ubject level)
NCS [12]
Object & acdinnkkymmizs-ells
Fully Fully-comestal
STEDIE (17)
Oljjett interacionreeevun"
Fully-cumeeted (object-level), factored (inemncioe-mleenn Sporse, dynamic factored (objesslevel attribee-evvlll
DAFT-RI
Obgect & wisibue
STUVE focuses specifically
attibutte: (position, velocity,
"TEDDE learns which attributes interact during object interactions, similar to our interaction pattern gmphss.
Table Al: Comparison of different object-oriented RL methods for compositional generalization. A check denotes that a method has an attribute, whereas a cross denotes the opposite.
In the context factored dynamics, most of these related works take into consideration the object- factored states. Approaches such as STEDIE [17]. NCS [12], and STOVE [8] go step further to disentangle the states related interaction, dynamics, or actions from other state variables, such as color, as mentioned in the Table. However, our model aims to offer more comprehensive approach factored dynamics, providing an attribute-level factorization for each object. For example this can include factorizing the transition function of dynamics-related factors in terms of position, velocity and
15

---

mass, and factorizing action- irrelvatte factors like color and shape. In terms of interaction modeling, many approaches assume fully- fully-connected interaction patterns, where interactions occur between every pair of objects within given scene. However, in many real-world scenarios, this is often inaccurate as the object interaction usually happens a sparse and dynamic manner. An exception is SRICS [6], which takes into account dynamic interaction modeling. This model emphasizes the sparsity of interactions among objects in a scene, with a dynamically changing structure at the object level. In contrast, STEDIE [17] employs learned factored dynamics for each object to model interactions, assuming that only interaction-irreles
objects will have an impact on others. Nonetheless, from the perspective of the object level, the pattern of interaction in STEDIE is still fully connected. Our method diverges from these approaches by integrating both dynamic and sparse interaction patterns at the object level, and factored interactions at the attribute level between pairs of objects. We provide a more fine-graing ame-ggained and realistic model of interactions as from the object level, the interaction pattern is dynamic and sparse, and for each object pair, the interaction only happens among some specific attributes. Finally, DAFT-RL is the only framework that considers the latent changing factors in the environment, which is also realistic setting where the agent sometimes cannot observe all essential attribute variables in the system but these unobserved factors do have effects on the dynamics or reward.
B Full Example with Summary of Notations
Notation Cc.....,CC {C C,.r.....
Definition
Set of object classes
Cz.so3
Class attributes for class Cj
C
Latent constant parameters for class C₁
(oi...., mm3
Set of objects in domain
C(i)
Class of object ot, other words C, such that o, E C, Object attributes for object o; with semantics defined by the class C(i)
[00.81, 00.82, 02.8n}
0:0
Latent constant parameters of object
Pixel observation at time step [
x
Pixel observation of object-i at time step Symbolic state of object
o f0 0:.82
0
time step f
{o]
Observable state at time step I for enviromment with objects Latent constant parameters for environment with O objects Action object at time step /
0 f0r Om θ
a' [af. J
Action at time step environment with objects
Reward on object o, at time step t
r Eiri
Reward at time step
Gc,
Class template graph for object with class Cj
gc.cc
Interaction pattern graph for objects with class C, and C₁
Action selector at time step
G: Piater
Instantiation of the class template graph for os with class C, at time Dynamic object interaction graph at time step
Table A2: Summary of notations in this paper.
In this section, we first provide a summary of the notation in Table A2, Then we show an example of the environment described in the main paper, and how the learned graphs are connected in a single ground graphical model, as described in Figure A1.
In our example, we consider two classes C = {C1, C2), which represent the classes of boxes and switches, respectively. We represent these two classes with eubes and cylinders in Figure A1.
For the class of boxes C1, represented in Figure AIA as a pink cube, we consider the attributes (Ci-s1; C1.82; C1-83) represent the color, velocity and position of abstract box in two consecu- tive timesteps t and t + 1. All of the relationship between the attributes are described in the class template graph Gc, The edge C1.sf - C1.s/+1 represents the fact that the color at timestep + 1 is only influenced by the color at timestep t (in this case being constant). The edge C1.s² + C2.s5 !+1 represents the fact that the velocity at timestep + 1 is influenced by the velocity at the previous timestep t. The edge C1.sh + C1.stt! 2+1 means that velocity can influence position in the next timestep.
16

---

Class Template Graphs
Interaction Pattern Graph
C. Interaction Graph
G-S
%
1
G
:
x
G
"
Py-ag
+
My 1
My. 2
Oh .
0,.0
Figure Al: The graphical representation of DAFT-MDP. Fig AI.A represents the class template graphs for boxes and switches. Fig AI.B represents the interaction pattern graphs between two boxes. Fig AI.C represents the dynamic interaction graph, that each timestep predicts which objects interact with each other and with the agent. The graph on the bottom shows an example of the instantiation of the complete ground graphical model for all of the objects in the environment from Fig Al.C, i.c. pink box, a blue switch and a green box at time steps f, + 1 and t + 2. The red lines describe the interactions of the agent with the objects (which follow the dotted lines the class template graphs for each object class in Fig A1.A). The blue lines represent the interactions between objects, which follow the interaction patterns described Fig A1.B.
which is also influenced by the previous position C1.5/+1 1+1 In this case the agent can optionally act (the dashed lines from a²) and influence the velocity Cr.s (+1 and position C1.s/t1 Finally the latent constant parameters for the abstract box C,.0, this case representing friction, influence the velocity ++1 at each timestep + 1 and modulate the effect of the action the position and velocity. The reward r' only depends on the position C,.sq.
C1.sh+1 ++1 and position C1
For the class of switches C2. represented in Figure AIA as a blue cylinder, we consider the at- tributes (C2-81, C2.82, C2.83) to represent the color, angle and position an abstract switch two consecutive timesteps and + 1. All of the relationship between the attributes described in
17

---

the class template graph Ic,, Similarly to boxes, the color C2.81 does not influence anything and only influenced by its previous value. In this example, we consider that none of the attributes influences each other, but that the reward re is only a function of the angle C2.sh. Additionally, the latent constant parameters C2.0 0 only influence the position C2-s2. as can potentially the action a (represented by the dashed red lines).
In Figure A IB we show the interaction pattern graph for boxes Gc,,C, that represents the way that attributes of two objects of the class box C₁ interact with each other. In the figure, we use pink to represent the first object and green to represent the second object. Specifically, in this case, the interaction between two boxes means that the velocity of an object influences the velocity of another object. Similarly, the velocity of an object influences the color of the other object. Additionally, the latent constant parameters of an object influence the position of the other object.
In Figure AIC we consider specific environment, specified by the objects O O= (0i, 02, 03), where 01 and 03 are boxes, while 02 is a switch. We show an unrolled version of the dynamic interaction graph Ginter for three consecutive timesteps for these objects. At timestep there an interaction between the two boxes and 03 (represented by blue lines), and the action a' binds to the switch 02, as selected by the action selector a². The interactions between the two objects are instantiated following the interaction pattern graphs in Figure A AIB. The action binding instead activates the red dashed lines in the class template graph in Figure AIA. In the next timestep t + there are no interactions between objects, but the action is now bound to object 01.
In the bottom graph in Figure AI, show how the three types of graphs are combined in this environment for these three consecutive timesteps. In particular, it can be seen that for each object, its attributes follow the same pattern deseribed by the class template graphs 9c; and 9c, (in grey). The interactions between the two boxes in timestep get instantiated following the interaction pattern graphs C (in light blue). The action binding specifies which object is affected by the action a' at each timestep. but the way the action affects the attributes of an object is based on its class template graph. This graph represents the complete graphical model representation for these three timesteps, that describes how the transition and reward functions factorize in these three timesteps.
As can be seen from the example, to learn the factorization of the transition and reward functions in new environment, we can reuse the class template graphs and the interaction pattern graphs across different domains with the same type of objects. We only need to relearn the dynamic interaction graph, which is also the only graph that is dynamic (i.e. the edges do not repeat across every couple of consecutive timesteps).
C Architecture and Implementation Details
In this section, we describe the architecture of DAFT-RL in more detail. Figure A2 summarizes the pipeline of learning the DAFT-RL framework, which uses the notation summarized in Table A2. We first provide a high-level illustration of the framework, and then provide the details in the following subsections. In the offline model estimation, the DAFT-MDP is learned through a two-step process, involving the estimation of both the class template graphs (Step 1), the interaction pattern graphs and dynamic interaction graphs (Step 2). These graphs capture the relationships and dependencies among objects from the attribute level in the environment. Once the DAFT-MDP is learned, policy learning performed using trajectory imagination and planning methods (Step 3). During the adaptation phase, the dynamic graphs are inferred for the new domain (Step 4). This inference step allows for the transfer of the previously trained policy to the new domain without extensive retraining. By adapting the dynamic graphs to the new domain, the policy can be directly deployed and applied to the new environment. In the following, we define the graph structures and losses for each of these steps.
C.1 Binary Matrices in DAFT-RL
We first introduce the definitions of all structures that are used in the loss functions in DAFT-RL:
Binary matrices in class template graphs. As described in the previous sections, we learn the class template graphs Gc, for each object class Cj. This graph is assumed to be constant in time. In practice, these graphs are learned as binary adjacency matrices. To help express the loss functions in a factorized way, we define the following binary matrices.
18

---

Step 1: Class learning single-object environments
Stage 3: Policy Learning
(ON of Di
10 Gi)
(0gg di Die
Object-centric model
Object-Dentric model
( af. Nali
(o), di. o,, M
class inforance
C.
Generate imagined trajectorios
(6), 8.. 6. PM
Learning cless templete graph
8g, Sc,
Learned policy
sataa s. veve)
Stage 2: Interaction Learning and Latents Inference
Stage Adaptation
10x. Bi
10.. Objact-cantricc modal
Objact-cantric model
(l. ol, d'r1)
[e], of. a', 1
Learn the object solector
a'
Interred constant latent parameters object inertaction graph Apply polley
0
Learn object interaction graph
inttrr
inter
pattern graphs
$ cc-r;:1 (1.23). (1.23)
s,0, sec)
Learn the iatent constant parameters
00.00..0000 0..
Figure A2: The learning pipelines of the DAFT-RL. framework.
For l = (1,...,. m2 we define:
Cy.Ge-to E {0, 1 as the subset of the adjacency matrix of Gc, that describes the edges incoming into c from each Cj.sf for E {I,..., nt.. In other words this binary vector is 1 only for the components Cj.s, for [1. that have edge Cy.st+3 in 9c,-
Co.Go-ter E {0,1) 1} as binary value that represents the existence of an edge between af and Cj.st denoting how the action will affect the dynamics at the attribute at next time step.
G..Gn-s E 10 as a binary value that represents the existence of an edge between C,.0 and C,.st 2+1
We additionally define:
Cy.G.mar {0, 13"h is the subset of the adjacency matrix of Gc, that describes the edges incoming into y from each Cj.s for E (1..... n.
Cy.Gatr {0, 1} is a binary value that represents the existence of an edge between a'.
For each object class Cj all of these five groups of binary matrices together represent the adjacency matrix of graph Gc,
Binary matrices in interaction pattern graphs. Interaction pattern graphs Gc. represent the interaction structures between attributes of the object with class C, and the object with class Cj. which are assumed to be constant in time. For each / = {1.... n. we define Gc,-c, E {0, 1}" represent the binary vector that describes how the attributes of an object with class C, influence the attribute Cj.st 1+1 of the object with class Cj. All of these vectors together represent the adjacency matrix of graph Gc.CG
Binary values in Dynamic interaction graphs The dynamic interaction graphs G Ginter represent the a }.
interactions between the objects O at timestep I, which include the action selector {a+ a² where o (0, 1)mh represents which object affected by a' at all time steps = We use {0, 1}
the action selection a' in our losses. For each o,, 0j O we define a binary value G (inte(t..))
19

---

to represent whether object and of will have the interaction at time step t. All of these values together represent the adjacency matrix of graph G inter
C.1.1 Detailed Loss Functions
We provide the detailed loss functions for learning the DAFT-MDP below. In each step. we mark in blue which of the structures defined the previous subsections are learn at that step.
Class learning in single-object environments As described in Sec 3.1, we learn the class template graphs {g gg ..... g } using the collected Deingle, We give the detailed loss functions, including the prediction loss Spred and the sparsity regularization sparsee for dynamics and reward models, for this step below:
Step 1:
T
C Spred
logpx, (or.s)+1 Cj.G.-ts ooi' Co Oaa)
(AI)
+logp, (r) Cy. O
Cj. Oaa))
C
Σ 1=1
I I. + IC + Σ 1Cr.Gamalll + Σ
1IC,.Ga-rll) Bacp:llllIII,,l////l
(A2)
sparse
where we update the prediction models including Ass A,, as well as the binary vectors defined previously G.- 1
Gu-: M
/i=1:,,.r. Cy.Gaorr Cy.Ga-r) (marked in blue). The complete
loss L for Step is:
f = -ppred + Taparse Searre
(A3)
where 2 Ysparse is a hyper-parameter.
Learning the action binding In this step, we learn the dynamic action selector of by using the soft attention networks. We have collected Dimutti which have multiple objects in the scene. We consider that the class template graphs learned in the previous phase are fixed and update the prediction loss with the action selector. The detailed loss function for this stage is:
Step 2.1:
T
2orod ΣΣ>
Po (oi-sf+1 C.GGs 2 Ooi, C. Ca.Ga-ts O fo(a!))
(A4)
where fo is the value network to be updated. We also need to learn the key and query networks fa and fg to update at for each object o (marked in blue). Additionally, we use the same dynamic network X and also update it in this stage.
dNRI [37] to learn the dynamic interaction graph. Additionally, we also learn all the interaction pattern graphs, as well as learn how to infer latent parameters, and how they influence each attribute, this step. We consider the class template graphs and action binding are fixed from the previous step. Detailed loss functions are:
Step 2.2: learning dynamic interaction graph As described in Sec 3.2.2,
T
n
Cpred
log Ps. (o,.) C.GGr o C..Ga-a O af fo(a'),
(AS)
C.Goo+9; 3 0. { GC @ o)
C3 = Σ
KL (gs (z' | lips (z" glit x(tt11)
(A6)
t=2
C3 sparee
GG Ml +
IC 11
(A7)
20

---

where X₂ is the dynamics model we have employed in the previous steps. We learn the encoder/prior networks & to generale the latent vectors z, where we sample the graph Gitter from. Through the binary vectors we also learn the interaction pattern graphs (Gc,,c, (i., latent parameters fo..0) and the binary vectors representing the edges from latent parameters to each attribute (Ci We have completed the loss for step 2.2:
_ pred + 2 CKL + 7
"sparse sparse
(A8)
where 3 and 3 Ssparse are the syper-parameters..
C.1.2 Details for the Other Algorithms
Object class learning For the case with pixel observation as input, we choose to learn the object classes in supervised manner, where the imput includes the extracted feature vectors o, for each object and we have the labels yi for each single objects. We apply cross-entropy loss to learn the mapping.
Interaction graph sampling During the learning of dynamic interaction graphs, we generate the edge distribution z' at each time step t. We sample the edge matrix M of the graph Ginter Specifically, M' ~ Bern(z") (z"), where Bern is the multivariate Bernoulli distribution and all elements mutually independent. We also employ the Gumbel-Softmax trick [40] make the sampling process differentiable.
D Experimental Details
In this section, we summarize all the experimental details. We first discuss our baselines and how we modified them to evaluate them a fair comparison with DAFT-RL, focusing particular on adapting symbolic input approaches to pixel inputs and image-based approaches to symbolic inputs. then describe the experimental setups, including each task description and the modifications we made to the benchmarks.
D.1 Baseline Modifications
Symbolic input approaches with pixel inputs. We adapt various approaches that rely symbolic inputs, such as self-attention [15], deep sets [15], GNN [4], SRICS [6] and LRN [7], to handle scenarios where the inputs raw pixels. To accomplish this, we leverage pre-trained object- centric models, specifically AIR [31], to obtain the state representation of each object within the scene. The pre-trained object-centrie models, specifically AIR, are employed extract object- specific information from the raw pixel inputs. These object-factored states, which represent the individual objects, are then used inputs for the approaches mentioned above. By incorporating the object-centric models and leveraging their extracted object representations, we enable the symbolic approaches originally designed for other input formats be applicable in scenarios where raw pixels serve as the input representation.
Image-based approaches with symbolic inputs. To modify the image-based approaches fit with the benchmarks with symbolic inputs, we adopt the following changes:
SMORL [5]: we remove the SCALOR encoder and directly obtain z from the simulator.
STOVE [8]: similarly, remove the SuPAIR model and directly observe object-f states for dynamics and policy modeling.
NCS [12]: directly observe the type and state variables from the simulator without the slot attention models.
Notably, we do not and cannot modify COBRA [9] for the symbolic case, because COBRA is mostly built upon the scene encoder and decoder.
D.2 Experimental Setups
In this section, we give the task description for each of the three benchmarks and describe the detailed experimental settings for both model estimation and policy learning.
21

---

D.2.1 Task Description
OpenAl Fetch Push and Switch. In this benchmarks the tasks are to push N boxes or flip M switches. The agent can obtain both the object factored states and the global (agent) states. Specifically,, the object states are the object's pose and velocity. For switches, the states are also the switch angle and the position. The action the 3D continuous control of the end effector displacement and ID control of the distance between gripper fingers. The goals are the target position for cubes (pushing task) and the target angle of the switches (switching task). In the experiment, we consider a dense reward, computed by the distance between the object's states and the desired states in the goals.
There are four tasks in the Spriteworld benchmark, as listed below. We follow the task sets as well as the modifications in [13]. The agent directly observes the 2D images as the observation states.
Spriteworld.
Object goal task: the goal is to move the object to the target position without touching other distractor objects. The action includes four cardinal directions. A positive reward will be given if the goal is achieved.
Object interaction task: reward will be obtained if the agent pushes the target to the goal position and the action also includes four cardinal directions.
the goal is push the target a given location in the scene. The
Object comparison task: the agent needs to figure out which object is different from other objects and move this object to the given position. The reward will be given if the goal is achieved and the action also includes four cardinal directions.
Property comparison task: similarly, the agent needs to find out the object with a different color or shape from other objects. This task is generally more challenging as the agent needs to reason the property-level difference while the object comparison task only requires object-level reasoning. The reward and action are the same as the object comparison task.
Stacking. We follow the Block-stacking experimental setups in [17]. In the block-stacking task, the agent can observe the images from the current scene and the goal. The action includes picking and placing the objects and coordinates [6 I.. The goal is to stack m objects with a given position.
D.2.2 Benchmark Modifications
Color, shape, friction and mass in Push & Switch To make the benchmark more realistic, we add two variables into the original states in the benchmark, the object color and shape, both are represented as one-hot vector. Additionally, we also use a different friction coefficient and mass for each object, We implement this by conducting a post-processing filter for each object in the simulator. Specifically, different friction coefficients will result in slower velocity and action effects. Different masses will result in different action effects of the agent. During the training of the estimation model, we use the objects with masses and friction coefficients uniformly sampled from {4, 6, 8, 10) and (0.4. 0.6, 0.8, 1.0} respectively. During testing, the object masses and friction coefficients are sampled from 21, 3, 11, 13) U {3, 5, 7, 9} and {0.1, 1.1, 1.3} U {0.5, 0.7, 0.9} respectively.
Unseen colors, shapes, and numbers of objects in Spriteworld To evaluate the generalization capabilities of our model, we take into account the number of unseen objects, shapes, and colors. We achieve this without directly modifying the benchmark, as it offers the flexibility to alter these properties. During the training phase for model estimation, we use the colors green, red, and yellow, along with a varied number of objects ranging from (3.5, 7). We consider as shapes circles and boxes. During the testing phase, we introduce new elements that were not seen during training. Specifically, we incorporate the color blue and triangle shapes. Additionally, the number of objects is extended to include {2, 4, 6, 83. By considering these changes in colors, shapes, and object numbers, aim to assess the model's ability to generalize to novel configurations and demonstrate its robustness.
Mass in Stacking We adjust the mass of each block by modifying the action effects on the object, treating them as post- prot-srocsssing filters applied to the image frames. particular, picking up a heavier object will require more time compared to selecting a lighter one due to gravity. Placing heavier one will make it faster than the lighter one. During training, we consider the masses ranging from {1, 1.5, 2, 2.5). During testing, we make the masses of the objects ranging from (0.5, 1.25, 2.25,3)
22

---

D.2.3 Offline Model Estimation
In Step 1, we use the GRU model [39] as both the dynamics and reward prediction model to learn the class template graph. The hidden size of the GRU is set to 128. Additionally, we incorporate two ML. layers before and after the GRU, each with a hidden size of 64. During training, we continue optimizing the model until the dynamics and reward models converge, employing the mean squared error (MSE) loss with the Adam optimizer. In Step 2, we use MLP layers to parameterize the soft attention layers [35]. Specifically, for the functions fu. f.. and fu. the hidden sizes are all set to 128, with only one hidden layer for each of them. To learn the dynamic interaction graph, we follow the hyperparameters and model architectures outlined in [37]. However, as opposed to them, we use the previously learned GRU in Step I as the decoder. For various experiments, we adopt different hyperparameters for the GRU decoder. The training details for each experiment are provided below.
Push & Switch In Step I, we gather total of 400 trajectories for both boxes and switches (200 for each). Each trajectory consists of 100 time steps, and the actions are generated using a random policy, The batch size for this stage is set to 80, and the weighting parameter before the regularization term 0.015. In this step, in each trajectory, there is only one single object and different trajectories have different objects with different types (e.g., box or cubes), different colors (e.g., blue, red, etc). and different latent parameters (e.g., friction coefficients). In Step 2. we collect 500 trajectories, where each trajectory consists of 50 time steps, for learning the action binding selector. The batch size for training the soft attention network set to 100. Regarding the learning of dynamic interaction graphs, we use the same parameters and model architectures as described in [37]. For the decoder, we reuse the decoder used in Step I, which includes the interaction pattern graphs and the class template graphs. The balancing weights before the KL divergence and the sparsity regularization are set to 0.9 and 0.2, respectively. During training, we set the learning rate to 2.5 x 10-5, and we apply a decaying factor of 0.5 every 100 training epochs.
Spriteworld We adopt the same pre-trained procedures outlined in [13] for the object-centric models. Specifically, for both SA (Slot Attention) [30] and AIR (Attend, Infer, Repeat) [31], we generate a dataset containing various objects with distinct colors and shapes. The dataset comprises 200 trajectories, each consisting of 50 time steps. Within each frame, there are typically 2-5 objects exhibiting diverse range of colors and shapes. For Step 1, we use the collected dataset of 300 trajectories (with 50 time steps each) to learn the class template graphs with different objects. Each trajectory has the transitions of one single object. The weighting parameter before the sparsity regularizer is set to 0.3. The batch size is 64 and the learning rate is 0.002. In our Spriteworld scenario, dense reward signals are absent. Therefore, we solely focus on learning the graphs between states and actions for most steps, excluding those that achieve the goal. Additionally, there is no direct interaction among objects, which means we do not train Step 2 for inferring the interactions. This setup aligns with [13] and is consistent across all the baseline methods we employ.
Stacking We also pre-train the SA and AIR using the randomly collected dataset containing 200 trajectories, each with 30 time steps. The blocks are initialized with different colors in the dataset. To train Step 1, we collect the single-block dataset with 200 trajectories, each with 30 time steps. The balancing weight for sparsity regularization is 0.05. For Step 2, we have the dataset with 300 trajectories, each with 30 time steps and multiple objects. We use the same set of for learning the soft attention networks as Push & Switch. For learning the dynamic interaction graph, we balanced weights before the KL divergence and the sparsity regularization is 0.2 and 0.1, respectively. We the same learning rate to 2.5 x 10-5, and we apply a decaying factor of 0.5 every 100 training epochs.
D.2.4 Policy Learning and Planning
For the Push & Switch and Spriteworld experiments, we use PPO [42] for all baselines with the trajectories generated by the learned dynamics and reward functions. The learning rate for PPO is 0.005 and 0.003 in these two benchmarks, respectively. The coefficient of the entropy term in PPO is 0.05. The policy is parameterized with MI.P layers of size 256 and 128 in both experiments. For Stacking, we follow all the MPC hyper-parameter settings in [17] for the planning.
23

---

D.3
Full Results
We provide the full results of the experiments in this section, including the quantitative results for all experiments, ablation studies, and the visualized graphs.
D.3.1 Quantitative Results
Table A3 and A4 give the results of the single training task in Push Switch and Spriteworld benchmarks. Table A5 and A6 provide the results of Spriteworld with 1) changing object numbers and 2) changing object colors and shapes simultaneously during the testing phase, respectively.
Expertment Settings
(Symbolic)
Deep
STOVE
2-Pash
0.033
0:981 0,923
0.973
0.968
0.906 0.064
3-Push
0972
0.205
+0112
Table A3: Average success rate over 3 random seeds for Push & Switch environments testing (single task training mode). The numbers in bold highlight the top- cop-performing method.
F Esperiment
Method
Settings
DAP7 (SA)
DAPT (AIRI
SMORL
SRICS
GNN
STOVE
COBRA
NCS
LRN
Ofjest
0.765
0.066
0.461
0.715 0.049
0.846
Object
Object Compunison
Property Comparione
Table A4: Average success rate over 3 random seeds for Spriteworld environments training (single task training mode). The numbers in bold highlight the top-performing method.
Experiment Settings
Method
DAPT
(AIR)
Object
0.759
Object Comperienn
01159
0.061
0.1002
0.911
0.070
0.865 0,083
0.782
0:353 0.173
0.711 0.087
0,875
0.70
0.772 0.997
Table A5: Average success rate over 3 random seeds for Spriteworld environments testing (with unseen object numbers). The numbers in bold highlight the top- hop-performing method.
24

---

Experiment
Method
Settings
DAPT (SA)
DAPT (AIR)
SMOKL
SRICS
GNN
STOVE
COBRA
NCS
LRN
Ofject Gual
0,902 = IJ136
0,557
=: 0.641
0:00 = 0.071
0.785 0.062
0.42 de 0.069
0.664 + 0.050
0.681 =: 0.053
0.836 0.023
0.775 0.060
Ofject Intenaction
0.895
IJ172
0,580
0.007
0.710
+ 0.065
0.772 0.0.00
0.350 d: 0.140
0.65 6.000
0.688 =: 0,076
0.839
0.864
0.769 =: 0.117
Objest Compasioon
0.878
0.00
0 0.63
0.864
4: 0.052
8.763 0.074
0.429 2: 0.125
B.81 2.0.02
0.708 4: 0.159
0.824 0.056
0.726:20000
Proponyy Comparinn
0.912
0.885
0.081
0.325 0.171
0.000 0.085
0.725 000000
Table A6: Average success rate over 3 random seeds for Spriteworld environments testing (with unseen object colors and shapes). The numbers in bold highlight the uop-performing method.
D.3.2 Full Ablation Studies
Fig, A3 gives the full ablation studies on Push & Switch and Stacking benchmarks. We consider the following cases:
DAFT-RL w/o latent parameters;
DAFT-RL w/o factored class template graph;
DAFT-RL w/o dynamic interaction graph:
DAFT-RL w/o factored interaction pattern;
Using single class learning instead of multiple object classes;
DAFT-RL w/o action binding graph;
Using hard attention networks for action binding:
2- Puá 2 Swiidui)
2- Pesli))
SwitchL)
2- Pish -2- Snitaat
0.3
0.5
0.7
0.6
0.7
0.0
10
VT
Vir Vini
11
211
V1 VITI
Vin
10
vin
Cascs
Cases
Clucs
Cases
Fush Swindll.
Pudiil
Swindll.
0.8
Succeo [tate
0.5
0,70
0.83
0.6
110
V11 VII1
VII
V2II
17
VIII
111
VII VIl
Casss
Clasos
Coacss
Cases
2 Iiluddss
Blods
0.825
0.05
0,800
0.70
0.776
0000
0.50
0.40
0.785
0.60
V1
V11I
vir
vhn
VI
VII
vin
111
10
vhr Vin
Figure A3: Ablation studies Push Switch and stacking tasks. I. without latent parameters: II. without factored class template graph; III. without dynamic interaction graph; IV. without factored interaction pattern; using single object class; VI. without action binding graph; VII. using hard attention networks; VIII. Original DAFT-RL. For those tasks where the latent parameters do not exist, we did not conduct ablations on case Similarly, for those tasks with only one object, did not include the ablations on case V.
Based on the obtained results, we observe that each learned component plays crucial role in optimizing the policy. Notably, the dynamic interaction graph consistently demonstrates significant contribution across the majority of experiments in Push & Switch and stacking benchmarks.
25

---

Task
DAFT (SA)
DAFT (SA) wño class template graphs
Object goal
0.897
0.767
Object interaction
0.739
Object comparison
U. 893
0.685
Property comparison
0.907
0.832
Table A7: Ablation studies on the Spriteworld benchmark.
For Spriteworld, since there is no object interaction and latent parameters in this task, we only conduct the ablation studies on the class template graph. The results (Table A7) suggest that the class template graph can benefit policy optimization as well.
These findings highlight the importance of factored and dynamic graphs in capturing and modeling the attribute-level interactions between objects. They suggest that understanding and incorporating the dynamic relationships and dependencies among objects have a substantial impact policy optimization.
D.4 Results on End-to-end Learning
Table A8 presents the results of end-to-end training. For fair comparison, of data samples as in the original multi-step training. In this end-to-end approach, only Dmalti is employed both the interaction model and object template graphs trained simultaneously. The results indicate that while the multi-step training consistently outperforms, the agent can still achieve respectable compositional generalization performance with the end-to-end training method.
the number
Experiment
Settings
2-P42-8(55
(S+3)
2-P.
(1.
3-P(1.+0)
(1-0)
2-P+2-5 (1+55
3-P+3-5 (1+0)5)
Original (main pupper)
0,881 1: 0.008
0.835
0.024
0.063
0.036
1. 0.065
0.921
0.0077
0.903
0.003
0.793
0.026
0.783
0.025
End to-end
0.821 = 0.076
01763
0.023
0.800
0.068
(1832 4 ILAKS
0.878
0,02
0.796 4 0.047
01833 #: 0.1124
0.707 0.002
Table A8: Results on end-to-end training in Push & Switch compositional task. m-P and n-S denote the m-Push and n-Switch task. The numbers in bold highlight the lop-performing method.
D.5 Effect of Imagination
Table A9 gives the results where we disable the imagination for all methods. The results indicate that imagination helps improve the performances of all methods.
Experiment Scttings
Method
DAFT-RI (SA)
DAFT-RI, (AIR)
SMORI
SRICS
GNN
STOVE
NCS
LRN
Original
0.3%
0.086
0.077
imagination
0.07
0.204
Table A9: Results on the cases where imagination is disabled. Average success rate over 3 random seeds for Block-stacking with 8 boxes and different masses.
D.6 Effect of the Choice of Object-centric Models for Symbolic Methods
Table A10 shows the results of adding AIR and SA to the symbolic methods on the 8- 8-blocking with variants in terms of masses, Results indicate that both choices are comparable in terms of the final success rate.
D.7 Effects of Learned Latent Parameters and Graphs
To assess the efficacy of the learned latent parameters, we carried out two sets experiments: (i) incorporating the actual latent parameters (in this case, the mass value) into the state vector for policy learning across all methods; and (ii) integrating the latent parameters learned by DAFT-RL. into the state vectors for all baseline methods. The results are presented in Table AL Notably, while models
26

---

Exporiment Settings
Merhed
DAFT-RI
SRICS
GNN
ERN
AIR
0.506 = C.8S
0.420
0.584 +i 0.047
0.453 0.077
SA
0.571 0.027
0,200
0.283 & 0.036
0.482 0,1122
Table A10: Comparisons on the using AIR or SA as the object-centric model for symbolic methods.
using the true latent parameters did outperform those utilizing DAFT-RL's learned latents, the latents derived from DAFT-RL notably enhanced the performance of all baseline methods. These findings suggest the significance of accurate latent parameter estimation and suggest that DAFT-RL can learn a set of valuable latents to improve generalization.
Additionally, we investigated the quality of learned graphs and the latent parameters, and their effect on the RL performance w.r.t. the number of samples for the 3-Push + 3-Switch (L+O+S) task. In particular, we varied the amount of training data, ranging from 10%, 20%, 40%, 60%, 80% of the original sample size (900 trajectories). We measured the R2 coefficient of our learned parameters with the true latent parameters, both for random policy (Fig. A4a) and for pre-trained policy (Fig. A4b). These results show that the performance degrades with smaller sample sizes, but it is still acceptable with 60% data, and that the difference between data collected with random pre-trained policy is negligible at higher sample sizes; the normalized Structural Hamming Distance between the reconstructed graph and the ground truth graph, as well as the success rate of the learned policy. As expected, the more data, the more accurate the graph, and the better the performance of the policy trained with the is perfectly reconstructed, which also means a good RL performance. Additionally, for limited data (e.g- 0.2 0.4 of the original dataset) leveraging pre-trained policy enhances both graph and policy learning. However, as the amount of data increases, the benefits of using a pre-trained policy diminish.
accurate world model. For the number of samples in the main paper, the graph
Expertment Settings
Method
DAFT-RL (SA)
DAFERL (AIK)
SMOKL
SRICS
GNN
STOVE
NCS
LRN
Oliginal
0.506 #I.ISS
0.571 =0,039
0.386 #6.066
0.420
#0.061
0.334
20.47
0.27%
HOLING
0.307 #0,052
0.463 S0U77
True batent parameter
0.525 20. 40.062
0.598 e0.023
0.478
60.081
0.501
+0.072
0.386 20.005
0.315
20.049
0.410
+0.075
0.102 20
Purmmeter lesrned by DAFT-RL
0.506 2I1,18
0,571 =00039
0.415 20.059
0.498 20,073
20059200041
0.263 20,085
0.309 20L049
0,484 20061
Table A11: Effects of the latent parameter modeling. Original: Results in the original experiments, where for DAFT-RL, use the learned latent and the others without latent parameter estimations; True latent parameter: All methods directly use the true latent as part of the states. Parameters learned by DAFT-RL: all baselines are using the latent learned by DAFT-RL part of the state vectors. Average success rate over random seeds for Block-stacking with 8 boxes and different masses. The numbers in bold highlight the top-performing method.
(9)
(h)
47
0.2
0.0
Samplee
Samples
Figure A4: Quality of learned graphs w.r.t. the number of samples for 3-Push + 3-Switch (L+O+S). We plot the success rate of the RI. task, the R2 coefficient for learned representation vs the true latent parameters, and normalized Structured Hamming Distance (nSHD) between the learned graph and true graph under different number of training samples (measured as a ratio with the ones in the main paper) for (a) data collected by random policy, and (b) data collected by pre-trained policy.
D.7.1 Visualization of the Learned Graphs
In Figure A5, which illustrates the learned graphs for boxes and switches in the Push & Switch benchmarks, we denote the class template graphs for boxes and switches as Al and A2. The variables
27

---

C1.s1, C1.82, C1.S3. and C1.31 represent the color, shape, velocity, and position of the boxes, respectively. Similarly, C2.51, C2.82, S2. and C2.54 correspond the color, shape, position, and angle of the switches. In both cases, C..0 and Cj.0 indicate the friction coefficients.
As expected, in the learned class template graph, observe that shape and color do not have direct effect on the dynamics, and they remain unchanged throughout the dynamics. On the other hand, velocity directly influences the position of the boxes, while the position itself does not affect the velocity. Regarding the switches, their positions do not directly impact the angles since they are fixed specific positions on the panel. Also, as expected, actions applied the switches have an effect solely on their angles. The reward is based on the actions, positions, and angles for both boxes and switches. Additionally, the latent parameters influence the velocity and position of the boxes, while they specifically affect the angle of the switches.
A1
A2
B
99
C2S
cIS
C
C352
CSS
C
++
9
2
C
9
Cas
a
c,0
c.0
Gia
C2
c1.0
c1.0
Figure A5: Learned class template graph and interaction pattern graph in push & switch task. A1-A2: Class template graphs for box and switch. B: interaction pattern graph among two boxes.
As for the interaction graph between two boxes, shown in Figure A5B, the velocity of one box impact the position of the other box, while the other variables do not have direct influence on each other. This learned interaction graph successfully captures the underlying physical processes of the system, enabling an attribute-level understanding of the object interactions. Overall, these learned graphs effectively recover the underlying physical processes of the system.
28


Reviews:
[{"review_id": "QkQEAcuIv4", "reviewer": "Reviewer_akCD", "summary": "The paper presents DAFT-RL, which is a object-based dynamics model for reinforcement learning that learns factorized relationships between attributes of objects.  These relationships are learned in several steps.  First, single-object episodes are created using a random policy, and these are used to train a model of how the attributes of each object class interact with each other over time.  These graphs are then held fixed, and multi-object interaction data is used to train additional graphs representing interactions between the attributes of each pair of classes, and a graph representing which object instances interact with each other.  These learned graphs then help train a next-state dynamics model that predicts each object's attribute values conditioned on the relevant attribute values read off from the graphs.  Finally, this model is used to either train a policy, or for a planning method like Model Predictive Control.  DAFT-RL is evaluated in several experiments, where the number of objects and the latent factors of thos objects change.  The paper highlights that DAFT-RL performs favorably compared to other object-centered reinforcement learning approaches on three simulated robotic benchmarks. ", "questions": "- What attributes are extracted from objects in the pixel observation case?\n- What might the latent parameters of objects be in pixel observation environments, and would it be difficult to learn meaningful latent representations from a small set of data?\n- How sensitive is the RL performance to the quality of the graphs extracted in the first two phases of training?  On a similar note, how sensitive are the graph learning steps to the amount of data used to train them?  I imagine there may be situations where it is difficult to get a full sense of the class template graphs with data coming only from a random policy.  \n", "limitations": "The paper would improve from a more forthright examination of the limitations of the method.  How does the method scale to more complicated tasks with more objects?  What is holding the method back from being able to train end to end?", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The proposed model exploits object structure that exists in many RL problems to get better performance than a broad set of baselines.  To my knowledge, this is a novel approach to an interesting problem.\n- DAFT-RL can generalize to environments with different numbers of objects than what was seen during training.\n- Ablation studies are done to show relative importance of each piece of the proposed approach", "weaknesses": "- Clarity of writing can be improved.  In particular, Section 2 and Figure 1 were difficult to follow because the running example was not explained clearly.  $C_1.s_3$ for example is not defined, and context suggests there are only two attributes of boxes.  The appendix clears up the issues, so this may be fixable with some rewriting.    \n- The multi-phase training seems a bit cumbersome.  How realistic is it to be given interactions with single objects before observing multi-object interactions?\n- It is unclear how DAFT-RL would scale to scenes with many classes of objects, or objects with many attributes.  Would the graphs become very difficult to learn? \n- One of the main contributions is the decomposition of object interactions to the interactions of individual attributes.  Is there some ablation that could be done to analyze the benefits over modeling object-level interactions only? ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback, it will surely improve our paper, especially in terms of clarity. We answer each issue in the following.\n\n**Interactions with single objects and end-to-end-training**\n\nWe think that in many real-world cases, one would first train an agent in simple environments (e.g. single objects), before actually using it in complicated settings, similar to curriculum learning. Nevertheless, we performed an additional ablation with an end-to-end version of our method, which combines Step 1, 2.1 & 2.2 into a single step and uses $\\mathcal{D}^\\text{multi}$ as the training set. We show the results for this method and baselines in Table 1 in the pdf of the general response. The end-to-end method is slightly worse than our multi-phase approach, but it is still comparable to the best baseline (NCS) across all tasks, as shown by the average. Notably, our end-to-end version is better than all other baselines on the more complex tasks like 2 Push+2 Switch (combination of skills S + changing latent parameters L) and 3 Push+3 Switch (S+L+ changing number of objects O).\n\n**Scaling with many classes of objects, or objects with many attributes**\n\nCurrently we are limited by the benchmarks we use, which consist of a restricted number of classes (e.g. 2 classes in OpenAI Fetch) and where each object has only attributes such as velocity, position, color and shape, While we extended the number of attributes by adding the latent parameters as mass and friction coefficients, we think using more realistic simulators with multiple types of objects and more complex attributes is an important next step for this research direction. As a proof of concept, we will include a numerical experiment in which we simulate multiple classes and large attribute sets with random graphs in our final version, and provide an analysis of the learned graphs. We do not expect the graphs to be difficult to learn, but that depends on having enough samples, similar to what shown in one of the following answers.\n\n**Ablation with only object-level interactions**\n\nWe have conducted this ablation examining interactions only at the object level, and show an example in Fig. 2D, while we provide the results for all Push and Switch tasks in Appendix Fig. A3. In both figures, \"DAFT-RL w/o factored interaction pattern\" or “IV” refers to this scenario, where we use a dense (fully connected) interaction pattern graph. In all tasks, not modeling the factored interaction patterns worsens the performance. We will clarify the description of the ablations in the paper to make this clearer.\n\n**Attributes in the pixel observation case**\n\nWhen using slot attention, the extracted attributes are vectors that do not have direct physical interpretations. This allows us to deal with attributes that are not predetermined. When using AIR, the attributes are fixed and they are position, velocity, size, and latent parameters.\n\n**Latent parameters in pixel observation environments**\n\nIn pixel environments, the latent parameters are properties of objects that are not directly observable from a single image, but influence the dynamics. In our experiments, we considered as latent parameters the mass of an object or its friction coefficient. To learn these latent parameters, we consider sequences of images, capturing the temporal dynamics and interactions of objects. \n\nWe detailed the number of samples used for learning these latent parameters across benchmarks in Appendix D.2.3. We also performed an additional ablation, where we varied the amount of training data, ranging from {10%, 20%, 40%, 60%, 80%} of the original size (900 trajectories), and measured the $R^2$ coefficient of our learned parameters with the true latent parameters (a coefficient of 1 is perfect correlation). The results can be seen for 3-Push+3-Switch (L+O+S) in Fig. 1 in the pdf of the general response, where Fig. 1a shows the results for a random policy and Fig.1b shows the results for a pretrained policy. These results show that the performance degrades with smaller sample sizes, but it is still acceptable with 60% data. We also report the results for the pixel case for the random policy in the 8-block stacking task below (where the original trajectories are 700), with similar insights:\n|      Ratio    | 0.8  | 0.6  | 0.4  | 0.2  | 0.1  |\n|----------|------|------|------|------|------|\n| $R^2$ | 0.75 | 0.62 | 0.31 | 0.19 | -0.14 |\n\n**Quality of graphs - impacts on RL, sensitivity to small samples**\n\nThe quality of the graphs, which represent our world model, directly impacts the RL performance. Our analysis of the learned graph for OpenAI Fetch in Appendix Fig. A4 indicates that the extracted graph closely mirrors the underlying physical processes, which helps the RL performance.\n\nWe performed additional ablations on the graph estimation with lower sample sizes and reported them in Fig. 1(a) in the pdf of the general response. In these experiments, we varied the amount of training data for the model estimation stage, ranging from {10%, 20%, 40%, 60%, 80%} of the original sample size, and reported the normalized Structural Hamming Distance between the reconstructed graph and the ground truth graph, as well as the success rate of the learned policy. As expected, the more data, the more accurate the graph and the better the performance of the policy trained with the more accurate world model. For the number of samples we considered in the main paper, the graph is perfectly reconstructed, which also means a good RL performance. We also tested the impact of using a random policy to collect the trajectories vs. a pre-trained policy in Fig 1.(b). For limited data (e.g. 0.2 or 0.4 of the original dataset) leveraging a pre-trained policy enhances both graph and policy learning. This is likely because the pre-trained policy provides a more informative data sample compared to random actions. However, as the amount of data increases, the benefits of using a pre-trained policy diminish."}, {"review_id": "GCLxTEFbMB", "reviewer": "Reviewer_NUSp", "summary": "The authors has proposed Dynamic Attribute FacTored RL (DAFT-RL). For each class of object, class template graphs and interaction pattern graph and interaction graphs are computed. Through this template world model a policy can be learned for direct application in a new environment by just estimating the interactions and latent parameters. Authors showed that DAFT-RL outperforms the state-of-the-art in three compositional generalization benchmarks. Ablation studies ", "questions": "Why the same paper content is repeated in supplementary upto pg 9? remove.\nI was wondering why Dynamic Bayesian Networks (DBNs) was chosen?\nHow can standard knowledge graph and graph network theory approaches be leveraged in this context?", "limitations": "How this method will leverage for multi object multi instance objects with ambiguity?", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper brings in the philosophy of class template graphs and interaction pattern graph and interaction graphs for multi-object RL. This is overall well written paper with good supplementary material details. The results validate the claims.", "weaknesses": "No video demo or code is provided.\nThe definition looks cluttered line 163.\nA pictorial example at the intro of a sample task is expected. (borrow from supplementary add here)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your useful feedback. We will answer each issue in the following.\n\n**Code**\n\nFollowing the NeurIPS policies, we will provide the code in an anonymized link as a separate message to the AC.\n\n**Clarity, visual examples**\n\nWe will try to improve the clarity of the paper, including simplifying the definition and adding a visual example of a sample task.\n\n**Supplementary**\n\nThis was intentional, we have combined the main paper and the supplementary for the sake of convenience in the review phase (so the links work). For the final published version, we will separate the supplementary materials.\n\n**DBNs, KGs**\n\nDynamic Bayesian Networks are a standard way to model Factored MDPs and POMDPs (see for example the references [18-23]) and their extension, relational MDPs [24-26], which are the basis of our factored approach. In our setting, we use only one type of relation (as most MDPs) in a closed-world context, so we do not need knowledge graphs (which also do not have an obvious temporal aspect, at least in their standard form). Using graphical representations that are closer to first-order logic, might be an interesting future extension.\n\n**Multi-object multi-instance objects with ambiguity**\n\nIn our setting, there is no issue of ambiguity in terms of separating the single objects. If we have multiple objects with the same attributes, we can still leverage SA or AIR to extract a representation for each object and then model them as individual objects in our framework. This is even more clear in the settings in which the objects have different latent parameters, e.g. mass or friction coefficients.\n"}, {"review_id": "5v3CmiEjM4", "reviewer": "Reviewer_kNiT", "summary": "The paper presents a new framework for learning world models of environments with multi-object interactions. The presented DAFT-RL framework identifies objects as instances of classes, for which class template graphs describe how the object's dynamics and reward depend on it's attributes, and interaction pattern graphs describe how instances of different classes influence each other in terms of their attributes. The manuscript also presents how to learn a policy from this graph-based world model and shows that this policy can be adapted without additional policy learning to novel configurations by inference of object interactions and latent parameters. The proposed method is evaluated on several datasets and is shown to yield strong performance.", "questions": "3. Line 66: I think it should say \"Objects of the same class will have the same **set of**\" to be as explicit as possible. I had to reread the sentence a couple of times.\n1. Line 121: \"some of the attributes of each object after the other\". I think it should be \"affect\" instead of \"after\"\n1. Line 797-798 (Appendix D.2.2): Is that two layers before and two after, or two, one of which before, the other after?\n\n## Acknowledgement of rebuttal\nI have read the rebuttal and other reviews. The rebuttal addressed my concerns by describing how hyperparameters were selected, providing a discussion of limitations, and some ideas on how to address these limitations. The rebuttal has convinced me of the appropriateness of my accept rating.", "limitations": "I did not see a discussion of limitations. I'd suggest discussing for instance \n6. how the assumption that actions only have a direct effect on one object at a time limits the generality of the framework.\n1. how the framework could be applied, if the sets of classes and attributes are not available", "rating": 7, "confidence": 2, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "* To the best of my knowledge the proposed method is novel \n* Design decisions in DAFT-RL are well-motivated and the overall idea is well explained. \n* The performance improvements are impressive, especially in the environments with more objects.", "weaknesses": "1. I think it would be a lot easier for readers to parse the examples in Figure 1 if they were annotated with example attributes and parameters, e.g. position, velocity, friction coefficient, etc.\n1. I didn't see a description of how the hyper-parameters were selected.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback, it will definitely improve our paper. We answer your questions inline below:\n\n> I think it would be a lot easier for readers to parse the examples in Figure 1 if they were annotated with example attributes and parameters, e.g. position, velocity, friction coefficient, etc.\n\nThanks for the suggestion. Our original intention with Fig. 1 was to also introduce the notation, that’s why we had used the symbols. We also wanted to use it to clarify that the attributes are not predetermined, but that we can use any attribute extracted by OCRL methods like Slot Attention. On the other hand, using explicit parameters would help the readers to understand the setting better, so we will consider it, as well as other changes aimed at improving the clarity and readability of the paper.\n\n> I didn't see a description of how the hyper-parameters were selected.\n\nThanks for the suggestion, we have reported the values in Appendix D.2.3, and we will add the hyperparameter selection procedure in our final version. In particular, for the GRU and attention networks, we optimized the number of hidden units by searching within the options {64, 128, 256}. The learning rate and batch size were set to commonly used values. As for the balancing parameters (sparsity regularization and KL divergence), we varied them from 0.1 to 0.9 in increments of 0.1. However, this was the approach for all tasks except step 1 in stacking. For step 1 in the stacking task specifically, our search range for these parameters was {0.01, 0.05, 0.1, 0.5}.\n\n> Line 66: I think it should say \"Objects of the same class will have the same set of\" to be as explicit as possible. I had to reread the sentence a couple of times.\n\nThanks for the suggestion. We will clarify it in the final version as below:\n\n“Objects of the same class will have the same set of attributes, the same transition and reward functions, but can differ in the values of the attributes (e.g. they are at different positions) and in the value of the latent parameters (e.g. they have different friction coefficients).”\n\n> Line 121: \"some of the attributes of each object after the other\". I think it should be \"affect\" instead of \"after\"\n\nThanks for pointing this out. We will correct it in the final version. \n\n> Line 797-798 (Appendix D.2.2): Is that two layers before and two after, or two, one of which before, the other after?\n\nWe apologize for the confusion. It should be one before and one after, in total two layers. We will clarify it in the final version. \n\n> I did not see a discussion of limitations. I'd suggest discussing for instance 6. how the assumption that actions only have a direct effect on one object at a time limits the generality of the framework.\n\nThank you for your suggestion, we will include a better discussion of the limitations of the method in the main paper, including the assumptions that we use for action binding, but also the assumptions that the classes are predetermined. We also wanted to point out that the assumption of an action having an effect on one object at a time is common in some robotic manipulation benchmarks like OpenAI-Fetch and it can be realistic in those scenarios.\n\n> How the framework could be applied, if the sets of classes and attributes are not available\n\nThank you for raising this interesting point. We first wanted to point out that when combined with slot attention, our method automatically extracts the attributes from the representations, and is therefore able to learn templates and interaction graphs on top of these attributes. Instead, for our implementation with AIR, the attributes are predetermined and they represent position, velocity, size, and latent parameters. We wanted to clarify that we also do not need labels for each new object, but we need to train a classifier that can classify the object based on its extracted features, as explained in Appendix C.1.2. We will improve the explanation of this in the main part of the paper.\n\nIn our current approach, we have not explicitly addressed situations where the sets of classes are not predefined or known. In this case, our method would revert to considering all objects of the same class and learn a class template graph and interaction pattern graphs that would be the union of all of the underlying graphs for each single (unknown) class. For completeness, we performed an additional ablation on 3-Push+3-Switch (L+O+S) and reported its results in Fig. 2 in the rebuttal pdf in the general response. In that case, “V: DAFT-RL with a single class” represents the case of using a single class for this environment, which seems to impact our current method quite strongly and therefore requires a rethinking of the approach for this setting.\n\nExtending our method to unknown classes is an interesting research direction that we plan to explore in subsequent research. For example, if, during testing, we encounter an object of an unknown class, we could employ an open-set recognition method, such as the one described in [1], to detect this novel class. Subsequently, it would be necessary to accumulate a few trajectories pertaining to this new object. This would allow us to learn its class template graph, interaction graph, and any associated latent parameters.\n\n[1] Geng et al. \"Recent advances in open set recognition: A survey.\" IEEE transactions on pattern analysis and machine intelligence 43.10 (2020): 3614-3631.\n"}, {"review_id": "H9Tcxx9cFV", "reviewer": "Reviewer_aVkM", "summary": "The authors propose DAFT-RL, which uses object-centric representations to improve generalization to object attributes, when training world-models for model-based RL. For each class of object they learn a class template graph that maps class attributes to dynamics and rewards. The authors claim that this representation outperforms state-of-the-art generalizing to unseen objects with new attributes on a simulated block stacking task.\n\nFYI: Supplementary Material PDF just links to the paper pdf for me. Maybe intentional?", "questions": "None really. I didn't find any parts of the paper confusing.", "limitations": "Is there a limitations section? The \"modelling assumptions\" is perhaps closely related. There's also no societal impact statements, although I don't think it's necessary in this case given the scope of the experimental results (toy simulations).", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 1, "strengths": "- The object-centric attribute factorization is a reasonable approach when the number of attributes that affect environment behavior (dynamics and rewards) is small, or when there are multiple objects that share attributes.\n\n- On the benchmarks shown, the model does outperform baselines (although more on this in weaknesses).\n\n- I found the explosion easy to follow and the model section is well written (there's a fair bit of notation but at least it's clear).", "weaknesses": "- \"e.g. an object’s position and reward are affected by its previous position, but not by its appearance or activation state\" - nit: appearance and dynamics are not independent in the real world. A metal object appears metal. A rubber object appears rubber. One can infer a lot about dynamics from appearance. I get the authors point, but consider rephrasing for clarity.\n\n- \"Often an object’s transition and reward functions\" - Agreed but a reward for a given object is also task dependent. \"Stack the red block\" means the \"blue block\" has no reward (or even a penalty). Reword for clarity.\n\n- At a high level, the overall method is really just an extension to Object-Oriented (PO)MDPs to include object attributes. I'm not sure this crosses the bar for me for technical contribution. Adding attributes seems like an implementation detail that will help in certain applications, rather than a fundamental algorithmic contribution.\n\n- Experimental evidence is probably the weakest component. OpenAI Fetch, SpriteWorld and Block Stacking are all very simple simulated environments. I think the world-model RL folks have moved passed these and they stand more as unit tests than real RL problems. Real-world results would be nice, but not necessarily required. In particular I can see this method failing when the factored attribute graph is hard to infer or ambiguous (think - a \"rubber duck\" that is just painted metal to look rubbery). I don't see how these simulated environments is really stressing the generative model.\n\n- Likewise, the environments were largely chosen to result in the baselines not working, so I'm not entirely sure seeing margin there is all that compelling. i.e. choose a set of environments with varying attributes until baselines fail, then explicitly predict those attributes so that our model works, doesn't seem overly compelling as an experimental platform. If it was on a real-world experimental domain (with all the associated complexity of real data), that would be more compelling.\n\n- I also worry about how far this general approach scales. Assuming a known set of classes with known observable attributes isn't going to work for complex scenes. Likewise, how how can you generalize to a unknown attribute that is not within your set? The rigidity of the object classes and attributes feels perfectly tailored to simple simulated environments.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback, addressing it will help us make the paper better. We address your issues one by one in the following.\n\n**Supplementary Material**\n\nThis was intentional, we combined the main paper and supplementary material in one document for the sake of readability during the review process, so one can click to get to the right Appendix. We will remove it for the final version of the paper.\n\n**Inferring the dynamics from the appearance**\n\nThat’s a fair point, in many real-world scenarios appearance and dynamics will be strongly correlated, as you mention. We will rephrase and clarify the sentence by specifying that this is what happens in the benchmarks we consider. In particular, in these cases, once the object of interest is identified (“the red block”), its position (and hence in this case the reward) are primarily influenced by its position, and not e.g. by its color or texture.\n\n**Just an extension to Object-Oriented (PO)MDPs to include object attributes**\n\nWe disagree with this point. First, following this logic, one could say that factored MDPs are just an extension of MDPs with factors/attributes, and hence also object-oriented MDPs are just an extension of MDPs with objects. Second, we think our approach neatly combines several ideas from disparate communities: starting from Factored MDPs and Object-Oriented MDPs, to dynamic Neural Relational Inference for modeling interactions, soft attention networks for action binding, object-centric representation learning for extracting attributes from pixels and latent parameters that vary across objects from factored adaptation literature. We also want to point out that our approach focuses on continuous states, while most Object-Oriented MDPs focus on discrete states. We have described the relation to these methods in detail in Appendix A. \n\n**Benchmark environments are too simple and not realistic**\n\nFor our evaluation, we used similar simulated benchmarks from the state-of-the-art methods (like NCS, SMORL, SRICS), which also allows us to compare fairly. To the best of our knowledge, almost all of the related work is evaluated on similar, or even simpler simulated environments. We do agree that more realistic environments would be beneficial to the field, but we are not aware of any current realistic simulator/benchmark that would fit our setting. We would be happy if the reviewer would point us to a real-world setting, so we can use it in our final version.\n\n**The environments were largely chosen to result in the baselines not working**\n\nWe disagree on this point. First, we proposed some modifications to the existing benchmarks, which we think are more realistic and in which each object has individual characteristics that can be latent, e.g. mass and friction coefficients for a task of stacking blocks. This is common in presenting a new idea or method, otherwise one could never introduce a new benchmark. Second, we had also evaluated the baselines on the original settings of the benchmarks (refer to Table 1, A3-A4). Our findings indicate that our framework performs on par with other state-of-the-art methods even in these standard conditions. \n\n**Unknown set of classes or unknown attributes**\n\nFirst, we would like to point out that we do not need to have a priori known or observable attributes, since the attributes are extracted by object-centric approaches at training time and we are also able to model latent per-object parameters. \n\nOn the other hand, we do have these limitations in terms of unknown attributes or classes:\n\n- *Unknown attributes during testing*: In our current setup, the model will not inherently generalize to attributes that are not part of the training attributes. We plan to address this limitation in the future. One potential approach could involve leveraging active learning or open-set learning to enable the model to recognize and handle new attributes, even with limited exposure.\n\n- *Recognizing new object classes during testing*: For situations where a new class emerges during testing, we could consider using open-set recognition methods. These methods can identify when an object does not belong to any known class, prompting us to then gather more data on this object to learn its class template graph and interaction graph. \n\n- *Unknown classes*: Our ablation study on 3-Push+3-Switch (L+O+S) in Fig. 2 of the rebuttal pdf in the general response shows the importance of modeling known classes. In particular, “V: DAFT-RL with a single class” represents our method with a single class (which is the case if we cannot classify the objects), and shows a considerable decline in performance. This confirms that considering multiple object classes, even in a simplified environment, greatly aids in achieving policy learning.\n\n- *Scaling concerns*: We recognize the challenge of extending our approach to complex real-world scenes with many object classes and attributes. Incorporating more flexible representation learning techniques or utilizing unsupervised or semi-supervised learning might be potential directions to explore to make the framework more adaptable.\n\n**Limitations section**\n\nWhile there is no explicit limitation section, we will extend and improve the discussion of the limitations of the method in the main paper, including the assumptions on the classes and attributes mentioned in the previous answer."}, {"review_id": "3b394gFl2O", "reviewer": "Reviewer_xWTh", "summary": "This paper proposes a methodology to tackle compositional generalization, allowing the policy to be able to generalize to previous unseen combinations of objects or compose previously learned tasks. More specifically, the authors propose the Dynamic Attribute Factored RL (DAFT-RL) framework, which involves learning various graphs which model the objects and their interactions. Experimental results reveal superior performance relative to many baseline methods from the literature.", "questions": "- The lack of the dynamic interaction graph leads to the biggest drop in performance in the ablations -- If there are two blocks and two switches, and the number of interactions is limited, why does modeling the dynamic interactions in a graph help so much?\n- How would an end-to-end approach do?", "limitations": "There is no discussion about the limitations of the work.", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Well-written. The paper is well-written, and the figures aid in the understanding of the methodology.\n- Strong experimental results. The proposed method is compared to a number of methods from the literature and outperforms all of them.\n- Ablations for major components. The paper includes ablations for the major components of the design, showing that all of the major components help in performance.", "weaknesses": "- Comparison to prior work. It is unclear whether the comparison to prior work is completely fair. For instance, an imagination component is added to the baselines -- could this hurt the performance of the baselines? Furthermore, Table A4 in the Appendix shows strong results for some of the baselines (SMORL, NCS) during training, outperforming the proposed methodology, potentially implying overfitting. Is the capacity of the baselines comparable to that of the proposed method? Finally, is it possible that by using SA / AIR or the ground truth latent vector z from the simulator for image-based methods, you are actually hurting performance of these methods which use a visual encoder other than SA / AIR?\n- Choice of tasks. The best performing methods apart from the proposed method (NCS and LRN) evaluate their approaches on a different set of tasks to the one used in this work. It is unclear why the authors chose to deviate from this choice of tasks. NCS [12] may be the SOTA and the most recently published, so it would be beneficial to compare to NCS under their settings.\n- Incomprehensive ablations. The proposed method seems to be rather complex, and it is not convincing that such a complex method is needed. While ablations have been provided for the utility of each component at a higher level of abstraction, more ablations on the design choices for each module (factored class template graph, dynamic interaction graph, factored interaction pattern) would be beneficial. For instance, how critical is the action binding in the dynamic interaction graph? How important are the design choices in the dynamic action binding network, e.g. the choice of using soft-attention networks?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for the suggestions and feedback. We have run several experiments to answer your questions and we include a complete version of these analyses in the final version. \n\n**Effect of imagination on baselines**\n\nWe conducted an additional ablation study without the imagination component for each of the baselines. Our findings in Table 2 of the rebuttal pdf in the general response, indicate that removing the imagination slightly improves the performance of SMORL and SRICS (still not better than our method DAFT-RL), but worsens the performance of the other baselines, e.g. NCS.\n\n**Good baseline performance during training without latent parameters or skill generalization**\n\nIn simpler settings without latent parameters or skills generalization, some baselines like SMORL and NCS indeed showcase strong performances, in certain instances surpassing our method. On the other hand, the goal of our method is to improve compositional generalization capabilities, ensuring that learned policies can be applied across various skills and latent parameters. \n\n**Capacity of baselines (e.g. SMORL, NCS) w.r.t. to DAFT-RL**\n\nFor the RL part, the number of parameters is almost the same, while for the model estimation part we use roughly 2-2.5x parameters compared to SMORL and NCS. We will add a detailed comparison in the appendix in the final version.\n\n**Ablation study in which we use the same (learned or ground truth) $ \\mathbf{z} $**\n\nWe conducted two additional experiments: (1) adding the ground truth latent parameters $ \\mathbf{z} $ to the baselines; (2) adding the parameters $ \\mathbf{z} $learned by DAFT-RL to the baselines. The results for the 8-block stacking task are in Table 2 in rebuttal pdf in the general response. They indicate that adding the ground truth $ \\mathbf{z} $ to the baseline models unsurprisingly benefits the performance of all methods. Moreover, using the learned $ \\mathbf{z} $ (extracted via DAFT-RL) improves the performance for all of the baselines. This is testament to the quality of the latent parameters derived by DAFT-RL. Additionally, when we compare with the baselines with access to the true latent parameters, the DAFT-RL approach with the learned $ \\mathbf{z} $ still manages to outperform them. This underlines the effectiveness of the proposed framework in not just extracting meaningful latent parameters, but also leveraging them effectively in the task.\n\n**Choice of tasks - comparing with NCS under their settings**\n\nWe run the comparisons with NCS on its original tasks “Robogym” and report the results below. We plan to report the complete results in the final version of the paper.\n\nWe show that DAFT-RL has comparable performance to NCS in these settings. The preliminary results are given below (for all cells, the first one is ours and the second is NCS, and the number in parentheses is the std):\n\n|  | 4 | 5 | 6 | 7 |\n|---|---|---|---|---|\n| Robogym (complete) | 0.60 (0.03)/0.64 (0.01) | 0.45 (0.02)/0.47 (0.01) | 0.46(0.04)/0.49 (0.01) | 0.38(0.03)/0.41 (0.01) |\n| Robogym (partial) | 0.43 (0.04)/0.47 (0.01) | 0.32(0.03)/0.33 (0.01) | 0.25 (0.04)/0.27 (0.01) | 0.19(0.04)/0.22 (0.01) |\n\nWe also wanted to point out that in our setting, we wanted to explore objects of different classes (or types in NCS) with different types of attributes (e.g. boxes have position and velocity, switches have position and activation), factored class template graphs and interaction patterns, as well as the effect of latent parameters like mass and friction. These settings therefore extend the tasks considered by NCS.\n\n**Ablations for more design choices, e.g. action binding - soft attention networks**\n\nWe conducted additional ablations to compare the effectiveness of soft attention networks against hard attention ones [1] for action binding. We also consider the case where there is no action binding in the system, which means the action will affect all objects. As captured in Case VI and VII of Fig. 2 in our rebuttal pdf in the general response, these studies underscore that action binding is indispensable for an efficient world model. Moreover, our empirical results clearly favor soft attention networks over their hard counterparts. For the other mentioned modules, our original submission has methodically outlined their ablation studies. For ease of reference, and to illustrate the method's robustness across scenarios, we have encapsulated the comprehensive findings pertaining to the 3-Push+3-Switch task in Fig. 2 of the rebuttal pdf.\n\n[1] Biza, Ondrej, et al. \"Binding actions to objects in world models.\" ICLR 2022 workshop on Objects, Structure and Causality\n\n**Why do dynamic interactions help so much, even for two boxes and two switches**\n\nThe interaction graph plays an important role in the model. Even with only two boxes and two switches, the six potential interactions can affect the state transition dynamics substantially, e.g. consider a collision between boxes. As expected and confirmed by our results in Fig A3 in Appendix D.3.2, the absence of an interaction graph leads to a more significant drop in performance as the number of objects increases (and as the number of symmetric interactions grows as n(n-1)/2).\n\n**End-to-end approach**\n\nWe performed an additional ablation with an end-to-end version of our method, which combines Step 1, 2.1 & 2.2 into a single step and uses $\\mathcal{D}^\\text{multi}$ as the training set. We show the results for this method and baselines in Table 1 in the pdf of the general response. The end-to-end method is slightly worse than our multi-phase approach, but it is still comparable to the best baseline (NCS) across all tasks, as shown by the average. Notably, our end-to-end version is better than all other baselines on the more complex tasks like 2 Push+2 Switch (combination of skills S + changing latent parameters L) and 3 Push+3 Switch (S+L+ changing number of objects O)."}, {"review_id": "tix3Lq74CX", "reviewer": "Reviewer_H9zq", "summary": "This paper proposes a new relational world model which is at an attribute level. Based on the state knowledge from the environment or unsupervised object-centric representation (OCR) learning [1,2], they collected attribute-level object-centric knowledge, thereby, based on them, they learned a relational world model. In the aspect that their model can infer the detailed relationship between action and the specific attributes of objects (e.g., only change on the object position through “push” action, not changing the object’s color or shape) or the objects, their world model is more compositional and can be well generalized. They evaluated it through diverse generalization experiments on diverse three benchmarks, some of which give the ground-truth state information and others only give visual inputs.\n\n\n[1] SM Eslami, Nicolas Heess, Theophane Weber, Yuval Tassa, David Szepesvari, Geoffrey E Hinton, et al. Attend, infer, repeat: Fast scene understanding with generative models. Advances in neural information processing systems, 29, 2016.\n\n[2] Francesco Locatello, Dirk Weissenborn, Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. Object-centric learning with slot attention. Advances in Neural Information Processing Systems, 33:11525–11538, 2020.\n", "questions": "- In the ablation study (D.3.2 in Appendix), how did you represent the interaction except for the class template or dynamic interaction graphs? For example, when the action is to push the object, then how is the action represented? Can we see the ablation without the class template graph that is under the assumption that every object is in a single class?\n- What is TFD in line 852?\n- In [1], the Transformer encoder (similar to the self-attention that you evaluated, but using a learnable token for policy) shows good performances for the Image benchmark. Could you compare it with your model?\n- You utilized AIR [2] for the approaches relying on symbolic inputs. Could you test them with SA [3], not AIR? A few tasks are enough to evaluate due to the time limitation for the rebuttal.\n\n[1] Jaesik Yoon, Yi-Fu Wu, Heechul Bae, and Sungjin Ahn. An investigation into pre-training object-centric representations for reinforcement learning. International Conference on Machine Learning (ICML), 2023.\n\n[2] SM Eslami, Nicolas Heess, Theophane Weber, Yuval Tassa, David Szepesvari, Geoffrey E Hinton, et al. Attend, infer, repeat: Fast scene understanding with generative models. Advances in neural information processing systems, 29, 2016.\n\n[3] Francesco Locatello, Dirk Weissenborn, Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. Object-centric learning with slot attention. Advances in Neural Information Processing Systems, 33:11525–11538, 2020.", "limitations": "The authors didn't discuss their model's limitations. As I roughly discuss it, their modeling requires a class template that can restrict the application of their modeling, not just for collecting the class label, but also for their assumption for the class that in the class, the attributes and interaction patterns are shared. Another limitation is that it requires the encoder or environment must give the attribute-level representations. Because their model utilizes them as an input and the target of their objective, one-step prediction. Except for those limitations on their design, there is no ethical issue with their work I think, because it is fundamental research work rather than applications that can affect the real world. Additionally, they didn't discuss the computational overhead of their training or progressing their work.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- This paper proposes a new attribute-level factorized world model for multi-object RL tasks. It is an excellent approach to step deeper from an object-wise world model such as C-SWM [1]. Through this, they can extend the compositionality to the attribute level.\n- They well discussed which model they designed and how to learn the learnable parameters to fit. \n- They evaluated their modeling with diverse world model methods and prior works which studied policy learning with unsupervised object-centric representation learning.\n- Through their evaluation, they can find their model outperforms baselines, significantly, when generalization is more required. \n- They also studied ablations to analyze which components are more important for the performance.\n\n[1] Kipf, Thomas, Elise Van der Pol, and Max Welling. \"Contrastive learning of structured world models.\" arXiv preprint arXiv:1911.12247 (2019).", "weaknesses": "- Their model requires the class label to learn the class template graph through multiple objects in the class. It is reasonable to learn the template graph per class and it makes sense to reduce the complexity of the modeling and improve the performance by merging to learn a graph with multiple objects in the same class. However, it is an unrealistic assumption for real-world data. For example, if we applied this model to real-world data, then we should require additional labor to do labeling for each object.\n- Their model is assumed that the attribute can be accessible. Similar to the class label, it could be hard to be assumed now. For example, for AIR [1], it could be reasonable, because the AIR encoder returns position and size separately, but it is not for recent OCR models such as SA [2] or SLATE [3]. Recently, a few works have tried to represent the OCR through disentangled representation vectors [4], so we can expect that the attribute (maybe) can be given through the encoders. It is curious how you applied SA to your model.\n- Defining the reward function as an object level is too limited to cover general tasks. Sometimes the reward could be defined through the relationship between objects such as the object comparison task in the Image benchmark. It is curious how you made your model solve the task. In your modeling, reward estimation is learned through step 1 (Class Learning in single-object environments).\n- Action binding is unclearly discussed. In the section for DAFT-MDP, you assumed that an action has only an effect on one object at a time (line 74). However, action binding can affect multiple objects through soft attention (section 3.2.1), and I cannot find the details to regulate the action effect to only a single object from the detailed criteria (line 690 in Appendix). \n- Fine-tuning to a new multi-object domain is unclearly discussed. Which parameters are fine-tuned? Does Fine-tuning mean the parameters in the world model are frozen and only learning the parameters for the policy?\n- (minor) typos\n\n    - In line 222, $i$ should be $m$.\n    - In line 235, $o_i^t$ should be $o_1^t$.\n    - In line 236, isn't the value $v^t = <f_v(o_1^t),...,f_v(o_m^t)>$? \n    - In line 257, the input, isn't it $h^{t+1}_{(i,j)}$?\n    - PPO is not MBRL, but it is mentioned as one of them in line 293.\n\n\n[1] SM Eslami, Nicolas Heess, Theophane Weber, Yuval Tassa, David Szepesvari, Geoffrey E Hinton, et al. Attend, infer, repeat: Fast scene understanding with generative models. Advances in neural information processing systems, 29, 2016.\n\n[2] Francesco Locatello, Dirk Weissenborn, Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. Object-centric learning with slot attention. Advances in Neural Information Processing Systems, 33:11525–11538, 2020.\n\n[3] Singh, Gautam, Fei Deng, and Sungjin Ahn. \"Illiterate dall-e learns to compose.\" arXiv preprint arXiv:2110.11405 (2021).\n\n[4] Singh, Gautam, Yeongbin Kim, and Sungjin Ahn. \"Neural systematic binder.\" The Eleventh International Conference on Learning Representations. 2022.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your feedback, it will surely improve our paper. We will answer your concerns in the following.\n\n**Requiring the class label**\n\nWe first wanted to clarify a potential misunderstanding: we do not need labels for each new object, but we need to train a classifier that can classify the object’s class based on its extracted features, as explained in Appendix C.1.2. We will clarify the explanation in the main part of the paper. On the other hand, our current method would not work if an object of a new class appears. In that case, one solution could be to use off-the-shelf open-set recognition methods [1] to detect and categorize novel classes, followed by learning the template graph for this new class using collected trajectories. This is an interesting research direction and we intend to explore solutions to this limitation in subsequent research.\n\n[1] Geng et al. \"Recent advances in open set recognition: A survey.\" IEEE TPMAI (2020).\n\n**Assuming that the attribute is accessible**\n\nWe again wanted to clarify a potential misunderstanding: the attributes in this work are not predefined and we can use pre-trained object-centric models like SA to extract them from visual inputs. This means that there may not be a direct mapping between the extracted features and ground truth attributes, or a disentangled representation of the ground truth attributes when using SA. While using such encoders might lead to less interpretable class template graphs or interaction patterns graphs, our empirical results show that our framework still provides performance benefits. The direction you mentioned about disentangled representation vectors, as in the work of Singh et al. [4], is indeed promising and we will consider it in future work.\n\n**The reward could be defined through the relationship between objects**\n\nThat’s a fair point. In this case, we can model the reward function as being affected by all objects. For the object comparison task in the Image benchmark, we indeed do not learn the sparse connections between the attributes and the reward in the class template graphs and instead assume that they are all fully connected. We will clarify this special case in the final version. \n\n**Unclear discussion of action binding**\n\nThanks for your careful reading and feedback. While it is accurate that soft attention can distribute action effects across multiple objects, the real-world scenarios we considered and the graphs we learned from data show that in most of cases, only one object is affected significantly by a particular action. The subtle diffusion of effect to other objects via soft attention is generally negligible. Our choice for soft attention is backed by its inherent flexibility, demonstrated by previous work [1], especially in robotic manipulation tasks. To further elucidate the effectiveness of soft attention, we conducted a comparative study with hard attention networks, as described in [1]. The results are in Fig. 2 VIII of the rebuttal pdf in the general response, demonstrating that soft attention outperforms its hard counterpart. We hope this clears up the ambiguity surrounding action binding with soft attention in our work.\n\n[1] Biza, Ondrej, et al. \"Binding actions to objects in world models.\" ICLR 2022 workshop on Objects, Structure and Causality\n\n**Fine-tuning to a new multi-object domain**\n\nThe parameters of the world model are indeed frozen, and we indeed only need to learn the parameters for the policy. Just to be completely clear, this does not mean that the world model itself is frozen, since the objects are not frozen and we need to infer their latent parameters, although the class template graphs and interaction graphs are the same. \n\n**In ablations in D.3.2 how did you represent the interaction except for class template or dynamic interaction graphs** \n\nIn the absence of the class template or dynamic interaction graphs, the interaction representation assumes a dense configuration. That is, every pair of objects interacts with each other, and this interaction remains consistent throughout the temporal spectrum. For situations without action binding, as illustrated in Case VI of Fig. 2 in our rebuttal pdf in the general response, the action node (e.g., push) is uniformly linked to all objects in the scene, implying that the action has the potential to affect every object.\n\n**Ablation with single class**\n\nWe have conducted this ablation for the 3-Push+3-Switch (L+O+S) in Fig. 2 in the rebuttal pdf in the general response. The results, presented in Case V, show a considerable decline in performance, reinforcing the importance of class-specific distinctions in our model. We will include these results in the final version. \n\n**TFD in Line 852**\n\nThanks for the pointer. It means the transformer decoder. We follow this design (introduced in NCS) only for Spriteworld and will explain this in detail in the final appendix.\n\n**Comparing with Transformer encoder [Yoon et al 2023]**\n\nThanks for the suggestion, we tried to run this experiment but did not finish in time for the rebuttal. We will include it in the final version.\n\n**Using SA for approaches with symbolic inputs**\n\nWe conducted preliminary tests using SA for the symbolic input-based approaches, particularly on the stacking benchmark with 8 blocks. This involved integrating SA with the SRICS, GNN, and LRN. From our experiments, we observed that models leveraging SA yielded slightly worse results to those employing AIR, except for LRN . Specifically, you can refer to the last row of Table 2 in the rebuttal pdf in the general response. The performance for using AIR is provided in the first row of the same table. We will expand our evaluation using SA across other tasks and benchmarks in the final paper.\n\n**Limitations and overhead**\n\nWe will improve the discussion of the limitations of the method in the main paper, including the assumptions on the classes and attributes. We will also discuss the computational cost of our method."}]

Author rebuttal:
We thank all reviewers for the insightful feedback. Your suggestions and questions for more experiments will improve the quality of the paper, while the misunderstandings will improve its clarity and readability, making it more accessible to a wider audience.

We are happy about the positive feedback, mentioning that our paper proposes a “novel approach to an interesting problem” and “an excellent approach to step deeper from an object-wise world model”, that the design decisions are “well-motivated and the overall idea is well explained” and that the paper provides “strong experimental results”, which show that “the performance improvements are impressive”. 

We are also thrilled that some reviewers found it an “overall well written paper” and “easy to follow and the model section is well written”, and we will strive to improve its clarity in the final version also for the other reviewers and any future reader. We will also improve the discussion of the limitations of our method, as suggested by several reviewers.

We wanted to clarify a few common points (which we will also clarify in the paper): we do not need predefined attributes for our method or labels for each new object. When combined with slot attention, our method automatically extracts the attributes from the representations. We also do not need labels for each new object, but we need to train a classifier that can classify the object based on its extracted features, as explained in Appendix C.1.2. On the other hand, in our current approach, we have not explicitly addressed situations where the classes are unknown, which is an interesting future direction.

In the rebuttal pdf, we provide several ablations, which we describe in the following.

**Table 1 - End-to-end** 

Several reviewers were interested in an end-to-end version of our method. We provide results in Table 1. DAFT-RL (End-to-end) combines Step 1, 2.1 & 2.2 into a single step and uses $\mathcal{D}^\text{multi}$ as the training set, so it does not require datasets with single objects. The end-to-end method is slightly worse than our multi-phase approach, but it is still comparable to the best baseline (NCS) across all tasks, as shown by the average on the last row. Notably, our end-to-end version is better than all other baselines on the more complex tasks like 2 Push+2 Switch (S+L) and 3 Push+3 Switch (S+L+O).

**Table 2 - Imagination, true/learned parameters, SA** 

We run several ablations on block stacking with 8 blocks, in particular:
- We switched off the imagination component for each baselines. Our findings in Table 2 of the rebuttal pdf in the general response, indicate that removing the imagination slightly improves the performance of SMORL and SRICS (still not better than our method DAFT-RL), but worsens the performance of the other baselines, e.g. NCS;
- We used the latent parameters learned by DAFT-RL to the baselines, which also improves the performance of the baselines (but still not better than our method) and shows the quality of the latent parameters that DAFT-RL learns;
- We used the ground truth latent parameters as inputs to the baselines, which unsurprisingly benefits all methods. Interestingly, the DAFT-RL method without access to the ground truth still outperforms the baselines with access to the ground truth. This underlines the effectiveness of our framework in not just extracting meaningful latent parameters, but also leveraging them effectively in the task.
- We used Slot Attention for the symbolic input-based approaches, SRICS, GNN, and LRN and observed that models leveraging SA yielded slightly worse results to those employing AIR, except for LRN.

**Fig. 1 - Quality of learned graphs**

We investigated the quality of learned graphs and latent parameters, and their effect on the RL performance w.r.t. the number of samples for the 3-Push + 3-Switch (L+O+S) task. In particular, we varied the amount of training data, ranging from {10%, 20%, 40%, 60%, 80%} of the original sample size (900 trajectories). We measured:
- the $R^2$ coefficient of our learned parameters with the true latent parameters, both for a random policy (Fig. 1a) and for a pretrained policy (Fig. 1b). These results show that the performance degrades with smaller sample sizes, but it is still acceptable with 60% data, and that the difference between data collected with a random or pretrained policy is negligible at higher sample sizes; 
- the normalized Structural Hamming Distance between the reconstructed graph and the ground truth graph, as well as the success rate of the learned policy. As expected, the more data, the more accurate the graph, and the better the performance of the policy trained with the more accurate world model. For the number of samples in the main paper, the graph is perfectly reconstructed, which also means a good RL performance. Additionally, for limited data (e.g. 0.2 or 0.4 of the original dataset) leveraging a pre-trained policy enhances both graph and policy learning. However, as the amount of data increases, the benefits of using a pre-trained policy diminish.

**Fig. 2 - Single class, action binding design choices**

We report the ablation studies from Appendix D.3.2 for 3-Push+3-Switch (L+O+S), improve the naming for I-IV to emphasize what the ablations are doing, and extend them with three additional ablations. All of these ablations show that our design choices have a positive effect on our method. In particular:
–  “V: DAFT-RL with a single class” represents the case of using a single class for this environment, which seems to have the strongest negative effect; 
– “VI. DAFT-RL with dense action binding” represents the case in which the action will affect all objects, which also substantially lowers the performance of our method;
– “VII. DAFT-RL with action binding using hard attention network” represents a design choice in terms of hard attention vs soft attention, which also reduces the performances.


Meta-review: The paper proposes a novel object-attribute centric World Model in RL. It has received 1 x Accept, 2 x Weak Accepts, 2 x Borderline Accept, 1 x Borderline Reject. All but one of the reviewer feel very positive about the paper. In particular, akCD, NUSp point out the ability of the approach to generalize to number of objects not seen during training; akCD, xWTh, H9zq point out the strength of the experiments. The only negative review points out limited novelty (a marginal extension of object-oriented MDP and simple evaluation. After the rebuttal, the reviewer raises their score to BR but remains unconvinced. Hence, the paper is accepted to Neurips.

The authors are encouraged to add the requested explanation (clarifications, hyper parameters, figure improvements) and discuss some of the limitations of the model.
Decision: Accept (poster)