SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\operatorname {a r g m a x} _ {\lambda , \left\{\mathcal {G} ^ {C _ {j}} \right\} _ {j = 1} ^ {k}} \sum_ {t = 1} ^ {T} \sum_ {i = 1} ^ {m} \sum_ {l = 1} ^ {n} \left(\log p _ {\lambda} \left(o _ {i}. s _ {l} ^ {t + 1}, r _ {i} ^ {t} \mid \mathrm {p a} _ {\mathcal {G} _ {C (i)}} \left(o _ {i}. s _ {l} ^ {t + 1}\right), \mathrm {p a} _ {\mathcal {G} _ {C (i)}} \left(r _ {i} ^ {t}\right)\right) - \sum_ {j = 1} ^ {k} \left\| \mathcal {G} ^ {C _ {j}} \right\| _ {1} \right.
$$
where  $m$  and  $k$  indicate the number of single-object environments and object type classes, respectively, and  $\mathrm{pa}_{\mathcal{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  $\{\mathcal{G}^{C_1},\mathcal{G}^{C_2},\ldots ,\mathcal{G}^{C_k}\}$  with the exception of the edges from the latent parameters  $\theta$ , which here we assume are disabled and we will learn in the next step. In later stages, we will reuse the learned reward model  $\lambda_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  $\mathcal{D}^{\mathrm{multi}}$  from  $N$  multi-object environments with different object configurations and in which the objects can have varying latent parameters. Formally, we define  $\mathcal{D}^{\mathrm{multi}} = \{\{\mathbf{o}_1^t,\mathbf{a}_1^t,r_1^t,\mathbf{o}_2^t,\mathbf{a}_2^t,r_2^t,\dots ,\mathbf{o}_m^t,\mathbf{a}_m^t,r_m^t\}_{t = 1}^T\}_{m = 1}^N$ . 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  $\mathcal{G}_{C_i,C_j}$  for each pair of classes  $C_i$  and  $C_j$  and the dynamic interaction graph  $\mathcal{G}_{\mathrm{inter}}$  by exploiting the previously learned class template graphs. In particular, we first learn the action binding  $\alpha^{t}$ , and at a second stage, we jointly learn the rest of the dynamic interaction graph  $\mathcal{G}_{\mathrm{inter}}$ , the interaction patterns  $\mathcal{G}_{C_i,C_j}$  for each pair of classes, the object-specific latent parameters  $o_i.\theta$  and their edges to the other attributes in  $\mathcal{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  $\alpha = \{\alpha^1, \alpha^2, \ldots, \alpha^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  $\mathbf{k}^t = \langle f_k(\mathbf{o}_1^t), f_k(\mathbf{o}_2^t), \ldots, f_k(\mathbf{o}_m^t) \rangle$ , query  $\mathbf{q}^t = f_q(\mathbf{a}^t)$ , and value  $\mathbf{v}^t = \langle f_v(\mathbf{o}_1^t), f_v(\mathbf{o}_2^t), \ldots, f_v(\mathbf{o}_m^t) \rangle$ , respectively. We then compute the attention weights  $\alpha^t = \text{softmax}((\mathbf{k}_1^t)^\top \mathbf{q}^t, (\mathbf{k}_2^t)^\top \mathbf{q}^t, \ldots, (\mathbf{k}_m^t)^\top \mathbf{q}^t)$ . We use the learned attention weights  $\alpha$  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  $\mathbf{v}^t$  with the attention weights  $\alpha^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  $\mathbf{h}^t = \mathrm{GNN}(\mathbf{o}_1^t,\mathbf{o}_2^t,\dots,\mathbf{o}_m^t)$ . For each object pair  $o_i$  and  $o_j$ , we then obtain  $\mathbf{h}_{(i,j)}^t$ . 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:  $\mathbf{h}_{(i,j),\mathrm{enc}}^t = \mathrm{GRU_{enc}}(\mathbf{h}_{(i,j)}^t,\mathbf{h}_{(i,j)}^{t + 1})$ . For the prior, we also use a GRU, but in this context, we append  $\mathbf{h}_{(i,j)}^t$  with the output of the GRU from the previous step  $\mathbf{h}_{(i,j),\mathrm{prior}}^{t - 1}$  as the input. In specific terms, we have  $\mathbf{h}_{(i,j),\mathrm{prior}}^t = \mathrm{GRU_{prior}}(\mathbf{h}_{(i,j)}^t,\mathbf{h}_{(i,j),\mathrm{prior}}^{t - 1})$ . 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_{\phi}(\mathbf{z}^t\mid \mathbf{s}^{1:T})$ , where  $\mathbf{s}^{1:T} = \{\mathbf{o}_1^t,\dots,\mathbf{o_m}^t\}_{t = 1}^T$ , and prior distribution  $p_{\phi}(\boldsymbol {z}^t\mid \boldsymbol {s}^{1:t},\boldsymbol {z}^{1:t - 1})$ , respectively. We assume that the encoder outputs a Bernoulli distribution for each edge and the graph  $\mathcal{G}_{\mathrm{inter}}$  is sampled using the Gumbel-Softmax trick [40].
Decoder We perform one-step prediction to generate  $\hat{\mathbf{o}}_i^{t + 1}$  for each object  $o_i$ , predicting the state dynamics with the learned graphs, including  $\mathcal{G}_{\mathrm{inter}}$ , as well as the interaction pattern graph  $\mathcal{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,  $\pmb{\theta}$  at this stage. Specifically, we also incorporate  $C_j.\pmb{\theta}$  and the graph from  $\mathcal{G}_{C_j,\pmb{\theta},\mathbf{o}_i^t}$  for each object  $o_i$  with class  $C_j$  into the dynamics model and use the sparsity regularization for  $\mathcal{G}_{C_j,\pmb{\theta},\mathbf{o}_i}$ .
Therefore, the learning objectives include maximizing the likelihood of the dynamics model, minimizing the KL divergence between  $p_{\phi}$  and  $q_{\phi}$  to estimate  $\pmb{\theta}$ , and encouraging the sparsity of the interaction pattern graph  $\mathcal{G}_{C_i,C_j}$  and the subgraph from latent parameters to states  $\mathcal{G}_{C_i,\pmb{\theta},\mathbf{o}_i}$ :
$$
\begin{array}{l} \operatorname {a r g m a x} _ {\lambda_ {s}, \phi , v, k, q, \theta , \mathcal {G}} \sum_ {t = 1} ^ {T} \sum_ {i = 1} ^ {m} \sum_ {l = 1} ^ {n} \log p _ {\lambda_ {s}} \left(o _ {i}. s _ {l} ^ {t + 1} \mid \operatorname {p a} _ {\mathcal {G} _ {o _ {i}} ^ {t}} \left(o _ {i}. s _ {l} ^ {t + 1}\right), \left\{\operatorname {p a} _ {G _ {o _ {i}, o _ {k}} ^ {t}} \left(o _ {i}. s _ {l} ^ {t + 1}\right)\right\} _ {o _ {k} \rightarrow o _ {i} \in \mathcal {G} _ {\text {i n t e r}} ^ {t}}\right) \\ - \sum_ {j = 1} ^ {k} \sum_ {i = 1} ^ {k} \left\| \mathcal {G} _ {C _ {i}, C _ {j}} \right\| _ {1} - \sum_ {i = 1} ^ {m} \sum_ {j = 1} ^ {k} \left\| \mathcal {G} _ {C _ {j}. \boldsymbol {\theta}, \boldsymbol {\mathbf {o}} _ {i}} \right\| _ {1} - \sum_ {t = 2} ^ {T} \operatorname {K L} \left(q _ {\phi} \left(\boldsymbol {\mathbf {z}} ^ {t} \mid \boldsymbol {\mathbf {s}} ^ {1: T}\right) \| p _ {\phi} \left(\boldsymbol {\mathbf {z}} ^ {t} \mid \boldsymbol {\mathbf {s}} ^ {1: t}, \boldsymbol {\mathbf {z}} ^ {1: t - 1}\right)\right) \\ \end{array}
$$
where  $\mathcal{G}$  is the dynamic interaction graph  $\mathcal{G}_{\mathrm{inter}}$ , interaction pattern graphs  $\{\mathcal{G}_{C_i,C_j} \mid i \in \{1,2,\ldots,k\}, j \in \{1,2,\ldots,k\}\}$  and the subgraph from latent parameters to states  $\{\mathcal{G}_{C_j,\theta,\mathbf{o}_i} \mid i \in \{1,2,\ldots,m\}, j \in \{1,2,\ldots,k\}\}$ . Similarly, to Definition 1,  $\mathrm{pa}_{\mathcal{G}_{o_i}^t}(o_i.s_l^{t+1})$  indicates the parents of the attribute  $s_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 or deactivates 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)$ .  $\lambda_s$  and  $\phi$  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  $\theta$  and the interaction graph  $\mathcal{G}_{\mathrm{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  $\pi^{*}(\mathbf{a}^{t}|\mathbf{s}^{t},\boldsymbol {\theta},\mathcal{G}_{\mathrm{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  $\pi^{*}$ . 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  $\pi^{*}(\mathbf{a}^{t}|\mathbf{s}^{t},\pmb{\theta},\mathcal{G}_{\mathrm{inter}})$  by inferring latent parameters  $\pmb{\theta}$  and dynamic interaction graphs  $\mathcal{G}_{\mathrm{inter}}$  based on a few trajectories, without any policy learning.
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.
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
Table 1: Average success rate over 3 random seeds for Push & Switch compositional generalization in terms of combination of skills (S), changing number of objects (O), and changing latent parameters (L) with respect to training. The numbers in bold highlight the top-performing method.
Experiment Settings	Method
DAFT-RL (SA)	DAFT-RL (AIR)	SMORL	SRICS	GNN	STOVE	COBRA	NCS	LRN
Object Goal	0.897 ± 0.034	0.891 ± 0.039	0.668 ± 0.073	0.710 ± 0.060	0.390 ± 0.101	0.636 ± 0.048	0.663 ± 0.051	0.847 ± 0.022	0.759 ± 0.067
Object Interaction	0.890 ± 0.070	0.874 ± 0.055	0.685 ± 0.063	0.742 ± 0.104	0.320 ± 0.138	0.605 ± 0.096	0.668 ± 0.073	0.837 ± 0.062	0.749 ± 0.115
Object Comparison	0.893 ± 0.067	0.879 ± 0.061	0.834 ± 0.050	0.735 ± 0.072	0.399 ± 0.121	0.724 ± 0.088	0.707 ± 0.057	0.716 ± 0.054	0.760 ± 0.098
Property Comparison	0.907 ± 0.066	0.879 ± 0.073	0.805 ± 0.089	0.732 ± 0.095	0.295 ± 0.169	0.640 ± 0.105	0.660 ± 0.083	0.831 ± 0.112	0.748 ± 0.089
Table 2: Average success rate over 3 random seeds for Spriteworld with unseen object numbers, color and shape combinations. The numbers in bold highlight the top-performing method.
in the difficult settings. In this setting, we do not compare with COBRA [9], since it expects pixel inputs. As a representative example, we show the smoothed learning curve for 2-Push+2-Switch  $(\mathrm{L} + \mathrm{S})$  in Fig. 3A, in which we display the top three methods in terms of success rate.
Image benchmark: Spriteworld We consider four tasks in Spriteworld [9]: object goal, interaction, comparison, and property comparison. Given that the observations are in pixel space, we use 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 object-centric model. We consider cases in which we vary the number of objects, have unseen combinations of colors for the objects or unseen combinations of shapes in the target task, and provide the results in Table 2, showing that DAFT-RL with Slot Attention outperforms all baselines, followed closely by DAFT-RL with AIR. In this setting, we do not show the results from DeepSets or Self-Attention, because of bad performance, but we provide the results in App. D.3. We show the smoothed learning curve for the object comparison task in Fig. 3B, in which we only display the top three methods.
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  $(\mathrm{L} + \mathrm{O} + \mathrm{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
Figure 4: Ablation study on the 3-Push+3-Switch task.
Experiment Settings	Method
DAFT-RL (SA)	DAFT-RL (AIR)	SMORL	SRICS	GNN	STOVE	NCS	LRN
2 Blocks	0.809 ± 0.019	0.838 ± 0.030	0.658 ± 0.028	0.704 ± 0.016	0.549 ± 0.016	0.728 ± 0.044	0.797 ± 0.035	0.649 ± 0.026
4 Blocks	0.735 ± 0.032	0.698 ± 0.022	0.605 ± 0.020	0.591 ± 0.049	0.526 ± 0.041	0.498 ± 0.013	0.571 ± 0.026	0.461 ± 0.028
6 Blocks	0.591 ± 0.025	0.664 ± 0.017	0.536 ± 0.040	0.509 ± 0.043	0.461 ± 0.088	0.475 ± 0.023	0.521 ± 0.049	0.602 ± 0.097
8 Blocks	0.506 ± 0.083	0.571 ± 0.039	0.386 ± 0.062	0.420 ± 0.061	0.334 ± 0.047	0.278 ± 0.086	0.397 ± 0.052	0.463 ± 0.077
Table 3: Average success rate over 3 random seeds for Block-stacking with unseen object numbers and mass combinations. The numbers in bold highlight the top-performing method.
Figure 3: A. The smoothed learning curve for 2-Push + 2-Switch (L+S) with different friction coefficients for each object (for clarity, we show only the top three methods in terms of the success rate); B. The smoothed learning curve for the object comparison task in Spriteworld with unseen object numbers, combinations of colors and shapes (for clarity, we show only the top three methods in terms of the success rate); C. Success rate versus number of blocks in the stacking task, where each block has distinct mass.
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 open-set 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].