Title: Learning Dynamic Attribute-factored World Models for Efficient Multi-object Reinforcement Learning
Venue: NeurIPS 2023
Accepted: yes
Keywords: multi-object RL

compositional generalization

factored representations
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-factored World Models for Efficient Multi-object Reinforcement Learning

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.

1 Introduction

Model-based reinforcement learning (MBRL) and world models [1-3] have demonstrated improved performance in many RL 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 realworld tasks, we often have environments with multiple objects that interact, and we are interested in generalizing to unseen combinations and numbers of objects. In recent years, there have been several studies exploring and learning object-oriented environment models or policy models [4-17] and tackling the problem of compositional or combinatorial generalization, in which we consider combining the modeling of multiple objects or tasks to solve a new task.

Although these methods have effectively leveraged object-centric and object-factored representations in RL 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 a sparse subset of attributes, e.g. an 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

often sparse, both in time and in terms of which attributes are affected, e.g. the position of a box is 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 dynamics, 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 FacTored RL (DAFT-RL), a framework that learns a 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 [18-23], Relational MDPs [24-26] and especially Object-Oriented (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-art 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 to learn a policy on a 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 it outperforms the baselines.

2 Dynamic Attribute-FacTored 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 ObjectOriented (PO)MDPs [27-29], but we propose a more 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 a 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 (e.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 a given attribute only a sparse subset of other attributes influences these functions.

In our setting, objects can interact with each other, which might influence 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 a formal definition of a Dynamic Attribute-FacTored MDP.

Class system, attributes, and objects. We assume a known set of classes C = { C 1 , . . . , C k } of objects, where each class C j describes a set of observable attributes { C j .s 1 , . . . , C j .s n } , 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 j . θ , 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 . color , BOX . pos , BOX . vel , BOX . θ } representing the object color, position, velocity, and friction coefficient of a box, and { SW . color , SW . pos , SW . active } , representing the color, position, and activation state of a switch. The class system specifies a template for a set of environments that can

C. Interaction Graph

A. Class Template Graphs

B. Interaction Pattern Graph

o t

1

o t

3

o t

2

C

1

.

s

t

2

C

1

.

s

t

+1

1

C

1

.

s

t

1

a

t

r

t

C

1

.

s

t

3

C

1

.

s

t

+1

2

C

1

.

s

t

+1

3

C

1

.

θ

C

2

.

s

t

+1

1

C

2

.

s

t

1

C

2

.

s

t

3

C

2

.

s

t

+1

2

C

2

.

s

t

+1

3

C

2

.

θ

C

2

.

s

t

2

a

t

r

t

o t

+1

1

o t

+1

2

o t

+2

3

o t

+2

2

α t

+1

a

t

+1

α t

a

t

o t

+2

1

C

1

.

s

t

1

C

1

.

s

t

+1

1

C

1

.

s

t

2

C

1

.

s

t

3

C

1

.

s

t

+1

2

C

1

.

s

t

+1

3

C

1

.

s

t

1

C

1

.

s

t

2

C

1

.

s

t

3

C

1

.

s

t

+1

1

C

1

.

s

t

+1

2

C

1

.

s

t

+1

3

o t

+1

3

C

1

.

θ

C

1

.

θ

vary in the number and characteristics of the objects. Each environment has a fixed set of objects O = { o 1 , . . . , o m } , where each object is an instance of a class in C , which we denote as o i ∈ C j for i ∈ { 1 , . . . , m } , j ∈ { 1 , . . . , k } . We denote the class of an object o i with C ( i ) . For example, using the class system from the previous example, we can represent a source environment as O ′ = { o 1 , o 2 } where o 1 ∈ BOX and o 2 ∈ SW, and a target environment as O ′′ = { o 1 , o 2 , o 3 } , where we add o 3 ∈ BOX to the original objects. For each object o i ∈ C j and timestep t = { 1 , . . . , T } , we denote its attributes at time t as o t i = { o i .s t 1 , . . . , o i .s t n } , which are instantiations of the class attributes { C j .s 1 , . . . , C j .s n } , and its constant parameters as o i . θ . In our example, for box o 1 the attributes o 1 .s t 1 , o 1 .s t 2 and o 1 . θ are its position and velocity at time t and its friction coefficient, while for the switch o 2 the attributes o 2 .s t 1 and o 2 .s t 2 are its position and activation at time t .

States, actions, transitions, and rewards. We define the state at time t as the collection of all the object states, i.e., s t = { o t 1 , . . . , o t m } with domain S . We collect all object-specific latent parameters in a global parameter θ = { o 1 . θ , . . . , o m . θ } with domain Θ . We define the space of actions A and use a t ∈ A to denote the action at time t . We denote the transition probability as p ( s t + 1 | s t , a t , θ ) . The reward at time t is denoted as r t and the reward probability is denoted as p ( r t | s t , a t , θ ) . 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 j , we learn a DBN G C j over the nodes { C j .s t 1 , . . . , C j .s t n , C j .s t +1 1 , . . . , C j .s t +1 n , C j . θ , a t , r t } that represents a template for the instance graph between the attributes for each object of this class G t o i for o i ∈ C j . In particular, the edges between the action a t and the attributes at timestep t + 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. 1a. In this example, C 1 .s t 2 influences C 1 .s t +1 2 and C 1 .s t +1 3 , but it's not influenced by C 1 .s t -1 3 . The reward r t is only influenced by C 1 .s t 3 and the action will only have an effect on C 1 .s t +1 2 and C 1 .s t +1 3 . Moreover, C 2 .s t 2 influences only C 2 .s t +1 2 and r t .

Interaction pattern graphs (Fig. 1B). 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 an attribute perspective, the interplay between objects during interaction is also factored

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 a 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 i and C j (possibly also with i = j ), we learn a DBN G C i ,C j over the nodes {{ C i .s t l } n l =1 , { C i .s t +1 l } n l =1 , C i . θ , { C j .s t l } n l =1 , { C j .s t +1 l } n l =1 , C j . θ } . We show an example of a DBN describing how boxes interact with other boxes in Fig. 1b. In this case, the interaction is limited to C 1 .s t 2 from one object influencing the C 1 .s t +1 1 and C 1 .s t +1 2 from the other object, and the latent parameters C 1 . θ from each object influencing C 1 .s t +1 3 from the other object. While the interaction patterns are time-invariant, we assume that for each pair of objects, the interactions are only switched on at some points in time, as we will now describe through the interaction graphs.

Dynamic interaction graph (Fig. 1C). The class template and interaction pattern graphs that we just described can model the general behavior of the classes of objects in a static, time-invariant way. On the other hand, in many multi-object environments object interactions occur sparsely, the pairs of interacted objects are not fixed, and the action has an effect only on a limited number of 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 G inter = {G t inter } T t =1 at the object level, which is a sequence of graphs G t inter with edges from a subset of { o t 1 , . . . , o t m , a t } to a subset of { o t +1 1 , . . . , o t +1 m } . Each edge o t i → o t +1 j represents an interaction between an object o i and an object o j . This interaction is instantiated in an instance interaction graph G t o i ,o j , following the interaction pattern graph G C i ,C j for each pair of objects o i ∈ C i and o j ∈ C j .

We also learn an object selector α t ∈ { 1 , . . . , m } that selects the single object on which the action a t has an effect at timestep t , which we represent with an edge from a t to the selected object o t α t in G t inter . In particular, this selection (or action binding ) activates the edges from a t to the object attributes described in the class template graph G C j in the instantiated version for object o i ∈ C j , which we denote G t o i . The graph G inter is dynamic because the edges of each graph in the sequence can change at each timestep t = 1 , . . . , T .

We show an example of interactions between three objects at different timesteps in Fig. 1C. In this example, at timestep t there is an interaction between o 1 and o 3 , which will follow the interaction pattern presented in Fig. 1B, i.e. o 1 .s t 2 →{ o 3 .s t +1 1 , o 3 .s t +1 2 } and o 1 . θ t → o 3 .s t +1 3 , and viceversa for o 3 towards o 1 . Moreover, a t has an effect on o 2 , and specifically o 2 .s t +1 3 , following the class template graph presented in Fig. 1A. Instead, at timestep t +1 , there are no interactions between the objects and a t +1 has an effect on o 1 , specifically o 1 .s t +2 2 and o 1 .s t +2 3 . For completeness, in App. B we provide an example of this graph combined with the instantiated class template graphs for each object G o 1 , G o 2 and G o 3 for three timesteps t, t + 1 , t + 2 , as well as the instantiated interaction pattern graph G o 1 ,o 3 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-FacTored Markov Decision Process (DAFT-MDP). We will assume that the class system is fixed and that the objects can vary in each environment, as can their interaction graphs. Under this assumption, a DAFT-MDP defines a family of MDPs, that is parametrized in the objects O and their dynamic interaction graph G inter .

Definition 1 (DAFT-MDP) . A Dynamic Attribute-FacTored Markov Decision Process (DAFT-MDP) is a tuple ( C , O , Θ , A , G , P s , R , γ ) , where C is the set of classes, O is the set of objects, Θ is the space of the constant latent parameters, A the action space, P s 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 {G C j } C j ∈C , the collection of interaction pattern graphs for each pair of classes {G C i ,C j } C i ,C j ∈C and the dynamic interaction graph G inter , as defined previously. These graphs define the factorization of the transition distribution per object and per attribute, as follows:

where s t is the collection of all attributes of all objects at time t , θ ∈ Θ is the collection of all latent constant parameters for all objects, a t ∈ A is the action. In the right-hand term, o i .s t +1 l is attribute s l of object o i at time t +1 , while pa G t o i ( o i .s t +1 l ) are the parents of the attribute l for

object o i based on the class template graph G C ( i ) , where C ( i ) is the class of o i , and where the action binding α t activates any potential connections from a t . In the second term of the conditioning, we iterate over the objects o k that are interacting with o i at time t in the dynamic interaction graph o k → o i ∈ G t inter . For each of these objects o k we collect the attributes that interact with o i .s l in the instance interaction pattern G t o i ,o k based on interaction pattern graph G C ( i ) ,C ( k ) for the respective classes C ( i ) and C ( k ) . Similarly, we define the factorization of the reward function per object and per attribute as R ( s t , a t , θ ) = R ( { pa G o i ( r t ) } o i ∈O ) , where for each object o i we collect all the attributes that have an edge to 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 θ . In the training phase, we will learn how to classify objects, the transition and reward functions based on the class template graphs {G C j } C j ∈C and the interaction patterns {G C i ,C j } C i ,C j ∈C . In the testing phase, we will infer the class and latent parameters of each object, as well as the interactions between the objects in the dynamic interaction graph G t inter , which specify the transition and reward functions in the new environment.

3 The DAFT-RL Framework

In the previous section, we introduced our model, DAFT-MDP (Fig. 2). In this section we provide a framework for estimating DAFT-MDPs, leveraging them for policy learning in a set of source environments, and adapting the policy to a 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 multiobject environments, (iii) policy learning and imagination in multi-object environments, and finally (iv) adaptation to a new multi-object environment. We present each step and its implementation in the following.

̂

Step 1: Class learning in single-object environments

{(

x

t

1

,

a t

1

,

r

t

1

)}

T

t

=1

{(

x

t

2

,

a t

2

,

r

t

2

)}

T

t

=1

{(

x

t

3

,

a t

3

,

r

t

3

)}

T

t

=1

Object-centric model

{

o

t

1

,

a t

1

,

r

t

1

}

T

t

=1

{

o

t

2

,

a t

2

,

r

t

2

}

T

t

=1

{

o

t

3

,

a t

3

,

r

t

3

}

T

t

=1

Learning class template graph

Object class inference

/u1D4A2

C

1

, /u1D4A2

C

2

, /u1D4A2

C

3

C

1

,

C

2

,

C

3

Stage 2: Interaction Learning and Latents Inference

Stage 3: Policy Learning

Stage 4: Adaptation

Apply policy

Inferred constant latent parameters

o

i

.

θ

Inferred object inertaction graph

/u1D4A2

inter

π

*(

a

t

|

s

t

,

θ

, /u1D4A2

inter

)

Learn the object interaction graph

Learn the interaction pattern graphs

Learn the latent constant parameters

Learn the object selector

/u1D4A2

inter

o

1

.

θ

,

o

2

.

θ

,

o

3

.

θ

α t

/u1D4A2

Ci

→

Cj

,

i

∈ {1,2,3},

j

∈ {1,2,3}

{

o

t

1

,

o

t

2

,

o

t

3

,

a t

,

r

t

)}

T

t

=1

{(

x

t

,

a t

,

r

t

)}

T

t

=1

Object-centric model

{(

x

t

,

a t

,

r

t

,)}

T

t

=1

{

o

t

1

,

o

t

2

,

o

t

3

,

o

t

4

,

a t

,

r

t

)}

T

t

=1

Object-centric model

{(

x

t

,

a t

,

r

t

)}

T

t

=1

{

o

t

1

,

o

t

2

,

o

t

3

,

o

t

4

,

o

t

5

,

a t

,

r

t

)}

T

t

=1

Object-centric model

Generate imagined trajectories

{(

o

t

1

,

o

t

2

,

o

t

3

,

o

4

,

a t

,

r

t

)}

T

2

t

=

T

1

Learned policy

π

*(

a

t

|

s

t

,

θ

, /u1D4A2

inter

)

{(

x

t

,

a t

,

r

t

)}

T

t

=1

{(

x

t

,

a t

,

r

t

)}

T

t

=1

̂

̂

̂

̂

̂

In all steps and all environments, if the input is an image, we extract the objects and their attributes s t = { o t 1 , o t 2 , . . . , o t m } from sequences of images x t 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 t = { o t 1 , o t 2 , . . . , o t m } . 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.

3.1 Step 1: Class Learning in Single-Object Environments

In this step we consider data from m single-object environments with different objects and no agent interaction. In particular, for each class C j we collect the transitions for several objects o i ∈ C j as { o t i , a t i , r t i } T t =1 , in environments in which there is only object o i , and a random policy is used to generate actions. We denote these data as D single = {{ o t i , a t i , r t i } T t =1 } i ∈ C j , ∀ j =1 ,...,k,i =1 ,...,m .

We initiate the class template G C j for each class C j randomly, and then use D single 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 p λ ( o t +1 i , r t i | o t i , a t i , G C j ) , where λ = { λ s , λ r } 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 t for each object i with class C j , the inputs to the GRU are { pa G Cj ( o t +1 i ) } i m =1 and { pa G Cj ( r t i ) } i m =1 , and the GRU outputs o t +1 i and r t i . The learning objective of this step is given below, where we maximize the log-likelihoods of dynamics and reward models and regularize the graph to be sparse:

where m and k indicate the number of single-object environments and object type classes, respectively, and pa G C ( i ) denotes the parents of a variable in the template class graph for the class C ( i ) of object o i . After this step, we fix the learned {G C 1 , G C 2 , . . . , G C k } with the exception of the edges from the latent parameters θ , which here we assume are disabled and we will learn in the next step. In later stages, we will reuse the learned reward model λ r and the class template graphs.

3.2 Step 2: Interaction Learning and Latent Parameter Inference in Multi-object Environments

In this step, we consider data D multi from N multi-object environments with different object configurations and in which the objects can have varying latent parameters. Formally, we define D multi = {{ o t 1 , a t 1 , r t 1 , o t 2 , a t 2 , r t 2 , . . . , o t m , a t m , r t m } T t =1 } N m =1 . In each of these environments we assume the agent can interact only with one object at a time. On these data, we again extract the objects and their attributes from a 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 G C i ,C j for each pair of classes C i and C j and the dynamic interaction graph G inter by exploiting the previously learned class template graphs. In particular, we first learn the action binding α t , and at a second stage, we jointly learn the rest of the dynamic interaction graph G inter , the interaction patterns G C i ,C j for each pair of classes, the object-specific latent parameters o i . θ and their edges to the other attributes in G C j . We describe these two stages in detail in the following.

3.2.1 Step 2.1: Learning the Action Binding

Motivated by [35], we learn the dynamic action binding α = { α 1 , α 2 , . . . , α T } 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) to derive the key k t = ⟨ f k ( o t 1 ) , f k ( o t 2 ) , . . . , f k ( o t m ) ⟩ , query q t = f q ( a t ) , and value v t = ⟨ f v ( o t 1 ) , f v ( o t 2 ) , . . . , f v ( o t m ) ⟩ , respectively. We then compute the attention weights α t = softmax (( k t 1 ) ⊺ q t , ( k t 2 ) ⊺ q t , . . . , ( k t m ) ⊺ q t ) . We use the learned attention weights α 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 v t with the attention weights α t , 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 f k , f q , and f v .

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 a 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 t = GNN ( o t 1 , o t 2 , . . . , o t m ) . For each object pair o i and o j , we then obtain h t ( i,j ) . 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: h t ( i,j ) , enc = GRU enc ( h t ( i,j ) , h t +1 ( i,j ) ) . For the prior, we also use a GRU, but in this context, we append h t ( i,j ) with the output of the GRU from the previous step h t -1 ( i,j ) , prior as the input. In specific terms, we have h t ( i,j ) , prior = GRU prior ( h t ( i,j ) , h t -1 ( i,j ) , 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 q ϕ ( z t | s 1: T ) , where s 1: T = { o 1 t , . . . , o m t } T t =1 , and prior distribution p ϕ ( z t | s 1: t , z 1: t -1 ) , respectively. We assume that the encoder outputs a Bernoulli distribution for each edge and the graph G inter is sampled using the Gumbel-Softmax trick [40].

Decoder We perform one-step prediction to generate ˆ o t +1 i for each object o i , predicting the state dynamics with the learned graphs, including G inter , as well as the interaction pattern graph G C i ,C j for two objects with class C i and C j . At the same time, our goal is also to learn the interaction pattern graph with sparsity regularization. We also learn the latent parameters, θ at this stage. Specifically, we also incorporate C j . θ and the graph from G C j . θ , o t i for each object o i with class C j into the dynamics model and use the sparsity regularization for G C j . θ , o i .

Therefore, the learning objectives include maximizing the likelihood of the dynamics model, minimizing the KL divergence between p ϕ and q ϕ to estimate θ , and encouraging the sparsity of the interaction pattern graph G C i ,C j and the subgraph from latent parameters to states G C i . θ , o i :

where G is the dynamic interaction graph G inter , interaction pattern graphs {G C i ,C j | i ∈ { 1 , 2 , . . . , k } , j ∈ { 1 , 2 , . . . , k }} and the subgraph from latent parameters to states {G C j . θ , o i | i ∈ { 1 , 2 , . . . , m } , j ∈ { 1 , 2 , . . . , k }} . Similarly, to Definition 1, pa G t o i ( o i .s t +1 l ) indicates the parents of the attribute s l for object o i based on the class template graph G C ( i ) , where C ( i ) is the class of o i , and where the action binding α t activates or deactivates any potential connections from a t . In the second term of the conditioning, we iterate over the objects o k that are interacting with o i at time t in the dynamic interaction graph o k → o i ∈ G t inter . For each of these objects o k we collect the attributes that interact with o i .s l in the instance interaction pattern G t o i ,o k based on interaction pattern graph G C ( i ) ,C ( k ) for the respective classes C ( i ) and C ( k ) . λ s 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.

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-specific latent parameters θ and the interaction graph G inter . 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-specific world model to create a set of imagined trajectories. Finally, we can learn a policy π ∗ ( a t | s t , θ , G inter ) across different environments, based on the real and imagined trajectories. We can apply policy learning or planning methods using any RL algorithms. To

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 π ∗ . Detailed information about the domain parameters are provided in App. C.

3.4 Step 4: Adaptation to the New multi-Object Environment

In a new environment, we apply the policy π ∗ ( a t | s t , θ , G inter ) by inferring latent parameters θ and dynamic interaction graphs G inter based on a few trajectories, without any policy learning.

4 Related Work

We shortly summarize the related work in this section and provide a 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 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 high-dimensional raw pixels as input. NCS [12] and STEDIE [17] go further by disentangling action/control-relevant 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.

5 Experimental Evaluation

We consider a diverse set of RL benchmarks and setups, including modified OpenAI Fetch environments [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-Attention as the inductive bias [15, 4], such as SRICS [6], STOVE [8], SMORL [5], NCS [12], LRN [7], and COBRA [9]. To ensure a fair comparison, 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 to 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: OpenAI Fetch - Push and Switch. Following [15], we modify the OpenAI 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 a 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 generalization , 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 ) and skills, e.g. 3 -Push+ 3 -Switch ( 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 a bigger gain

Experiment

Settings

Method

DAFT-RL(Symbolic)

DeepSets

Self-attention

SRICS

GNN

STOVE

SMORL

NCS

LRN

2-Push + 2-Switch (S)

0

.

881

±

0

.

038

0

.

498

±

0

.

024

0

.

367

±

0

.

027

0

.

717

±

0

.

039

0

.

633

±

0

.

027

0

.

808

±

0

.

032

0

.

788

±

0

.

021

0.912

±

0.022

0

.

813

±

0

.

047

3-Push + 3-Switch (S+O)

0.805

±

0.024

0

.

324

±

0

.

012

0

.

183

±

0

.

031

0

.

685

±

0

.

019

0

.

594

±

0

.

023

0

.

672

±

0

.

021

0

.

694

±

0

.

017

0

.

751

±

0

.

024

0

.

715

±

0

.

025

2-Push (L)

0.968

±

0.036

0.940

±

0

.

029

0.941

±

0

.

046

0.922

±

0

.

042

0.953

±

0

.

037

0.931

±

0

.

047

0.964

±

0

.

023

0.944

±

0

.

039

0.923

±

0

.

038

2-Switch (L)

0.923

±

0.065

0.891

±

0

.

043

0.912

±

0

.

054

0.919

±

0

.

032

0.873

±

0

.

046

0.824

±

0

.

027

0.892

±

0

.

035

0.864

±

0

.

062

0.907

±

0

.

048

3-Push (L+O)

0.921

±

0.037

0

.

862

±

0

.

036

0

.

907

±

0

.

036

0

.

824

±

0

.

036

0

.

852

±

0

.

036

0

.

915

±

0

.

036

0

.

828

±

0

.

036

0

.

890

±

0

.

036

0

.

886

±

0

.

036

3-Switch (L+O)

0.903

±

0.023

0

.

451

±

0

.

075

0

.

392

±

0

.

045

0

.

590

±

0

.

087

0

.

744

±

0

.

057

0

.

772

±

0

.

079

0

.

835

±

0

.

029

0

.

818

±

0

.

074

0

.

857

±

0

.

032

2-Push + 2-Switch (L+S)

0.793

±

0.026

0

.

384

±

0

.

077

0

.

351

±

0

.

032

0

.

549

±

0

.

089

0

.

476

±

0

.

062

0

.

595

±

0

.

043

0

.

463

±

0

.

057

0

.

591

±

0

.

081

0

.

634

±

0

.

045

3-Push + 3-Switch (L+O+S)

0.783

±

0.025

0

.

256

±

0

.

045

0

.

115

±

0

.

029

0

.

531

±

0

.

082

0

.

362

±

0

.

035

0

.

531

±

0

.

042

0

.

525

±

0

.

051

0

.

529

±

0

.

027

0

.

418

±

0

.

092

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 -Switch (L+S) in Fig. 3A, in which we display the top three methods in terms of success rate.

Image benchmark: Spriteworld Weconsider four tasks in Spriteworld [9]: object goal, interaction, comparison, and property comparison. Given that the observations are in pixel space, we use pretrained SA [30] or AIR [31] as pre-trained encoders to obtain the object factored states. Following [13], we generate the datasets with varying numbers of objects placed randomly for pre-training the objectcentric 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.

Image benchmark: Block-stacking In the Block-stacking 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 a sanity check, we show in App. D.3 that DAFT-RL provides comparable performances to the baselines in the single-task mode . We train on a 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.

Ablation studies To evaluate the effectiveness of each component 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 labels); (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-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 importance 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

I

II

III

IV

V

VI

VII

VIII

Cases

0

.

60

0

.

65

0

.

70

0

.

75

Success Rate

3-Push + 3-Switch(L + O + S)

0

.

60

0

.

65

0

.

70

0

.

75

Success Rate

I: w/o latents

II: w/o class template graphs

III: w/o dynamic interaction graphs

IV: w/o interaction pattern graphs

V: using single class

VI: w/o action binding

VII: using hard attention

VIII: Ours

0

D

±

±

±

±

±

±

±

.

6

2

4

6

Number of Objects

DAFT (SA)

DAFT (AIR)

SMORL

SRICS

GNN

STOVE

NCS

LRN

2

4

6

8

Number of Objects

0

.

4

0

.

6

0

.

8

Success Rate

0

.

0

0

.

5

1

.

0

1

.

5

2

.

0

Time Steps (10

6

)

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

Success Rate

DAFT

STOVE

COBRA

0

1

2

3

Time Steps (10

7

)

0

.

2

0

.

4

0

.

6

0

.

8

Success Rate

DAFT

NCS

LRN

A

B

C

DAFT (SA)

SMORL

GNN

NCS

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-truth causal graphs.

6 Conclusions, Limitations, and Future Work

We proposed Dynamic Attribute FacTored RL (DAFT-RL), a framework that leverages learned attribute-factored representations with dynamic graphs. For each class of 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 a policy that can then be directly applied in a 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 our model might encounter challenges when introduced to unseen object classes. One potential area for future exploration is the integration of openset recognition [53] to identify and adapt to unseen object classes. Furthermore, we could integrate 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 AI tasks [59, 60].

Acknowledgments

FF would like to acknowledge the CityU High-Performance Computing (HPC) resources in Hong Kong SAR and LISA HPC from the SURF.nl. SM was supported by the MIT-IBM Watson AI Lab and the Air Force Office of Scientific Research under award number FA8655-22-1-7155. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the author(s) and do not necessarily reflect the views of the United States Air Force.

Success Rate

0

0

0

0

.

.

.

.

8

75

70

65

60

I

II

III

IV

V

Cases

VI

VII

References

Danijar Hafner, Timothy Lillicrap, Ian 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.

Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In International Conference on Learning Representations (ICLR) , 2020. URL https://openreview.net/forum?id=S1lOTC4tDS .

Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample-efficient world models. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id=vhFu1Acb0xb .

Richard Li, Allan Jabri, Trevor Darrell, and Pulkit Agrawal. Towards practical multi-object manipulation using relational reinforcement learning. In 2020 IEEE international conference on robotics and automation (ICRA) , pages 4051-4058. IEEE, 2020.

Andrii Zadaianchuk, Maximilian Seitzer, and Georg Martius. Self-supervised visual reinforcement learning with object-centric representations. In International Conference on Learning Representations (ICLR) , 2021. URL https://openreview.net/forum?id= xppLmXCbOw1 .

Andrii Zadaianchuk, Georg Martius, and Fanny Yang. Self-supervised reinforcement learning with independently controllable subgoals. In Conference on Robot Learning (CoRL) , pages 384394. PMLR, 2022. URL https://openreview.net/forum?id=TEQWRlncJVm .

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://openreview.net/forum?id=HFUxPr_I5ec .

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://openreview.net/forum? id=B1e-kxSKDH .

Nicholas Watters, Loic Matthey, Matko Bosnjak, Christopher P Burgess, and Alexander Lerchner. Cobra: Data-efficient model-based rl through unsupervised object discovery and curiositydriven exploration. arXiv preprint arXiv:1905.09275 , 2019.

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://openreview.net/forum?id= mLcmdlEUxy-.

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://openreview.net/forum? id=dYjH8Nv81K .

Michael Chang, Alyssa Li Dayan, Franziska Meier, Thomas L. Griffiths, Sergey Levine, and Amy Zhang. Hierarchical abstraction for combinatorial generalization in object rearrangement. In The Eleventh International Conference on Learning Representations (ICLR) , 2023. URL https://openreview.net/forum?id=fGG6vHp3W9W .

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.

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.

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.

Ondrej Biza, Thomas Kipf, David Klee, Robert Platt, Jan-Willem van de Meent, and Lawson LS Wong. Factored world models for zero-shot generalization in robotic manipulation. arXiv preprint arXiv:2202.05333 , 2022.

Akihiro Nakano, Masahiro Suzuki, and Yutaka Matsuo. Interaction-based disentanglement of entities for object-centric world models. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id=JQc2VowqCzz .

Brian Sallans and Geoffrey E Hinton. Reinforcement learning with factored states and actions. The Journal of Machine Learning Research , 5:1063-1088, 2004.

Michael Kearns and Daphne Koller. Efficient reinforcement learning in factored mdps. In Proceedings of the 16th international joint conference on Artificial Intelligence (IJCAI) , pages 740-747, 1999.

Craig Boutilier, Richard Dearden, and Moisés Goldszmidt. Stochastic dynamic programming with factored representations. Artificial intelligence , 121(1-2):49-107, 2000.

Carlos Guestrin, Daphne Koller, Ronald Parr, and Shobha Venkataraman. Efficient solution algorithms for factored mdps. Journal of Artificial Intelligence Research , 19:399-468, 2003.

Assaf Hallak, François Schnitzler, Timothy Mann, and Shie Mannor. Off-policy model-based learning under unknown factored dynamics. In International Conference on Machine Learning (ICML) , pages 711-719. PMLR, 2015.

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.

Natalia Gardiol and Leslie Kaelbling. Envelope-based planning in relational mdps. Advances in Neural Information Processing Systems (NIPS) , 16, 2003.

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.

Martijn Van Otterlo. A survey of reinforcement learning in relational domains. Centre for Telematics and Information Technology (CTIT) University of Twente, Tech. Rep , 2005.

Carlos Diuk, Andre Cohen, and Michael L Littman. An object-oriented representation for efficient reinforcement learning. In International Conference on Machine Learning (ICML) , pages 240-247, 2008.

Jonathan Scholz, Martin Levihn, Charles Isbell, and David Wingate. A physics-based model prior for object-oriented mdps. In International Conference on Machine Learning (ICML) , pages 1089-1097. PMLR, 2014.

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/ICRA. 2019.8793888.

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 (NeurIPS) , 33:11525-11538, 2020.

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.

Biwei Huang, Fan Feng, Chaochao Lu, Sara Magliacane, and Kun Zhang. AdaRL: What, where, and how to adapt in transfer reinforcement learning. In International Conference on Learning Representations (ICLR) , 2022. URL https://openreview.net/forum?id= 8H5bpVwvt5 .

Fan Feng, Biwei Huang, Kun Zhang, and Sara Magliacane. Factored adaptation for nonstationary reinforcement learning. In Advances in Neural Information Processing Systems (NeurIPS) , 2022.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems (NIPS) , 30, 2017.

Ondrej Biza, Robert Platt, Jan-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://openreview.net/ forum?id=HImz8BuUclc .

Thomas Kipf, Ethan Fetaya, Kuan-Chieh Wang, Max Welling, and Richard Zemel. Neural relational inference for interacting systems. In International Conference on Machine Learning (ICML) , pages 2688-2697. PMLR, 2018.

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.

Kevin Murphy. Dynamic bayesian networks: Representation, inference and learning. UC Berkeley, Computer Science Division , 2002.

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.

Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. In International Conference on Learning Representations (ICLR) , 2017. URL https:// openreview.net/forum?id=rkE3y85ee .

Eduardo F Camacho and Carlos Bordons Alba. Model predictive control . Springer science & business media, 2013.

John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347 , 2017.

Gautam Singh, Fei Deng, and Sungjin Ahn. Illiterate DALL-e learns to compose. In International Conference on Learning Representations (ICML) , 2022. URL https://openreview. net/forum?id=h0OYV0We3oh .

Thomas Kipf, Elise van der Pol, and Max Welling. Contrastive learning of structured world models. In International Conference on Learning Representations (ICLR) , 2020. URL https: //openreview.net/forum?id=H1gax6VtDB .

Caleb Chuck, Kevin Black, Aditya Arjun, Yuke Zhu, and Scott Niekum. Granger-causal hierarchical skill discovery. arXiv preprint arXiv:2306.09509 , 2023.

Jiaheng Hu, Zizhao Wang, Peter Stone, and Roberto Martin-Martin. Elden: Exploration via local dependencies. arXiv preprint arXiv:2310.08702 , 2023.

Rim Assouel, Pau Rodriguez, Perouz Taslakian, David Vazquez, and Yoshua Bengio. Objectcentric compositional imagination for visual abstract reasoning. In ICLR2022 Workshop on the Elements of Reasoning: Objects, Structure and Causality , 2022.

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.

Somjit Nath, Gopeshh Subbaraj, Khimya Khetarpal, and Samira Ebrahimi Kahou. Discovering object-centric generalized value functions from pixels. In International Conference on Machine Learning (ICML) , pages 25755-25768. PMLR, 2023.

Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym. arXiv preprint arXiv:1606.01540 , 2016.

John Banister Lanier. Curiosity-driven multi-criteria hindsight experience replay . University of California, Irvine, 2019.

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 preprint arXiv:1802.09464 , 2018.

Chuanxing Geng, Sheng-jun Huang, and Songcan Chen. Recent advances in open set recognition: A survey. IEEE transactions on pattern analysis and machine intelligence , 43(10): 3614-3631, 2020.

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.

Anurag Ajay, Seungwook Han, Yilun Du, Shaung Li, Abhi Gupta, Tommi Jaakkola, Josh Tenenbaum, Leslie Kaelbling, Akash Srivastava, and Pulkit Agrawal. Compositional foundation models for hierarchical planning. arXiv preprint arXiv:2309.08587 , 2023.

Mengjiao Yang, Yilun Du, Kamyar Ghasemipour, Jonathan Tompson, Dale Schuurmans, and Pieter Abbeel. Learning interactive real-world simulators. arXiv preprint arXiv:2310.06114 , 2023.

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.

Jiayuan Mao, Tomás Lozano-Pérez, Joshua B Tenenbaum, and Leslie Pack Kaelbling. Learning reusable manipulation strategies. In 7th Annual Conference on Robot Learning (CoRL) , 2023.

Abhishek Padalkar, Acorn Pooley, Ajinkya Jain, Alex Bewley, Alex Herzog, Alex Irpan, Alexander Khazatsky, Anant Rai, Anikait Singh, Anthony Brohan, et al. Open x-embodiment: Robotic learning datasets and rt-x models. arXiv preprint arXiv:2310.08864 , 2023.

Chengshu Li, Ruohan Zhang, Josiah Wong, Cem Gokmen, Sanjana Srivastava, Roberto MartínMartín, Chen Wang, Gabrael Levine, Michael Lingelbach, Jiankai Sun, et al. Behavior-1k: A benchmark for embodied ai with 1,000 everyday activities and realistic simulation. In Conference on Robot Learning (CoRL) , pages 80-93. PMLR, 2023.

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.

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 RL methods for compositional generalization.

A.1 Factored, Relational and Object-oriented MDPs

Our modelling assumptions, formalized as a Dynamic Attribute FacTored Markov Decision Process (DAFT-MDP) in Definition 1, are related to the literature on factored and object-oriented (PO)MDPs. In particular DAFT-MDPs are an 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 can 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 a family of MDPs that vary in terms of objects and types. Most of these methods focus on discrete states and often define relations or states as first-order predicates, with the exception of [28], which propose a 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 a more fine-grained factorization of the transition and reward at the object attribute level, based on estimating class template graphs and interaction pattern graphs. Additionally, we consider object-specific latent parameters that modulate the dynamics of each object and use dynamic graphs to account for the sparsity of interactions between objects and between the agent and an object.

A.2 Compositional Generalization with Object-centric RL Approaches

Table A1 provides a comparison of the object-centric RL methods discussed in Section 4. 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.

In the context of factored dynamics, most of these related works take into consideration the objectfactored states. Approaches such as STEDIE [17], NCS [12], and STOVE [8] go a step further to disentangle the states related to interaction, dynamics, or actions from other state variables, such as color, as mentioned in the Table. However, our model aims to offer a more comprehensive approach to 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

mass, and factorizing action-irrelevant factors like color and shape. In terms of interaction modeling, many approaches assume fully-connected interaction patterns, where interactions occur between every pair of objects within a given scene. However, in many real-world scenarios, this is often inaccurate as the object interaction usually happens in 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-irrelevant 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-grained 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 a 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

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 = { C 1 , C 2 } , which represent the classes of boxes and switches, respectively. We represent these two classes with cubes and cylinders in Figure A1.

For the class of boxes C 1 , represented in Figure A1A as a pink cube, we consider the attributes { C 1 .s 1 , C 1 .s 2 , C 1 .s 3 } to represent the color, velocity and position of an abstract box in two consecutive timesteps t and t +1 . All of the relationship between the attributes are described in the class template graph G C 1 . The edge C 1 .s t 1 → C 1 .s t +1 1 represents the fact that the color at timestep t +1 is only influenced by the color at timestep t (in this case being constant). The edge C 1 .s t 2 → C 1 .s t +1 2 represents the fact that the velocity at timestep t +1 is influenced by the velocity at the previous timestep t . The edge C 1 .s t 2 → C 1 .s t +1 3 means that velocity can influence position in the next timestep,

C. Interaction Graph

A. Class Template Graphs

B. Interaction Pattern Graph

o t

1

o t

3

o t

2

C

1

.

s

t

2

C

1

.

s

t

+1

1

C

1

.

s

t

1

a

t

r

t

C

1

.

s

t

3

C

1

.

s

t

+1

2

C

1

.

s

t

+1

3

C

1

.

θ

C

2

.

s

t

+1

1

C

2

.

s

t

1

C

2

.

s

t

3

C

2

.

s

t

+1

2

C

2

.

s

t

+1

3

C

2

.

θ

C

2

.

s

t

2

a

t

r

t

o t

+1

1

o t

+1

2

o t

+2

3

o t

+2

2

α t

+1

a

t

+1

α t

a

t

o t

+2

1

C

1

.

s

t

1

C

1

.

s

t

+1

1

C

1

.

s

t

2

C

1

.

s

t

3

C

1

.

s

t

+1

2

C

1

.

s

t

+1

3

C

1

.

s

t

1

C

1

.

s

t

2

C

1

.

s

t

3

C

1

.

s

t

+1

1

C

1

.

s

t

+1

2

C

1

.

s

t

+1

3

o t

+1

3

C

1

.

θ

C

1

.

θ

o

1

.

s

t

1

o

1

.

s

t

+2

1

o

3

.

s

t

2

o

3

.

s

t

1

o

3

.

s

t

3

o

1

.

s

t

2

o

2

.

s

t

1

o

1

.

s

t

+1

1

o

1

.

s

t

3

o

2

.

s

t

3

o

2

.

s

t

2

o

2

.

θ

1

o

1

.

s

t

+1

2

o

1

.

s

t

+2

3

o

3

.

s

t

+1

2

o

1

.

s

t

+2

2

a

t

o

3

.

s

t

+1

3

o

2

.

s

t

+1

3

o

2

.

s

t

+1

2

o

2

.

s

t

+1

1

o

1

.

s

t

+1

3

o

2

.

s

t

+2

1

o

2

.

s

t

+2

2

o

2

.

s

t

+2

3

o

1

.

s

t

+2

1

o

1

.

s

t

+2

2

o

1

.

s

t

+2

3

r

t

a

t

+1

r

t

+1

o

3

.

θ

1

o

1

o

3

o

2

o

1

.

θ

1

o

3

.

s

t

+1

1

which is also influenced by the previous position C 1 .s t +1 3 . In this case the agent can optionally act (the dashed lines from a t ) and influence the velocity C 1 .s t +1 2 and position C 1 .s t +1 3 . Finally the latent constant parameters for the abstract box C 1 .θ , in this case representing friction, influence the velocity C 1 .s t +1 2 and position C 1 .s t +1 3 at each timestep t +1 and modulate the effect of the action on the position and velocity. The reward r t only depends on the position C 1 .s t 3 .

For the class of switches C 2 , represented in Figure A1A as a blue cylinder, we consider the attributes { C 2 .s 1 , C 2 .s 2 , C 2 .s 3 } to represent the color, angle and position of an abstract switch in two consecutive timesteps t and t +1 . All of the relationship between the attributes are described in

the class template graph G C 2 . Similarly to boxes, the color C 2 .s 1 does not influence anything and is only influenced by its previous value. In this example, we consider that none of the attributes influences each other, but that the reward r t is only a function of the angle C 2 .s t 2 . Additionally, the latent constant parameters C 2 .θ only influence the position C 2 .s t 3 , as can potentially the action a t (represented by the dashed red lines).

In Figure A1B we show the interaction pattern graph for boxes G C 1 ,C 1 that represents the way that attributes of two objects of the class box C 1 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 A1C we consider a specific environment, specified by the objects O = { o 1 , o 2 , o 3 } , where o 1 and o 3 are boxes, while o 2 is a switch. We show an unrolled version of the dynamic interaction graph G inter for three consecutive timesteps for these objects. At timestep t there is an interaction between the two boxes o 1 and o 3 (represented by blue lines), and the action a t binds to the switch o 2 , as selected by the action selector α t . The interactions between the two objects are instantiated following the interaction pattern graphs in Figure A1B. The action binding instead activates the red dashed lines in the class template graph in Figure A1A. In the next timestep t +1 there are no interactions between objects, but the action is now bound to object o 1 .

In the bottom graph in Figure A1, we 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 as described by the class template graphs G C 1 and G C 2 (in grey). The interactions between the two boxes in timestep t get instantiated following the interaction pattern graphs G C 1 ,C 1 (in light blue). The action binding specifies which object is affected by the action a t 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 a 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 is 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 G C j for each object class C j . 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.

̂

Step 1: Class learning in single-object environments

{(

x

t

1

,

a t

1

,

r

t

1

)}

T

t

=1

{(

x

t

2

,

a t

2

,

r

t

2

)}

T

t

=1

{(

x

t

3

,

a t

3

,

r

t

3

)}

T

t

=1

Object-centric model

{

o

t

1

,

a t

1

,

r

t

1

}

T

t

=1

{

o

t

2

,

a t

2

,

r

t

2

}

T

t

=1

{

o

t

3

,

a t

3

,

r

t

3

}

T

t

=1

Learning class template graph

Object class inference

/u1D4A2

C

1

, /u1D4A2

C

2

, /u1D4A2

C

3

C

1

,

C

2

,

C

3

Stage 2: Interaction Learning and Latents Inference

Stage 3: Policy Learning

Stage 4: Adaptation

Apply policy

Inferred constant latent parameters

o

i

.

θ

Inferred object inertaction graph

/u1D4A2

inter

π

*(

a

t

|

s

t

,

θ

, /u1D4A2

inter

)

Learn the object interaction graph

Learn the interaction pattern graphs

Learn the latent constant parameters

Learn the object selector

/u1D4A2

inter

o

1

.

θ

,

o

2

.

θ

,

o

3

.

θ

α t

/u1D4A2

Ci

→

Cj

,

i

∈ {1,2,3},

j

∈ {1,2,3}

{

o

t

1

,

o

t

2

,

o

t

3

,

a t

,

r

t

)}

T

t

=1

{(

x

t

,

a t

,

r

t

)}

T

t

=1

Object-centric model

{(

x

t

,

a t

,

r

t

,)}

T

t

=1

{

o

t

1

,

o

t

2

,

o

t

3

,

o

t

4

,

a t

,

r

t

)}

T

t

=1

Object-centric model

{(

x

t

,

a t

,

r

t

)}

T

t

=1

{

o

t

1

,

o

t

2

,

o

t

3

,

o

t

4

,

o

t

5

,

a t

,

r

t

)}

T

t

=1

Object-centric model

Generate imagined trajectories

{(

o

t

1

,

o

t

2

,

o

t

3

,

o

4

,

a t

,

r

t

)}

T

2

t

=

T

1

Learned policy

π

*(

a

t

|

s

t

,

θ

, /u1D4A2

inter

)

{(

x

t

,

a t

,

r

t

)}

T

t

=1

{(

x

t

,

a t

,

r

t

)}

T

t

=1

̂

̂

̂

̂

̂

For l = { 1 , . . . , n } we define:

C j .G s → s l ∈ { 0 , 1 } n as the subset of the adjacency matrix of G C j that describes the edges incoming into C j .s t +1 l from each C j .s t i for i ∈ { 1 , . . . , n } . In other words this binary vector is 1 only for the components C j .s t i for i ∈ { 1 , . . . , n } that have an edge to C j .s t +1 l in G C j .

C j .G a → s l ∈ { 0 , 1 } as a binary value that represents the existence of an edge between a t and C j .s t +1 l , denoting how the action will affect the dynamics at the attribute s l at next time step.

C j .G θ → s l ∈ { 0 , 1 } as a binary value that represents the existence of an edge between C j .θ t and C j .s t +1 l

We additionally define:

C j .G s → r ∈ { 0 , 1 } n is the subset of the adjacency matrix of G C j that describes the edges incoming into r t +1 from each C j .s t i for i ∈ { 1 , . . . , n }

C j .G a → r ∈ { 0 , 1 } is a binary value that represents the existence of an edge between a t .

For each object class C j all of these five groups of binary matrices together represent the adjacency matrix of graph G C j .

Binary matrices in interaction pattern graphs. Interaction pattern graphs G C i ,C j represent the interaction structures between attributes of the object with class C i and the object with class C j , which are assumed to be constant in time. For each l = { 1 , . . . , n } we define G C i → C j .s l ∈ { 0 , 1 } n to represent the binary vector that describes how the attributes of an object with class C i influence the attribute C j .s t +1 l of the object with class C j . All of these vectors together represent the adjacency matrix of graph G C i ,C j .

Binary values in Dynamic interaction graphs The dynamic interaction graphs G t inter represent the interactions between the objects O at timestep t , which include the action selector { α 1 , α 2 , . . . , α T } , where α i ∈ { 0 , 1 } m represents which object is affected by a t at all time steps t = 1 , . . . , T . We use the action selection α i in our losses. For each o i , o j ∈ O we define a binary value G t inter( i,j ) ∈ { 0 , 1 }

to represent whether object o i and o j will have the interaction at time step t . All of these values together represent the adjacency matrix of graph G t 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 in the previous subsections are learn at that step.

Step 1: Class learning in single-object environments As described in Sec 3.1, we learn the class template graphs {G C 1 , G C 2 , . . . , G C k } using the collected D single . We give the detailed loss functions, including the prediction loss L 1 pred and the sparsity regularization L 1 sparse for dynamics and reward models, for this step below:

where we update the prediction models including λ s , λ r , as well as the binary vectors defined previously {{ C j .G s → s l } n l =1 , { C j .G a → s l } n l =1 , C j .G s → r , C j .G a → r } (marked in blue). The complete loss L 1 for Step 1 is:

where γ 1 sparse is a hyper-parameter.

Step 2.1: Learning the action binding In this step, we learn the dynamic action selector α t by using the soft attention networks. We have collected D multi 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:

where f v is the value network to be updated. We also need to learn the key and query networks f k and f q to update α t i for each object o i (marked in blue). Additionally, we use the same dynamic network λ s and also update it in this stage.

Step 2.2: learning dynamic interaction graph As described in Sec 3.2.2, we use 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, at this step. We consider the class template graphs and action binding are fixed from the previous step. Detailed loss functions are:

where λ s is the dynamics model we have employed in the previous steps. We learn the encoder/prior networks ϕ to generate the latent vectors z , where we sample the graph G inter from. Through the binary vectors we also learn the interaction pattern graphs { G C i ,C j } i,j ∈{ 1 ,...,k } , latent parameters { o i .θ } i m =1 and the binary vectors representing the edges from latent parameters to each attribute { C i .G θ → s l } k i =1 . We have completed the loss for step 2.2:

where γ 3 KL and γ 3 sparse are the hyper-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 a supervised manner, where the input includes the extracted feature vectors o i for each object and we have the labels y i for each single objects. We apply a cross-entropy loss to learn the mapping.

Interaction graph sampling During the learning of dynamic interaction graphs, we generate the edge distribution z t at each time step t . We sample the edge matrix M t of the graph G t inter . Specifically, M t ∼ Bern( z t ) , where Bern is the multivariate Bernoulli distribution and all elements are mutually independent. We also employ the Gumbel-Softmax trick [40] to 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 in a fair comparison with DAFT-RL, focusing in particular on adapting symbolic input approaches to pixel inputs and image-based approaches to symbolic inputs. We 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 on symbolic inputs, such as self-attention [15], deep sets [15], GNN [4], SRICS [6] and LRN [7], to handle scenarios where the inputs are raw pixels. To accomplish this, we leverage pre-trained objectcentric models, specifically AIR [31], to obtain the state representation of each object within the scene. The pre-trained object-centric models, specifically AIR, are employed to extract objectspecific information from the raw pixel inputs. These object-factored states, which represent the individual objects, are then used as 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 to be applicable in scenarios where raw pixels serve as the input representation.

Image-based approaches with symbolic inputs. To modify the image-based approaches to 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, we remove the SuPAIR model and directly observe object-factored states for dynamics and policy modeling.

NCS [12]: we 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.

D.2.1 Task Description

OpenAI 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 is the 3D continuous control of the end effector displacement and 1D 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.

Spriteworld. 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.

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 : the goal is to push the target to a given location in the scene. The reward will be obtained if the agent pushes the target to the goal position and the action also includes four cardinal directions.

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 [61]. 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 { 1 , 2 , 3 , 11 , 13 } ∪ { 3 , 5 , 7 , 9 } and { 0 . 1 , 0 . 2 , 1 . 1 , 1 . 3 } ∪ { 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 , 8 } . By considering these changes in colors, shapes, and object numbers, we 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-processing filters applied to the image frames. In particular, picking up a heavier object will require more time compared to selecting a lighter one due to gravity. Placing a 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 } .

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 MLP 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 f k , f q , and f v , 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 1 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 1, we gather a 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 is 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 is 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 1, 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 × 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 a 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 hyper-parameters 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 use the same learning rate to 2 . 5 × 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 3 MLP layers of size 256 and 128 in both experiments. For Stacking, we follow all the MPC hyper-parameter settings in [17] for the planning.

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.

±

±

±

±

±

±

±

±

±

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;

I

II

III

IV

V

VI

VII VIII

Cases

0

.

5

0

.

6

0

.

7

0

.

8

Success Rate

2 - Push + 2 - Switch(L + S)

I

II

III

IV

V

VI

VII VIII

Cases

0

.

5

0

.

6

0

.

7

0

.

8

0

.

9

Success Rate

2 - Switch(L)

I

II

III

IV

V

VI

VII VIII

Cases

0

.

6

0

.

8

Success Rate

2 - Push(L)

II

III

IV

V

VI

VII

VIII

Cases

0

.

4

0

.

6

0

.

8

Success Rate

2 - Push + 2 - Switch(S)

I

II

III

IV

V

VI

VII VIII

Cases

0

.

60

0

.

65

0

.

70

0

.

75

Success Rate

3 - Push + 3 - Switch(L + O + S)

II

III

IV

V

VI

VII

VIII

Cases

0

.

4

0

.

5

0

.

6

0

.

7

0

.

8

Success Rate

3 - Push + 3 - Switch(S + O)

I

II

III

IV

V

VI

VII VIII

Cases

0

.

6

0

.

8

Success Rate

3 - Push(L + O)

I

II

III

IV

V

VI

VII VIII

Cases

0

.

4

0

.

6

0

.

8

Success Rate

3 - Switch(L + O)

I

II

III

IV

VI

VII

VIII

Cases

0

.

725

0

.

750

0

.

775

0

.

800

0

.

825

Success Rate

2 Blocks

I

II

III

IV

VI

VII

VIII

Cases

0

.

60

0

.

65

0

.

70

Success Rate

4 Blocks

I

II

III

IV

VI

VII

VIII

Cases

0

.

55

0

.

60

0

.

65

Success Rate

6 Blocks

I

II

III

IV

VI

VII

VIII

Cases

0

.

45

0

.

50

0

.

55

Success Rate

8 Blocks

Based on the obtained results, we observe that each learned component plays a crucial role in optimizing the policy. Notably, the dynamic interaction graph consistently demonstrates a significant contribution across the majority of experiments in Push & Switch and stacking benchmarks.

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 on policy optimization.

D.4 Results on End-to-end Learning

Table A8 presents the results of end-to-end training. For a fair comparison, we use the same number of data samples as in the original multi-step training. In this end-to-end approach, only D multi is employed as both the interaction model and object template graphs are 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.

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.

±

±

±

±

±

±

±

±

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-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 of 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 A11. Notably, while models

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 R 2 coefficient of our learned parameters with the true latent parameters, both for a random policy (Fig. A4a) and for a 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 a random or 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 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.

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

Ratio of Full Samples

-

0

.

75

-

0

.

50

-

0

.

25

0

.

00

0

.

25

0

.

50

0

.

75

Measurements

Success Rate

R

2

nSHD

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

Ratio of Full Samples

-

0

.

2

0

.

0

0

.

2

0

.

4

0

.

6

Measurements

Success Rate

R

2

nSHD

(a)

(b)

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 A1 and A2. The variables

C 1 .s 1 , C 1 .s 2 , C 1 .s 3 , and C 1 .s 4 represent the color, shape, velocity, and position of the boxes, respectively. Similarly, C 2 .s 1 , C 2 .s 2 , C 2 .s 3 , and C 2 .s 4 correspond to the color, shape, position, and angle of the switches. In both cases, C i .θ and C j .θ indicate the friction coefficients.

As expected, in the learned class template graph, we observe that shape and color do not have a 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 at specific positions on the panel. Also, as expected, actions applied to 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.

c

1

s

t

+1

1

c

1

s

t

+1

2

c

1

s

t

+1

3

c

1

s

t

+1

4

c

1

s

t

1

c

1

s

t

2

c

1

s

t

3

c

1

s

t

4

a t

r

t

c

,

θ

c

2

s

t

+1

1

c

2

s

t

+1

2

c

2

s

t

+1

3

c

2

s

t

+1

4

c

2

s

t

1

c

2

s

t

2

c

2

s

t

3

c

2

s

t

4

a t

r

t

c

,

θ

c

1

s

t

+1

4

c

1

s

t

1

c

1

s

t

2

c

1

s

t

3

c

1

s

t

4

c

1

s

t

1

c

1

s

t

2

c

1

s

t

3

c

1

s

t

4

c

1

s

t

+1

1

c

1

s

t

+1

2

c

1

s

t

+1

3

c

1

s

t

+1

4

c

1

s

t

+1

1

c

1

s

t

+1

2

c

1

s

t

+1

3

c

1,

θ

A1

A2

B

c

1,

θ

As for the interaction graph between two boxes, shown in Figure A5B, the velocity of one box can impact the position of the other box, while the other variables do not have a 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\n- 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).\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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?\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nThere is no discussion about the limitations of the work.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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?\n\nSTRENGTHS:\n- 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).\n\nWEAKNESSES:\n- \"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.\n\nQUESTIONS:\nNone really. I didn't find any parts of the paper confusing.\n\nLIMITATIONS:\nIs 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).", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-08"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n* 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n3. 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.\n\nLIMITATIONS:\nI 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", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-09"}

comments: 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

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.
The definition looks cluttered line 163.
A pictorial example at the intro of a sample task is expected. (borrow from supplementary add here)

QUESTIONS:
Why the same paper content is repeated in supplementary upto pg 9? remove.
I was wondering why Dynamic Bayesian Networks (DBNs) was chosen?
How 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?
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-10

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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\n\nWEAKNESSES:\n- 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?\n\nQUESTIONS:\n- 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\nLIMITATIONS:\nThe 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?", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-21"}