SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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.
[49] 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.
[50] Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym. arXiv preprint arXiv:1606.01540, 2016.
[51] John Banister Lanier. Curiosity-driven multi-criteria hindsight experience replay. University of California, Irvine, 2019.
[52] Matthias Plappert, Marcin Andrychowicz, Alex Ray, Bob McGrew, Bowen Baker, Glenn Powell, Jonas Schneider, Josh Tobin, Maciek Chogiej, Peter Welinder, et al. Multi-goal reinforcement learning: Challenging robotics environments and request for research. arXiv preprint arXiv:1802.09464, 2018.
[53] 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.
[54] Phillip Lippe, Sara Magliacane, Sindy Löwe, Yuki M Asano, Taco Cohen, and Efstratios Gavves. BISCUIT: Causal representation learning from binary interactions. In Proceedings of the Thirty-Ninth Conference on Uncertainty in Artificial Intelligence (UAI), volume 216 of Proceedings of Machine Learning Research, pages 1263-1273. PMLR, 2023.
[55] Anurag Ajay, Seungwook Han, Yilun Du, Shaung Li, Abhi Gupta, Tommi Jaakkola, Josh Tenenbaum, Leslie Kaelbling, Akash Srivastava, and Pulkit Agrawal. Compositional foundation models for hierarchical planning. arXiv preprint arXiv:2309.08587, 2023.
[56] Mengjiao Yang, Yilun Du, Kamyar Ghasemipour, Jonathan Thompson, Dale Schuurmans, and Pieter Abbeel. Learning interactive real-world simulators. arXiv preprint arXiv:2310.06114, 2023.
[57] Yifeng Zhu, Zhenyu Jiang, Peter Stone, and Yuke Zhu. Learning generalizable manipulation policies with object-centric 3d representations. In 7th Annual Conference on Robot Learning (CoRL), 2023.
[58] Jiayuan Mao, Tomás Lozano-Pérez, Joshua B Tenenbaum, and Leslie Pack Kaelbling. Learning reusable manipulation strategies. In 7th Annual Conference on Robot Learning (CoRL), 2023.
[59] Abhishek Padalkar, Acorn Pooley, Ajinkya Jain, Alex Bewley, Alex Herzog, Alex Irpan, Alexander Khazatsky, Anant Rai, Anikait Singh, Anthony Brohan, et al. Open x-embodiment: Robotic learning datasets and rt-x models. arXiv preprint arXiv:2310.08864, 2023.
[60] Chengshu Li, Ruohan Zhang, Josiah Wong, Cem Gokmen, Sanjana Srivastava, Roberto Martín-Martín, Chen Wang, Gabrael Levine, Michael Lingelbach, Jiankai Sun, et al. Behavior-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.
[61] Rishi Veerapaneni, John D Co-Reyes, Michael Chang, Michael Janner, Chelsea Finn, Jiajun Wu, Joshua Tenenbaum, and Sergey Levine. Entity abstraction in visual model-based reinforcement learning. In Conference on Robot Learning (CoRL), pages 1439-1456. PMLR, 2020.
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  $\mathcal{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.
Methods	Factored Dynamics	Interaction Modeling	Generalize to Changing Latent Factors
GNN [4]	Object	Fully-connected	X
Self-attention [15]	Object	Fully-connected	X
Deep Sets [15]	Object	Fully-connected	X
LRN [7]	Object	Fully-connected	X
COBRA [9]	Object	X	X
FWM [16]	Object	Fully-connected	X
SMORL [5]	Object	X	X
STOVE [8]	Object & pre-determined attributesa	Fully-connected	X
SRICS [6]	Object	Sparse & dynamic (object-level)	X
NCS [12]	Object & action/dynamics-relevant	Fully-connected	X
STEDIE [17]	Object & interaction-relevantb	Fully-connected (object-level), factored (interaction-relevant)	X
DAFT-RL	Object & attribute	Sparse, dynamic & factored (object-level & attribute-level)	✓
${}^{a}$  STOVE focuses specifically on these attributes: (position,velocity,size,others).
${}^{b}$  STEDIE learns which attributes interact during object interactions,similar to our interaction pattern graphs.
Table A1: Comparison of different object-oriented RL methods for compositional generalization. A check denotes that a method has an attribute, whereas a cross denotes the opposite.
In the context of factored dynamics, most of these related works take into consideration the object-factored states. Approaches such as STEDIE [17], NCS [12], and STOVE [8] go 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
Notation	Definition
C={C1,...,Ck}	Set of object classes
{Cj.s1,Cj.s2,...,Cj.sn}	Class attributes for class Cj
Cj.θ	Latent constant parameters for class Cj
O={o1,...,om}	Set of objects in a domain
C(i)	Class of object o,i, in other words Cj such that o,i∈Cj
{o.i.s1,o.i.s2,...,o.i.sn}	Object attributes for object o,i with semantics defined by the class C(i)
o.i.θ	Latent constant parameters of object i
xt	Pixel observation at time step t
xi	Pixel observation of object-i at time step t
oist= {o.i.s1, o.i.s2, ..., o.i.snt}	Symbolic state of object o,i at time step t
st= {o1, o1, ..., omt}	Observable state at time step t for environment with objects
θ={o1.θ, ..., omt}	Latent constant parameters for environment with O objects
aith	Action on object o,i at time step t
at= {a1, ..., am}	Action at time step t for environment with O objects
rti	Reward on object o,i at time step t
rt= Σi rti	Reward at time step t
Gcj	Class template graph for object with class Cj
GCr,i,Cj	Interaction pattern graph for objects with class Ci and Cj
αt	Action selector at time step t
Goi	Instantiation of the class template graph for o,i with class Cj at time t
Ginter	Dynamic object interaction graph at time step t
Table A2: Summary of notations in this paper.
In this section, we first provide a summary of the notation in Table A2. Then we show an example of the environment described in the main paper, and how the learned graphs are connected in a single ground graphical model, as described in Figure A1.
In our example, we consider two classes  $\mathcal{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  $\mathcal{G}_{C_1}$ . The edge  $C_1.s_1^t\to C_1.s_1^{t + 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_2^t\rightarrow C_1.s_2^{t + 1}$  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_2^t\rightarrow C_1.s_3^{t + 1}$  means that velocity can influence position in the next timestep,
Figure A1: The graphical representation of DAFT-MDP. Fig A1.A represents the class template graphs for boxes and switches. Fig A1.B represents the interaction pattern graphs between two boxes. Fig A1.C represents the dynamic interaction graph, that at each timestep predicts which objects interact with each other and with the agent. The graph on the bottom shows an example of the instantiation of the complete ground graphical model for all of the objects in the environment from Fig A1.C, i.e. a pink box, a blue switch and a green box at time steps  $t$ ,  $t + 1$  and  $t + 2$ . The red lines describe the interactions of the agent with the objects (which follow the dotted lines in the class template graphs for each object class in Fig A1.A). The blue lines represent the interactions between objects, which follow the interaction patterns described in Fig A1.B.
which is also influenced by the previous position  $C_1.s_3^{t + 1}$ . In this case the agent can optionally act (the dashed lines from  $\mathbf{a}^t$ ) and influence the velocity  $C_1.s_2^{t + 1}$  and position  $C_1.s_3^{t + 1}$ . Finally the latent constant parameters for the abstract box  $C_1.\theta$ , in this case representing friction, influence the velocity  $C_1.s_2^{t + 1}$  and position  $C_1.s_3^{t + 1}$  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_3^t$ .
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  $\mathcal{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_2^t$ . Additionally, the latent constant parameters  $C_2.\theta$  only influence the position  $C_2.s_3^t$ , as can potentially the action  $\mathbf{a}^t$  (represented by the dashed red lines).
In Figure A1B we show the interaction pattern graph for boxes  $\mathcal{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  $\mathcal{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  $\mathcal{G}_{\mathrm{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  $\mathbf{a}^t$  binds to the switch  $o_2$ , as selected by the action selector  $\alpha^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  $\mathcal{G}_{C_1}$  and  $\mathcal{G}_{C_2}$  (in grey). The interactions between the two boxes in timestep  $t$  get instantiated following the interaction pattern graphs  $\mathcal{G}_{C_1,C_1}$  (in light blue). The action binding specifies which object is affected by the action  $\mathbf{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  $\mathcal{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.
Figure A2: The learning pipelines of the DAFT-RL framework.
For  $l = \{1, \dots, n\}$  we define:
-  $C_j.G_{s\to s_l} \in \{0,1\}^n$  as the subset of the adjacency matrix of  $\mathcal{G}_{C_j}$  that describes the edges incoming into  $C_j.s_l^{t + 1}$  from each  $C_j.s_i^t$  for  $i \in \{1,\dots ,n\}$ . In other words this binary vector is 1 only for the components  $C_j.s_i^t$  for  $i \in \{1,\dots ,n\}$  that have an edge to  $C_j.s_l^{t + 1}$  in  $\mathcal{G}_{C_j}$ .
-  $C_j.G_{a\to s_l} \in \{0,1\}$  as a binary value that represents the existence of an edge between  $\mathbf{a}^t$  and  $C_j.s_l^{t + 1}$ , denoting how the action will affect the dynamics at the attribute  $s_l$  at next time step.
-  $C_j.G_{\theta \to s_l} \in \{0,1\}$  as a binary value that represents the existence of an edge between  $C_j.\theta^t$  and  $C_j.s_l^{t + 1}$
We additionally define:
-  $C_j.G_{s\rightarrow r} \in \{0,1\}^n$  is the subset of the adjacency matrix of  $\mathcal{G}_{C_j}$  that describes the edges incoming into  $r^{t + 1}$  from each  $C_j.s_i^t$  for  $i \in \{1,\dots ,n\}$
-  $C_j.G_{a\to r}\in \{0,1\}$  is a binary value that represents the existence of an edge between  $\mathbf{a}^t$
For each object class  $C_j$  all of these five groups of binary matrices together represent the adjacency matrix of graph  $\mathcal{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,\dots ,n\}$  we define  $G_{C_i\to C_j.s^l}\in \{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_l^{t + 1}$  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_{\mathrm{inter}}^{t}$  represent the interactions between the objects  $\mathcal{O}$  at timestep  $t$ , which include the action selector  $\{\alpha^1, \alpha^2, \dots, \alpha^T\}$ , where  $\alpha^i \in \{0, 1\}^m$  represents which object is affected by  $\mathbf{a}^t$  at all time steps  $t = 1, \dots, T$ . We use the action selection  $\alpha^i$  in our losses. For each  $o_i$ ,  $o_j \in \mathcal{O}$  we define a binary value  $G_{\mathrm{inter}(i,j)}^{t} \in \{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_{\mathrm{inter}}^t$ .
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  $\{\mathcal{G}^{C_1},\mathcal{G}^{C_2},\ldots ,\mathcal{G}^{C_k}\}$  using the collected  $\mathcal{D}^{\mathrm{single}}$ . We give the detailed loss functions, including the prediction loss  $\mathcal{L}_{\mathrm{pred}}^{1}$  and the sparsity regularization  $\mathcal{L}_{\mathrm{sparse}}^{1}$  for dynamics and reward models, for this step below:
$$
\begin{array}{l} \mathcal {L} _ {\text {p r e d}} ^ {1} = \sum_ {t = 1} ^ {T} \sum_ {i = 1} ^ {m} \sum_ {l = 1} ^ {n} \left(\log p _ {\lambda_ {s}} \left(o _ {i}. s _ {l} ^ {t + 1} \mid C _ {j}. G _ {s \rightarrow s _ {l}} \odot \mathbf {o} _ {\mathbf {i}} ^ {t}, C _ {j}. G _ {a \rightarrow s _ {l}} \odot \mathbf {a} _ {i} ^ {t}\right)\right. \tag {A1} \\ + \log p _ {\lambda_ {r}} \left(r _ {i} ^ {t} \mid C _ {j}. G _ {s \rightarrow r} \odot \mathbf {o _ {i}} ^ {t}, C _ {j}. G _ {a \rightarrow r} \odot \mathbf {a} _ {i} ^ {t}\right)) \\ \end{array}
$$
$$
\mathcal {L} _ {\text {s p a r s e}} ^ {1} = \sum_ {j = 1} ^ {k} \left(\sum_ {l = 1} ^ {n} \| C _ {j}. G _ {s \rightarrow s _ {l}} \| _ {1} + \| C _ {j}. G _ {s \rightarrow r} \| _ {1} + \sum_ {l = 1} ^ {n} \| C _ {j}. G _ {a \rightarrow s _ {l}} \| _ {1} + \| C _ {j}. G _ {a \rightarrow r} \| _ {1}\right) \tag {A2}
$$
where we update the prediction models including  $\lambda_s, \lambda_r$ , as well as the binary vectors defined previously  $\{\{C_j.G_{s\rightarrow s_l}\}_{l = 1}^n,\{C_j.G_{a\rightarrow s_l}\}_{l = 1}^n,C_j.G_{s\rightarrow r},C_j.G_{a\rightarrow r}\}$  (marked in blue). The complete loss  $\mathcal{L}^1$  for Step 1 is:
$$
\mathcal {L} ^ {1} = \mathcal {L} _ {\text {p r e d}} ^ {1} + \gamma_ {\text {s p a r s e}} ^ {1} \mathcal {L} _ {\text {s p a r s e}} ^ {1} \tag {A3}
$$
where  $\gamma_{\mathrm{sparse}}^{1}$  is a hyper-parameter.
Step 2.1: Learning the action binding In this step, we learn the dynamic action selector  $\alpha^t$  by using the soft attention networks. We have collected  $\mathcal{D}^{\mathrm{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:
$$
\mathcal {L} _ {\text {p r e d}} ^ {2} = \sum_ {t = 1} ^ {T} \sum_ {i = 1} ^ {m} \sum_ {l = 1} ^ {n} \log p _ {\lambda_ {s}} \left(o _ {i}. s _ {l} ^ {t + 1} \mid C _ {i}. G _ {s \rightarrow s _ {l}} \odot \mathbf {o} _ {\mathbf {i}} ^ {t}, C _ {i}. G _ {a \rightarrow s _ {l}} \odot \alpha_ {i} ^ {t} \cdot f _ {v} (\mathbf {a} ^ {t})\right) \tag {A4}
$$
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  $\alpha_{i}^{t}$  for each object  $o_{i}$  (marked in blue). Additionally, we use the same dynamic network  $\lambda_{s}$  and also update it in this stage.