SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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:
$$
\mathcal {L} _ {\text {p r e d}} ^ {3} = \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 {A5}
$$
$$
C _ {i}. G _ {\theta \rightarrow s _ {l}} \odot o _ {i}. \theta , \left\{G _ {\operatorname {i n t e r} (i, j)} ^ {t} \cdot G _ {C _ {j} \rightarrow C _ {i}. s _ {l}} \odot \mathbf {o} _ {j} ^ {t} \right\} _ {j \in \{1, \dots , m \} \backslash i})
$$
$$
\mathcal {L} _ {\mathrm {K L}} ^ {3} = \sum_ {t = 2} ^ {T} \operatorname {K L} \left(q _ {\phi} \left(\mathbf {z} ^ {t} \mid \mathbf {s} ^ {1: T}\right) \| p _ {\phi} \left(\boldsymbol {z} ^ {t} \mid \boldsymbol {s} ^ {1: t}, \boldsymbol {z} ^ {1: t - 1}\right)\right) \tag {A6}
$$
$$
\mathcal {L} _ {\text {s p a r s e}} ^ {3} = \sum_ {i = 1} ^ {k} \sum_ {j = 1} ^ {k} \sum_ {l = 1} ^ {n} \left\| G _ {C _ {i} \rightarrow C _ {j}. s _ {l}} \right\| _ {1} + \sum_ {i = 1} ^ {k} \sum_ {l = 1} ^ {n} \| C _ {i}. G _ {\theta \rightarrow s _ {l}} \| _ {1} \tag {A7}
$$
where  $\lambda_{s}$  is the dynamics model we have employed in the previous steps. We learn the encoder/prior networks  $\phi$  to generate the latent vectors  $\mathbf{z}$ , where we sample the graph  $G_{\mathrm{inter}}$  from. Through the binary vectors we also learn the interaction pattern graphs  $\{G_{C_i,C_j}\}_{i,j\in \{1,\dots,k\}}$ , latent parameters  $\{o_i.\theta \}_{i = 1}^m$  and the binary vectors representing the edges from latent parameters to each attribute  $\{C_i.G_{\theta \rightarrow s_l}\}_{i = 1}^k$ . We have completed the loss for step 2.2:
$$
\mathcal {L} ^ {3} = \mathcal {L} _ {\text {p r e d}} ^ {3} + \gamma_ {\mathrm {K L}} ^ {3} \mathcal {L} _ {\mathrm {K L}} ^ {3} + \gamma_ {\text {s p a r s e}} ^ {3} \mathcal {L} _ {\text {s p a r s e}} ^ {3} \tag {A8}
$$
where  $\gamma_{\mathrm{KL}}^3$  and  $\gamma_{\mathrm{sparse}}^3$  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  $\mathbf{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  $\mathbf{z}^t$  at each time step  $t$ . We sample the edge matrix  $\mathbf{M}^t$  of the graph  $\mathcal{G}_{\mathrm{inter}}^t$ . Specifically,  $\mathbf{M}^t \sim \mathrm{Bern}(\mathbf{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.
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 object-centric 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 object-specific 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  $\mathbf{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 Sprinteworld 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\} \cup \{3,5,7,9\}$  and  $\{0.1,0.2,1.1,1.3\} \cup \{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 \times 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 Sprinteworld 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 \times 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.
Experiment Settings	Method
DAFT (Symbolic)	Deep Set	Self-attention	SRICS	GNN	STOVE	Self-attention	NCS	Relational RL
2-Push	0.959 ± 0.031	0.985 ± 0.025	0.943 ± 0.024	0.981 ± 0.015	0.923 ± 0.047	0.973 ± 0.028	0.968 ± 0.036	0.993 ± 0.013	0.916 ± 0.054
2-Switch	0.982 ± 0.013	0.869 ± 0.032	0.954 ± 0.015	0.978 ± 0.029	0.931 ± 0.028	0.916 ± 0.045	0.943 ± 0.016	0.977 ± 0.038	0.945 ± 0.042
3-Push	0.961 ± 0.035	0.753 ± 0.031	0.670 ± 0.023	0.931 ± 0.027	0.784 ± 0.041	0.954 ± 0.033	0.972 ± 0.019	0.893 ± 0.024	0.929 ± 0.022
3-Switch	0.907 ± 0.066	0.879 ± 0.077	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 A3: Average success rate over 3 random seeds for Push & Switch environments testing (single task training mode). The numbers in bold highlight the top-performing method.
Experiment Settings	Method
DAFT (SA)	DAFT (AIR)	SMORL	SRICS	GNN	STOVE	COBRA	NCS	LRN
Object Goal	0.916 ± 0.032	0.920 ± 0.037	0.745 ± 0.066	0.784 ± 0.058	0.464 ± 0.091	0.643 ± 0.049	0.715 ± 0.049	0.925 ± 0.021	0.846 ± 0.067
Object Interaction	0.909 ± 0.068	0.896 ± 0.053	0.758 ± 0.063	0.824 ± 0.096	0.396 ± 0.146	0.683 ± 0.094	0.746 ± 0.073	0.931 ± 0.061	0.812 ± 0.115
Object Comparison	0.917 ± 0.065	0.902 ± 0.060	0.923 ± 0.049	0.812 ± 0.071	0.476 ± 0.128	0.625 ± 0.091	0.738 ± 0.057	0.901 ± 0.051	0.693 ± 0.097
Property Comparison	0.930 ± 0.034	0.905 ± 0.075	0.918 ± 0.088	0.810 ± 0.095	0.369 ± 0.174	0.602 ± 0.108	0.732 ± 0.083	0.897 ± 0.112	0.644 ± 0.091
Table A4: Average success rate over 3 random seeds for Spriteworld environments training (single task training mode). The numbers in bold highlight the top-performing method.
Experiment Settings	Method
DAFT (SA)	DAFT (AIR)	SMORL	SRICS	GNN	STOVE	COBRA	NCS	LRN
Object Goal	0.928 ± 0.038	0.899 ± 0.043	0.720 ± 0.069	0.766 ± 0.064	0.450 ± 0.095	0.672 ± 0.052	0.699 ± 0.055	0.883 ± 0.028	0.791 ± 0.071
Object Interaction	0.900 ± 0.074	0.886 ± 0.059	0.735 ± 0.067	0.802 ± 0.100	0.380 ± 0.142	0.665 ± 0.100	0.728 ± 0.077	0.851 ± 0.066	0.789 ± 0.119
Object Comparison	0.904 ± 0.071	0.890 ± 0.065	0.876 ± 0.054	0.795 ± 0.076	0.459 ± 0.129	0.674 ± 0.096	0.724 ± 0.061	0.864 ± 0.058	0.800 ± 0.102
Property Comparison	0.911 ± 0.070	0.875 ± 0.081	0.865 ± 0.093	0.782 ± 0.099	0.355 ± 0.173	0.680 ± 0.113	0.711 ± 0.087	0.875 ± 0.116	0.772 ± 0.097
Table A5: Average success rate over 3 random seeds for Spriteworld environments testing (with unseen object numbers). The numbers in bold highlight the top-performing method.
Experiment Settings	Method
DAFT (SA)	DAFT (AIR)	SMORL	SRICS	GNN	STOVE	COBRA	NCS	LRN
Object Goal	0.902 ± 0.036	0.897 ± 0.041	0.694 ± 0.071	0.738 ± 0.062	0.420 ± 0.099	0.654 ± 0.050	0.681 ± 0.053	0.835 ± 0.024	0.775 ± 0.069
Object Interaction	0.895 ± 0.072	0.880 ± 0.057	0.710 ± 0.065	0.772 ± 0.102	0.350 ± 0.140	0.635 ± 0.098	0.698 ± 0.075	0.859 ± 0.064	0.769 ± 0.117
Object Comparison	0.878 ± 0.069	0.895 ± 0.063	0.854 ± 0.052	0.765 ± 0.074	0.429 ± 0.125	0.631 ± 0.092	0.704 ± 0.059	0.824 ± 0.056	0.724 ± 0.100
Property Comparison	0.912 ± 0.068	0.867 ± 0.079	0.885 ± 0.091	0.742 ± 0.097	0.325 ± 0.171	0.660 ± 0.109	0.690 ± 0.085	0.809 ± 0.114	0.725 ± 0.093
Table A6: Average success rate over 3 random seeds for Sprinteworld environments testing (with unseen object colors and shapes). The numbers in bold highlight the top-performing method.
D.3.2 Full Ablation Studies
- 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;
Fig. A3 gives the full ablation studies on Push & Switch and Stacking benchmarks. We consider the following cases:
Figure A3: Ablation studies on Push & Switch and stacking tasks. I. without latent parameters; II. without factored class template graph; III. without dynamic interaction graph; IV. without factored interaction pattern; V. using single object class; VI. without action binding graph; VII. using hard attention networks; VIII. Original DAFT-RL. For those tasks where the latent parameters do not exist, we did not conduct ablations on case I. Similarly, for those tasks with only one object, we did not include the ablations on case V.
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.
Task	DAFT (SA)	DAFT (SA) w/o class template graphs
Object goal	0.897	0.767
Object interaction	0.890	0.739
Object comparison	0.893	0.685
Property comparison	0.907	0.832
For Sprinteworld, 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  $\mathcal{D}^{\mathrm{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.
Table A7: Ablation studies on the Spriteworld benchmark.
Experiment Settings	Method
2-P+2-S (S)	3-P+3-S (S+O)	2-P (L)	2-S (L)	3-P (L+O)	3-S (L+O)	2-P+2-S (L+S)	3-P+3-S (L+O+S)
Original (main paper)	0.881 ± 0.038	0.805 ± 0.024	0.968 ± 0.036	0.923 ± 0.065	0.921 ± 0.037	0.903 ± 0.023	0.793 ± 0.026	0.783 ± 0.025
End-to-end	0.821 ± 0.076	0.765 ± 0.023	0.806 ± 0.048	0.852 ± 0.038	0.875 ± 0.029	0.796 ± 0.047	0.635 ± 0.024	0.707 ± 0.062
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.
Table A8: Results on end-to-end training in Push & Switch compositional task. m-P and n-S denote the m-Push and n-Switch task. The numbers in bold highlight the top-performing method.
Experiment Settings	Method
DAFT-RL (SA)	DAFT-RL (AIR)	SMORL	SRICS	GNN	STOVE	NCS	LRN
Original	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
No imagination	0.469 ± 0.063	0.508 ± 0.064	0.401 ± 0.055	0.434 ± 0.078	0.319 ± 0.040	0.224 ± 0.071	0.264 ± 0.059	0.398 ± 0.061
Table A9: Results on the cases where imagination is disabled. Average success rate over 3 random seeds for Block-stacking with 8 boxes and different masses.
D.6 Effect of the Choice of Object-centric Models for Symbolic Methods
Table A10 shows the results of adding AIR and SA to the symbolic methods on the 8-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
Experiment Settings	Method
DAFT-RL	SRICS	GNN	LRN
AIR	0.506 ± 0.083	0.420 ± 0.061	0.334 ± 0.047	0.463 ± 0.077
SA	0.571 ± 0.039	0.395 ± 0.046	0.293 ± 0.036	0.482 ± 0.052
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.