SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Fan Feng
ffeng1017@gmail.com
Sara Magliacane
sara.magliacane@gmail.com
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.
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 real-world 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 Object-Oriented (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  $\mathcal{C} = \{C_1,\dots ,C_k\}$  of objects, where each class  $C_j$  describes a set of observable attributes  $\{C_{j}.s_{1},\ldots ,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.\theta$ , 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  $\mathcal{C} = \{\mathrm{BOX},\mathrm{SW}\}$ , boxes and switches, with class attributes  $\{\mathrm{BOX.color,BOX(pos,BOX vel,BOX.}\theta \}$  representing the object color, position, velocity, and friction coefficient of a box, and  $\{\mathrm{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
A. Class Template Graphs
Figure 1: The graphical representation of DAFT-MDP. The colors denote the attributes for an object or a class, the red dashed lines denote edges that can be switched on or off at different timesteps.
B. Interaction Pattern Graph
C. Interaction Graph
vary in the number and characteristics of the objects. Each environment has a fixed set of objects  $\mathcal{O} = \{o_1,\dots,o_m\}$ , where each object is an instance of a class in  $\mathcal{C}$ , which we denote as  $o_i\in C_j$  for  $i\in \{1,\ldots ,m\},j\in \{1,\ldots ,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  $\mathcal{O}' = \{o_1,o_2\}$  where  $o_1\in \mathrm{BOX}$  and  $o_2\in \mathrm{SW}$ , and a target environment as  $\mathcal{O}'' = \{o_1,o_2,o_3\}$ , where we add  $o_3\in \mathrm{BOX}$  to the original objects. For each object  $o_i\in C_j$  and timestep  $t = \{1,\dots,T\}$ , we denote its attributes at time  $t$  as  $\mathbf{o}_i^t = \{o_i.s_1^t,\dots,o_i.s_n^t\}$ , which are instantiations of the class attributes  $\{C_j.s_1,\dots,C_j.s_n\}$ , and its constant parameters as  $o_i.\theta$ . In our example, for box  $o_1$  the attributes  $o_1.s_1^t,o_1.s_2^t$  and  $o_1.\theta$  are its position and velocity at time  $t$  and its friction coefficient, while for the switch  $o_2$  the attributes  $o_2.s_1^t$  and  $o_2.s_2^t$  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.,  $\mathbf{s}^{\mathbf{t}} = \{\mathbf{o}_{\mathbf{1}}^{\mathbf{t}}, \dots, \mathbf{o}_{\mathbf{m}}^{\mathbf{t}}\}$  with domain  $S$ . We collect all object-specific latent parameters in a global parameter  $\pmb{\theta} = \{o_1, \pmb{\theta}, \dots, o_m, \pmb{\theta}\}$  with domain  $\Theta$ . We define the space of actions  $\mathcal{A}$  and use  $\mathbf{a}^t \in \mathcal{A}$  to denote the action at time  $t$ . We denote the transition probability as  $p(\mathbf{s}^{t+1} | \mathbf{s}^t, \mathbf{a}^t, \pmb{\theta})$ . The reward at time  $t$  is denoted as  $r^t$  and the reward probability is denoted as  $p(r^t | \mathbf{s}^t, \mathbf{a}^t, \pmb{\theta})$ . 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  $\mathcal{G}_{C_j}$  over the nodes  $\{C_j.s_1^t,\ldots,C_j.s_n^t,C_j.s_1^{t+1},\ldots,C_j.s_n^{t+1},C_j.\theta,\mathbf{a}^t,r^t\}$  that represents a template for the instance graph between the attributes for each object of this class  $\mathcal{G}_{o_i}^t$  for  $o_i \in 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_2^t$  influences  $C_1.s_2^{t + 1}$  and  $C_1.s_3^{t + 1}$ , but it's not influenced by  $C_1.s_3^{t - 1}$ . The reward  $r^t$  is only influenced by  $C_1.s_3^t$  and the action will only have an effect on  $C_1.s_2^{t + 1}$  and  $C_1.s_3^{t + 1}$ . Moreover,  $C_2.s_2^t$  influences only  $C_2.s_2^{t + 1}$  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  $\mathcal{G}_{C_i,C_j}$  over the nodes  $\{\{C_i.s_l^t\}_{l=1}^n, \{C_i.s_l^{t+1}\}_{l=1}^n, C_i.\theta, \{C_j.s_l^t\}_{l=1}^n, \{C_j.s_l^{t+1}\}_{l=1}^n, C_j.\theta\}$ . 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_2^t$  from one object influencing the  $C_1.s_1^{t+1}$  and  $C_1.s_2^{t+1}$  from the other object, and the latent parameters  $C_1.\theta$  from each object influencing  $C_1.s_3^{t+1}$  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  $\mathcal{G}_{\mathrm{inter}} = \{\mathcal{G}_{\mathrm{inter}}^t\}_{t=1}^T$  at the object level, which is a sequence of graphs  $\mathcal{G}_{\mathrm{inter}}^t$  with edges from a subset of  $\{o_1^t,\dots,o_m^t,\mathbf{a}^t\}$  to a subset of  $\{o_1^{t+1},\dots,o_m^{t+1}\}$ . Each edge  $o_i^t\rightarrow o_j^{t+1}$  represents an interaction between an object  $o_i$  and an object  $o_j$ . This interaction is instantiated in an instance interaction graph  $\mathcal{G}_{o_i,o_j}^t$ , following the interaction pattern graph  $\mathcal{G}_{C_i,C_j}$  for each pair of objects  $o_i\in C_i$  and  $o_j\in C_j$ .
We also learn an object selector  $\alpha^t\in \{1,\dots ,m\}$  that selects the single object on which the action  $\mathbf{a}^t$  has an effect at timestep  $t$ , which we represent with an edge from  $\mathbf{a}^t$  to the selected object  $o_{\alpha^t}^t$  in  $\mathcal{G}_{\mathrm{inter}}^t$ . In particular, this selection (or action binding) activates the edges from  $\mathbf{a}^t$  to the object attributes described in the class template graph  $\mathcal{G}_{C_j}$  in the instantiated version for object  $o_i\in C_j$ , which we denote  $\mathcal{G}_{o_i}^t$ . The graph  $\mathcal{G}_{\mathrm{inter}}$  is dynamic because the edges of each graph in the sequence can change at each timestep  $t = 1,\ldots ,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_2^t \to \{o_3.s_1^{t+1}, o_3.s_2^{t+1}\}$  and  $o_1.\theta^t \to o_3.s_3^{t+1}$ , and viceversa for  $o_3$  towards  $o_1$ . Moreover,  $\mathbf{a}^t$  has an effect on  $o_2$ , and specifically  $o_2.s_3^{t+1}$ , following the class template graph presented in Fig. 1A. Instead, at timestep  $t+1$ , there are no interactions between the objects and  $\mathbf{a}^{t+1}$  has an effect on  $o_1$ , specifically  $o_1.s_2^{t+2}$  and  $o_1.s_3^{t+2}$ . For completeness, in App. B we provide an example of this graph combined with the instantiated class template graphs for each object  $\mathcal{G}^{o_1}, \mathcal{G}^{o_2}$  and  $\mathcal{G}^{o_3}$  for three timesteps  $t, t+1, t+2$ , as well as the instantiated interaction pattern graph  $\mathcal{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  $\mathcal{O}$  and their dynamic interaction graph  $\mathcal{G}_{\mathrm{inter}}$ .
Definition 1 (DAFT-MDP). A Dynamic Attribute-FacTored Markov Decision Process (DAFT-MDP) is a tuple  $(\mathcal{C},\mathcal{O},\Theta ,\mathcal{A},\mathcal{G},\mathbb{P}_s,\mathcal{R},\gamma)$ , where  $\mathcal{C}$  is the set of classes,  $\mathcal{O}$  is the set of objects,  $\Theta$  is the space of the constant latent parameters,  $\mathcal{A}$  the action space,  $\mathbb{P}_s$  is the transition distribution,  $\mathcal{R}$  is the reward function and  $\mathcal{G}$  is a set of graphs that contains the collection of class template graphs for each class  $\{\mathcal{G}_{C_j}\}_{C_j\in \mathcal{C}}$ , the collection of interaction pattern graphs for each pair of classes  $\{\mathcal{G}_{C_i,C_j}\}_{C_i,C_j\in \mathcal{C}}$  and the dynamic interaction graph  $\mathcal{G}_{\mathrm{inter}}$ , as defined previously. These graphs define the factorization of the transition distribution per object and per attribute, as follows:
$$
\mathbb {P} _ {s} (\pmb {s} ^ {t + 1} | \pmb {s} ^ {t}, \pmb {\theta}, \pmb {a} ^ {t}) = \prod_ {i = 1} ^ {m} \prod_ {l = 1} ^ {n} \mathbb {P} _ {s} \left(o _ {i}. s _ {l} ^ {t + 1} | \mathrm {p a} _ {\mathcal {G} _ {o _ {i}} ^ {t}} (o _ {i}. s _ {l} ^ {t + 1}), \{\mathrm {p a} _ {\mathcal {G} _ {o _ {i}, o _ {k}} ^ {t}} (o _ {i}. s _ {l} ^ {t + 1}) \} _ {o _ {k} \to o _ {i} \in \mathcal {G} _ {\mathrm {i n t e r}} ^ {t}}\right)
$$
where  $s^t$  is the collection of all attributes of all objects at time  $t$ ,  $\theta \in \Theta$  is the collection of all latent constant parameters for all objects,  $\mathbf{a}^t \in \mathcal{A}$  is the action. In the right-hand term,  $o_i.s_l^{t+1}$  is attribute  $s_l$  of object  $o_i$  at time  $t+1$ , while  $\mathrm{pa}_{\mathcal{G}_{o_i}^t}(o_i.s_l^{t+1})$  are the parents of the attribute  $l$  for
object  $o_i$  based on the class template graph  $\mathcal{G}_{C(i)}$ , where  $C(i)$  is the class of  $o_i$ , and where the action binding  $\alpha^t$  activates any potential connections from  $\mathbf{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 \to o_i \in \mathcal{G}_{\mathrm{inter}}^t$ . For each of these objects  $o_k$  we collect the attributes that interact with  $o_i$ .  $s_l$  in the instance interaction pattern  $\mathcal{G}_{o_i,o_k}^t$  based on interaction pattern graph  $\mathcal{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  $\mathcal{R}(\mathbf{s}^t,\mathbf{a}^t,\boldsymbol{\theta}) = \mathcal{R}(\{\mathrm{pa}_{\mathcal{G}_{o_i}}(r^t)\}_{o_i \in \mathcal{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  $\mathcal{C}$  are known and fixed across environments, while the objects  $\mathcal{O}$  can vary, as can the latent parameters  $\theta$ . In the training phase, we will learn how to classify objects, the transition and reward functions based on the class template graphs  $\{\mathcal{G}_{C_j}\}_{C_j\in \mathcal{C}}$  and the interaction patterns  $\{\mathcal{G}_{C_i,C_j}\}_{C_i,C_j\in \mathcal{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  $\mathcal{G}_{\mathrm{inter}}^t$ , 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 multi-object 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.
Figure 2: The learning pipelines of the DAFT-RL framework.
In all steps and all environments, if the input is an image, we extract the objects and their attributes  $\mathbf{s}^t = \{\mathbf{o}_1^t,\mathbf{o}_2^t,\dots ,\mathbf{o}_m^t\}$  from sequences of images  $\mathbf{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  $\mathbf{s}^t = \{\mathbf{o}_1^t,\mathbf{o}_2^t,\dots ,\mathbf{o}_m^t\}$ . For each set of objects  $\mathcal{O}$ , we learn to classify the objects into their classes  $\mathcal{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 \in C_j$  as  $\{\mathbf{o}_i^t, \mathbf{a}_i^t, r_i^t\}_{t=1}^T$ , in environments in which there is only object  $o_i$ , and a random policy is used to generate actions. We denote these data as  $\mathcal{D}^{\mathrm{single}} = \{\{\mathbf{o}_i^t, \mathbf{a}_i^t, r_i^t\}_{t=1}^T\}_{i \in C_j, \forall j=1,\dots,k, i=1,\dots,m}$ .
We initiate the class template  $\mathcal{G}_{C_j}$  for each class  $C_j$  randomly, and then use  $\mathcal{D}^{\mathrm{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_{\lambda}\left(\mathbf{o}_i^{t + 1},r_i^t\mid \mathbf{o}_i^t,\mathbf{a}_i^t,\mathcal{G}^{C_j}\right)$ , where  $\lambda = \{\lambda_s,\lambda_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  $\{\mathrm{pa}_{\mathcal{G}_{C_j}}(\mathbf{o}_i^{t + 1})\}_{i = 1}^m$  and  $\{\mathrm{pa}_{\mathcal{G}_{C_j}}(r_i^t)\}_{i = 1}^m$ , and the GRU outputs  $\mathbf{o}_i^{t + 1}$  and  $r_i^t$ . 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: