Original: C OMPOSITIONAL
Revised: C OMPOSITIONAL

Original: M ULTI -O BJECT R EINFORCEMENT L EARNING WITH L INEAR R ELATION N ETWORKS
Revised: M ULTI -O BJECT R EINFORCEMENT L EARNING WITH L INEAR R ELATION N ETWORKS

Original: Anonymous authors

Original: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Although reinforcement learning has seen remarkable progress over the last years, solving robust dexterous object-manipulation tasks in multi-object settings remains a challenge.
Revised: Although reinforcement learning has seen remarkable progress over the last years, solving robust dexterous object-manipulation tasks in multi-object settings remains a challenge.

Original: In this paper, we focus on models that can learn manipulation tasks in fixed multi-object settings and extrapolate this skill zero-shot without any drop in performance when the number of objects changes.
Revised: In this paper, we focus on models that can learn manipulation tasks in fixed multi-object settings and extrapolate this skill zero-shot without any drop in performance when the number of objects changes.

Original: We consider the generic task of bringing a specific cube out of a set to a goal position.
Revised: We consider the generic task of bringing a specific cube out of a set to a goal position.

Original: We find that previous approaches, which primarily leverage attention and graph neural network-based architectures, do not generalize their skills when the number of input objects changes while scaling as K 2 .
Revised: We find that previous approaches, which primarily leverage attention and graph neural network-based architectures, do not generalize their skills when the number of input objects changes while scaling as K 2 .

Original: We propose an alternative plug-and-play module based on relational inductive biases to overcome these limitations.
Revised: We propose an alternative plug-and-play module based on relational inductive biases to overcome these limitations.

Original: Besides exceeding performances in their training environment, we show that our approach, which scales linearly in K , allows agents to extrapolate and generalize zero-shot to any new object number.
Revised: Besides exceeding performances in their training environment, we show that our approach, which scales linearly in K , allows agents to extrapolate and generalize zero-shot to any new object number.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Deep reinforcement learning (RL) has witnessed remarkable progress over the last years, particularly in domains such as video games or other synthetic toy settings (Mnih et al., 2015; Silver et al., 2016;
Revised: Deep reinforcement learning (RL) has witnessed remarkable progress over the last years, particularly in domains such as video games or other synthetic toy settings (Mnih et al., 2015; Silver et al., 2016;

Original: Vinyals et al., 2019).
Revised: Vinyals et al., 2019).

Original: On the other hand, applying deep RL on real-world grounded robotic setup such as learning seemingly simple dexterous manipulation tasks in multi-object settings is still confronted with many fundamental limitations being the focus of many recent works (Duan et al., 2017; Janner et al., 2019;
Revised: On the other hand, applying deep RL on real-world grounded robotic setup such as learning seemingly simple dexterous manipulation tasks in multi-object settings is still confronted with many fundamental limitations being the focus of many recent works (Duan et al., 2017; Janner et al., 2019;

Original: Deisenroth et al., 2011; Kroemer et al., 2018; Andrychowicz et al., 2020; Rajeswaran et al., 2018; Lee et al., 2021; Funk et al., 2022).
Revised: Deisenroth et al., 2011; Kroemer et al., 2018; Andrychowicz et al., 2020; Rajeswaran et al., 2018; Lee et al., 2021; Funk et al., 2022).

Original: The reinforcement learning problem in robotics setups is much more challenging (Dulac-Arnold et al., 2021).
Revised: The reinforcement learning problem in robotics setups is much more challenging (Dulac-Arnold et al., 2021).

Original: Compared to discrete toy environments, state and action spaces are continuous, and solving tasks typically requires long-horizon time spans, requiring the agent to apply very long sequences of precise low-level control actions.
Revised: Compared to discrete toy environments, state and action spaces are continuous, and solving tasks typically requires long-horizon time spans, requiring the agent to apply very long sequences of precise low-level control actions.

Original: Accordingly, exploration under easy-to-define sparse rewards becomes nearly impossible without horrendous amounts of data.
Revised: Accordingly, exploration under easy-to-define sparse rewards becomes nearly impossible without horrendous amounts of data.

Original: This is usually impossible in the real world but has been likewise hard to achieve with computationally demanding realistic physics simulators.
Revised: This is usually impossible in the real world but has been likewise hard to achieve with computationally demanding realistic physics simulators.

Original: One promising path is to constrain the possible solution space of a learning agent by encoding suitable inductive biases in their neural architecture (Geman et al., 1992).
Revised: One promising path is to constrain the possible solution space of a learning agent by encoding suitable inductive biases in their neural architecture (Geman et al., 1992).

Original: Choosing good inductive biases that leverage the underlying structure of the problem setting can help to learn solutions that facilitate desired generalization capabilities (Mitchell, 1980; Baxter, 2000; Hessel et al., 2019).
Revised: Choosing good inductive biases that leverage the underlying structure of the problem setting can help to learn solutions that facilitate desired generalization capabilities (Mitchell, 1980; Baxter, 2000; Hessel et al., 2019).

Original: In the case of robotics, environments of multi-object manipulation tasks naturally accommodate a compositional description of their current state in terms of symbol-like entities (such as physical objects, robot parts, etc.).
Revised: In the case of robotics, environments of multi-object manipulation tasks naturally accommodate a compositional description of their current state in terms of symbol-like entities (such as physical objects, robot parts, etc.).

Original: These representations can be directly obtained in simulator settings and ultimately hoped to be inferred robustly from learned object perception modules (Greff et al., 2019; Kipf et al., 2021;
Revised: These representations can be directly obtained in simulator settings and ultimately hoped to be inferred robustly from learned object perception modules (Greff et al., 2019; Kipf et al., 2021;

Original: Locatello et al., 2020).
Revised: Locatello et al., 2020).

Original: While such a compositional understanding is in principle considered crucial for any systematic generalization ability (Greff et al., 2020; Spelke, 1990; Battaglia et al., 2018; Garnelo et al., 2016) it remains an open question how to design an agent that can process this type of input data to leverage this promise.
Revised: While such a compositional understanding is in principle considered crucial for any systematic generalization ability (Greff et al., 2020; Spelke, 1990; Battaglia et al., 2018; Garnelo et al., 2016) it remains an open question how to design an agent that can process this type of input data to leverage this promise.

Original: Consider transporting a cube to a specified target location and assume the agent has mastered solving this task with three available cubes (see Figure 1 left).
Revised: Consider transporting a cube to a specified target location and assume the agent has mastered solving this task with three available cubes (see Figure 1 left).

Original: We refer to cubes different from the one to transport as distractors.
Revised: We refer to cubes different from the one to transport as distractors.

Original: This problem setting forms the basis of our work.
Revised: This problem setting forms the basis of our work.

Original: If we now introduced two additional cubes, the solution to solve this task would remain the same, but the input distribution changed.
Revised: If we now introduced two additional cubes, the solution to solve this task would remain the same, but the input distribution changed.

Original: So why not just train with much more cubes and use a sufficiently large input state?
Revised: So why not just train with much more cubes and use a sufficiently large input state?

Original: First, the learning problem requires exponentially more data when solving such a task with more available cubes (see Figure 10, red line), making it infeasible already for only half a dozen cubes.
Revised: First, the learning problem requires exponentially more data when solving such a task with more available cubes (see Figure 10, red line), making it infeasible already for only half a dozen cubes

Original: Second we do not want to put any constraint on the possible number of cubes at test time but simply might not have this flexibility at train time.
Revised: Second, we do not want to put any constraint on the possible number of cubes at test time but simply might not have this flexibility at train time.

Original: Therefore, we demand to learn challenging manipulation tasks in fixed multi-cube settings and extrapolate this skill zero-shot without any drop in performance when the number of cubes changes.
Revised: Therefore, we demand to learn challenging manipulation tasks in fixed multi-cube settings and extrapolate this skill zero-shot without any drop in performance when the number of cubes changes.

Original: Achieving this objective will be the primary goal of this work.
Revised: Achieving this objective will be the primary goal of this work.

Original: An essential prerequisite for achieving this is to endow the agents with the ability to process a variable number of input objects by choosing a suitable inductive bias.
Revised: An essential prerequisite for achieving this is to endow the agents with the ability to process a variable number of input objects by choosing a suitable inductive bias.

Original: One potential model class for these object encoding modules are graph neural networks (GNNs) and a vast line of previous approaches builds upon the subclass of attentional GNNs to achieve this (Zambaldi et al., 2018;
Revised: One potential model class for these object encoding modules are graph neural networks (GNNs) and a vast line of previous approaches builds upon the subclass of attentional GNNs to achieve this (Zambaldi et al., 2018;

Original: Li et al., 2020; Wilson and Hermans, 2020; Zadaianchuk et al., 2021).
Revised: Li et al., 2020; Wilson and Hermans, 2020; Zadaianchuk et al., 2021).

Original: As a first contribution, we will show in Section 4 that utilizing such attention-based GNN approaches will not achieve the desiderata described above on our problem setting.
Revised: As a first contribution, we will show in Section 4 that utilizing such attention-based GNN approaches will not achieve the desiderata described above on our problem setting.

Original: Attentional approaches typically fail to extrapolate any learned behavior to a changing number of objects and require fine-tuning, resulting in catastrophic forgetting of previous skills.
Revised: Attentional approaches typically fail to extrapolate any learned behavior to a changing number of objects and require fine-tuning, resulting in catastrophic forgetting of previous skills.

Original: Lastly, using attention scales quadratically in the number of input objects which becomes impractical for training and inference with many objects.
Revised: Lastly, using attention scales quadratically in the number of input objects which becomes impractical for training and inference with many objects.

Original: Main contributions.
Revised: Main contributions.

Original: In this work, we present support to build agents upon relational reasoning inductive biases.
Revised: In this work, we present support to build agents upon relational reasoning inductive biases.

Original: Specifically, we propose an object encoder module based on a linear relation neural network (LRN) that is able to learn challenging manipulation tasks in fixed multi-object settings.
Revised: Specifically, we propose an object encoder module based on a linear relation neural network (LRN) that is able to learn challenging manipulation tasks in fixed multi-object settings.

Original: We show that our module consistently outperforms an attention-based baseline while scaling only linear in the number of objects.
Revised: We show that our module consistently outperforms an attention-based baseline while scaling only linear in the number of objects.

Original: Most importantly, we show that our agents can extrapolate and therefore generalize zero-shot without any drop in performance when the number of objects changes.
Revised: Most importantly, we show that our agents can extrapolate and therefore generalize zeroshot without any drop in performance when the number of objects changes.

Original: We present supporting evidence for this compositional generalization capability and necessary requirements for multi-object manipulation in two challenging task environments.
Revised: We present supporting evidence for this compositional generalization capability and necessary requirements for multi-object manipulation in two challenging task environments.

Original:  P ROBLEM S ETTING AND R ELATED W ORK
Revised:  P ROBLEM S ETTING AND R ELATED W ORK

Original: Consider the following problem setting that we summarize in Figure 1: There are 3 identical cubes in an arena.
Revised: Consider the following problem setting that we summarize in Figure 1: There are 3 identical cubes in an arena.

Original: A random cube is selected to be transported to an episode-specific location during each episode.
Revised: A random cube is selected to be transported to an episode-specific location during each episode.

Original: The remaining 2 cubes are then distractors for this episode.
Revised: The remaining 2 cubes are then distractors for this episode.

Original: This work aims to answer the following question: Can we train RL agents that learn such manipulation tasks from a fixed multi-cube setting and extrapolate this skill zero-shot when the number of distractors changes?
Revised: This work aims to answer the following question: Can we train RL agents that learn such manipulation tasks from a fixed multi-cube setting and extrapolate this skill zero-shot when the number of distractors changes?

Original: This work considers tasks where the goal is always to transport one of the N o movable objects, i.e. cubes, to a target location, and all other N o − 1 cubes act as distractors.
Revised: This work considers tasks where the goal is always to transport one of the N o movable objects, i.e. cubes, to a target location, and all other N o − 1 cubes act as distractors.

Original: Therefore, given their straightforward relation, we refer to each environment either using the total number of cubes or the number of distractors, the latter called N .
Revised: Therefore, given their straightforward relation, we refer to each environment either using the total number of cubes or the number of distractors, the latter called N .

Original: Specifically, we define the objective of any learning agent as follows:
Revised: Specifically, we define the objective of any learning agent as follows:

Original: First, the agent needs to learn to solve the task at train time under a constant and a minimal number of manipulable objects N traino .
Revised: First, the agent needs to learn to solve the task at train time under a constant and a minimal number of manipulable objects N traino .

Original: Second, the agent needs to be able to zero-shot extrapolate solving the equivalent task under N testo ̸
Revised: Second, the agent needs to be able to zero-shot extrapolate solving the equivalent task under N testo ̸

Original: = N traino objects in the scene.
Revised: = N traino objects in the scene.

Original: Therefore, the agent needs to learn an invariant behavior with respect to the episode-specific distractors.
Revised: Therefore, the agent needs to learn an invariant behavior with respect to the episode-specific distractors.

Original: Note that our task formulation can be seen as generic subtasks of more general manipulation tasks.
Revised: Note that our task formulation can be seen as generic subtasks of more general manipulation tasks.

Original: If we can solve this task, any more complex multi-object goal can be solved by specifying one object-goal pose after another.
Revised: If we can solve this task, any more complex multi-object goal can be solved by specifying one object-goal pose after another.

Original: Crucially, we want to learn such tasks with a constant number of available objects because varying the number of objects during training might limit us in the real world.
Revised: Crucially, we want to learn such tasks with a constant number of available objects because varying the number of objects during training might limit us in the real world.

Original: We ultimately need adapting agents that can generalize when additional objects are added or removed from the environment scene.
Revised: We ultimately need adapting agents that can generalize when additional objects are added or removed from the environment scene.

Original: We begin by choosing train environment with at least three movable objects to let agents also experience potentially important three-body interactions during training.
Revised: We begin by choosing train environment with at least three movable objects to let agents also experience potentially important three-body interactions during training.

Original: Background on compositional multi-object manipulation tasks.
Revised: Background on compositional multi-object manipulation tasks.

Original: Given the compositional environment structure of such real-world grounded manipulation tasks, we are concerned with input states s t ∈ S that comprise a set of objects describing individual rigid objects in the scene such as movable cubes { s o 1 , s o 2 , ..., s oN o } or the individual parts of a robot such as each finger in our setup { s r 1 , s r 2 , ..., s rN r } .
Revised: Given the compositional environment structure of such real-world grounded manipulation tasks, we are concerned with input states s t ∈ S that comprise a set of objects describing individual rigid objects in the scene such as movable cubes { s o 1 , s o 2 , ..., s oN o } or the individual parts of a robot such as each finger in our setup { s r 1 , s r 2 , ..., s rN r } .

Original: To facilitate a goal-conditioned formulation, we similarly feed additional environment-specific information in the form of a goal object s g .
Revised: To facilitate a goal-conditioned formulation, we similarly feed additional environment-specific information in the form of a goal object s g .

Original: Concerning the objective of our tasks, the goal information can be provided via the object state s g that describes the target pose of this cube.
Revised: Concerning the objective of our tasks, the goal information can be provided via the object state s g that describes the target pose of this cube.

Original: To identify the object to be transported, we can simply add a unique identifier feature to each object state and specify the object to be moved through the same identifier in the goal object.
Revised: To identify the object to be transported, we can simply add a unique identifier feature to each object state and specify the object to be moved through the same identifier in the goal object.

Original: Translated to a robotic RL setting, the above state-space factorization requires a policy that is parameterized by a function a t =
Revised: Translated to a robotic RL setting, the above state-space factorization requires a policy that is parameterized by a function a t =

Original: π ( { s o 1 , ..., s oN o , s r 1 , ..., s rN r , s g } ) .
Revised: π ( { s o 1 , ..., s oN o , s r 1 , ..., s rN r , s g } ) .

Original: Our demand is that this policy can deal with a variable number N o of input objects independent of their order.
Revised: Our demand is that this policy can deal with a variable number N o of input objects independent of their order.

Original: In other words, the policy needs to be permutation-invariant.
Revised: In other words, the policy needs to be permutation-invariant.

Original: Domain specific related works.
Revised: Domain specific related works.

Original: Flavours of this problem have a long history and are typically summarized as relational reinforcement learning, which tries to approach this by various suitable inductive biases that enable reasoning upon such object-oriented structured representations (Dˇzeroski et al., 1998; Dˇzeroski et al., 2001; Kaelbling et al., 2001; Van Otterlo, 2002).
Revised: Flavours of this problem have a long history and are typically summarized as relational reinforcement learning, which tries to approach this by various suitable inductive biases that enable reasoning upon such object-oriented structured representations (D ˇ zeroski et al., 1998; D ˇ zeroski et al., 2001; Kaelbling et al., 2001; Van Otterlo, 2002).

Original: A popular class of neural architecture that offers these properties are various sub-types of graph neural networks (GNNs) with attention-based approaches or interaction networks being one special case (Bronstein et al., 2021; Veliˇckovi´c et al., 2018; Vaswani et al., 2017; Battaglia et al., 2016; Santoro et al., 2017;
Revised: A popular class of neural architecture that offers these properties are various sub-types of graph neural networks (GNNs) with attention-based approaches or interaction networks being one special case (Bronstein et al., 2021; Veli ˇ ckovi ´ c et al., 2018; Vaswani et al., 2017; Battaglia et al., 2016; Santoro et al., 2017;

Original: Kipf and Welling, 2017).
Revised: Kipf and Welling, 2017).

Original: Zambaldi et al.
Revised: Zambaldi et al.

Original: (2018) was one of the first to leverage the self-attention approach in the context of deep reinforcement learning for toy and game environments from a constant number of features extracted from a CNN.
Revised: (2018) was one of the first to leverage the self-attention approach in the context of deep reinforcement learning for toy and game environments from a constant number of features extracted from a CNN.

Original: The ATTN baseline in our experiments directly resembles their object-processing module, but we will also evaluate this approach when the number of input objects changes.
Revised: The ATTN baseline in our experiments directly resembles their object-processing module, but we will also evaluate this approach when the number of input objects changes.

Original: Interestingly, most of the previous work for learning similar manipulation tasks in multi-object settings uses such object-encoding modules based on multi-head dot product attention (MHDPA) layers to process the object-based input representations Li et al.
Revised: Interestingly, most of the previous work for learning similar manipulation tasks in multi-object settings uses such object-encoding modules based on multi-head dot product attention (MHDPA) layers to process the object-based input representations Li et al.

Original: ; Zadaianchuk et al.
Revised: ; Zadaianchuk et al.

Original: ; Wilson and Hermans (2020).
Revised: ; Wilson and Hermans (2020).

Original: In summary, previous approaches suffer from various fundamental limitations regarding our desiderata.
Revised: In summary, previous approaches suffer from various fundamental limitations regarding our desiderata.

Original: First, the compute of MHDPA-based architectures scales quadratically with the number of objects, which is especially limiting for training and inference with N o ≫ 1 .
Revised: First, the compute of MHDPA-based architectures scales quadratically with the number of objects, which is especially limiting for training and inference with N o ≫ 1 .

Original: Second, to achieve extrapolation for different object numbers, the training domain must cover a variable and typically large number of objects.
Revised: Second, to achieve extrapolation for different object numbers, the training domain must cover a variable and typically large number of objects.

Original: Third, the performance at test time typically decreases for a new number of objects, and fine-tuning is required to recover the original performance.
Revised: Third, the performance at test time typically decreases for a new number of objects, and fine-tuning is required to recover the original performance.

Original: Related work on relation networks.
Revised: Related work on relation networks.

Original: Our method is motivated by means of an alternative perspective on this problem via relation networks (Raposo et al., 2017; Santoro et al., 2017).
Revised: Our method is motivated by means of an alternative perspective on this problem via relation networks (Raposo et al., 2017; Santoro et al., 2017).

Original: Relation networks process input sets by computing pairwise relations for every object-object pair, followed by a permutation invariant aggregation function, such as the sum, mean, or max.
Revised: Relation networks process input sets by computing pairwise relations for every object-object pair, followed by a permutation invariant aggregation function, such as the sum, mean, or max.

Original: Relation networks generally scale quadratically in the number of input objects (Santoro et al., 2017).
Revised: Relation networks generally scale quadratically in the number of input objects (Santoro et al., 2017).

Original: They are predominantly used for various non-RL applications such as visual and text-based question answering, few-shot image comparison, or object detection (Santoro et al., 2017; Sung et al., 2018;
Revised: They are predominantly used for various non-RL applications such as visual and text-based question answering, few-shot image comparison, or object detection (Santoro et al., 2017; Sung et al., 2018;

Original: Hu et al., 2018).
Revised: Hu et al., 2018).

Original: Importantly, we are not aware of RL applications approaching this problem utilizing relation networks as relational inductive biases.
Revised: Importantly, we are not aware of RL applications approaching this problem utilizing relation networks as relational inductive biases.

Original: In addition, standard relation networks typically operate on a fixed number of input objects and suffer from the same scalability issues as attention.
Revised: In addition, standard relation networks typically operate on a fixed number of input objects and suffer from the same scalability issues as attention.

Original:  A GENTS WITH O BJECT R EASONING M ODULES
Revised:  A GENTS WITH O BJECT R EASONING M ODULES

Original: This section describes the linear relation network (LRN) that we propose as an alternative object reasoning module for learning compositional multi-object manipulation skills.
Revised: This section describes the linear relation network (LRN) that we propose as an alternative object reasoning module for learning compositional multi-object manipulation skills.

Original: See Figure 2 for a summary of the overall architecture, and the agent model is described in Appendix A
Revised: See Figure 2 for a summary of the overall architecture, and the agent model is described in Appendix A.4.

Original: .4. Linear Relation Network Reasoning Module.
Revised: Linear Relation Network Reasoning Module.

Original: Similar to the original Relation Network from Santoro et al.
Revised: Similar to the original Relation Network from Santoro et al.

Original: (2017) used for visual question answering tasks, our module builds upon the idea of reasoning about the future actions from learned relations between the given objects.
Revised: (2017) used for visual question answering tasks, our module builds upon the idea of reasoning about the future actions from learned relations between the given objects.

Original: It seems reasonable that in the context of goal-conditioned RL, the next optimal action is primarily governed by the relations of the available objects with the particular goal of interest.
Revised: It seems reasonable that in the context of goal-conditioned RL, the next optimal action is primarily governed by the relations of the available objects with the particular goal of interest.

Original: In contrast, most other relations might be only important in a sparse and context-dependent way.
Revised: In contrast, most other relations might be only important in a sparse and context-dependent way.

Original: We, therefore, argue to focus on these relations, and we propose the following adjusted composite representation function
Revised: We, therefore, argue to focus on these relations, and we propose the following adjusted composite representation function

Original: [Equation]
Revised: [Equation]

Original: where the o i ∈
Revised: where the o i ∈

Original: R
Revised: R

Original: d are common-size embeddings of the K input objects.
Revised: d are common-size embeddings of the K input objects.

Original: The objects are the scene’s moveable elements in our setting, which means cubes and the three robotic finger states.
Revised: The objects are the scene’s moveable elements in our setting, which means cubes and the three robotic finger states.

Original: The goal information of the task is being embedded into a separate object o g ∈ R d .
Revised: The goal information of the task is being embedded into a separate object o g ∈ R d .

Original: First, the shared learnable function g θ computes the K embeddings r i ∈ R d for each object-goal relation.
Revised: First, the shared learnable function g θ computes the K embeddings r i ∈ R d for each object-goal relation.

Original: For our experiment g θ is a 4-layer MLP.
Revised: For our experiment g θ is a 4-layer MLP.

Original: Next, a dimension-wise sum operation over the K relations is applied, which serves as a permutation-invariant aggregation function.
Revised: Next, a dimension-wise sum operation over the K relations is applied, which serves as a permutation-invariant aggregation function.

Original: The final representation z ∈ R d z is obtained after processing this intermediate representation by another learnable function f ϕ followed by a layer normalization operation (Ba et al., 2016).
Revised: The final representation z ∈ R d z is obtained after processing this intermediate representation by another learnable function f ϕ followed by a layer normalization operation (Ba et al., 2016).

Original: The function f ϕ can serve the purpose of learning further necessary relations from the K object-goal relations.
Revised: The function f ϕ can serve the purpose of learning further necessary relations from the K object-goal relations.

Original: In our experiments, we will use a simple linear transformation for f ϕ .
Revised: In our experiments, we will use a simple linear transformation for f ϕ .

Original: The layer normalization ensures that the final representation z remains in distribution for a changing number of input objects.
Revised: The layer normalization ensures that the final representation z remains in distribution for a changing number of input objects.

Original: The overall complexity of this goal-conditioned linear relation network reasoning module is linear in K, i.e. O ( K ) .
Revised: The overall complexity of this goal-conditioned linear relation network reasoning module is linear in K, i.e. O ( K ) .

Original: To assess the impact of inferring the K object-goal relations only, we will also implement an equivalent variant that considers all possible K 2 relations, similar to the naive implementation of relation networks (’RN’).
Revised: To assess the impact of inferring the K object-goal relations only, we will also implement an equivalent variant that considers all possible K 2 relations, similar to the naive implementation of relation networks (’RN’).

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: We now aim to evaluate the previously described and proposed architectures for learning challenging manipulation tasks in fixed multi-object settings and their ability to extrapolate this skill zero-shot when the number of objects changes.
Revised: We now aim to evaluate the previously described and proposed architectures for learning challenging manipulation tasks in fixed multi-object settings and their ability to extrapolate this skill zero-shot when the number of objects changes.

Original: For an outline of the experimental setup refer to Appendix A.7.
Revised: For an outline of the experimental setup refer to Appendix A.7.

Original: Here, we will discuss the learning behaviour in terms of sample efficiency and performance of our proposed modules against the common attention-based approach from previous works in Section 4.1.
Revised: Here, we will discuss the learning behaviour in terms of sample efficiency and performance of our proposed modules against the common attention-based approach from previous works in Section 4.1.

Original: We close this section by demonstrating the models’ zero-shot extrapolation skills in Section 4.2.
Revised: We close this section by demonstrating the models’ zero-shot extrapolation skills in Section 4.2.

Original: 4. L EARNING TO MANIPULATE OBJECTS FROM FIXED MULTI - OBJECT TRAIN DOMAINS .
Revised: 4. L EARNING TO MANIPULATE OBJECTS FROM FIXED MULTI - OBJECT TRAIN DOMAINS .

Original: Quantitative results for learning to manipulate objects from the 3 cubes train domain are summarized in Figure 3 for PushingCube in the top panel and LiftingCube in the lower panel.
Revised: Quantitative results for learning to manipulate objects from the 3 cubes train domain are summarized in Figure 3 for PushingCube in the top panel and LiftingCube in the lower panel.

Original: We refer to the project-site (sites.google.com/view/compositional-rrl/) for additional videos.
Revised: We refer to the project-site (sites.google.com/view/compositional-rrl/) for additional videos.

Original: Across tasks, we can see that our proposed backbone performs very favourably over the attention-based ATTN baseline.
Revised: Across tasks, we can see that our proposed backbone performs very favourably over the attention-based ATTN baseline.

Original: Specifically, both relation network variants – RN and LRN – reliably succeed in learning to solve the tasks almost perfectly with very high scores.
Revised: Specifically, both relation network variants – RN and LRN – reliably succeed in learning to solve the tasks almost perfectly with very high scores.

Original: On the other hand, ATTN exhibits much higher variance across different seeds and takes longer to train.
Revised: On the other hand, ATTN exhibits much higher variance across different seeds and takes longer to train.

Original: ATTN is clearly inferior in terms of final performances.
Revised: ATTN is clearly inferior in terms of final performances.

Original: Interestingly, RN and LRN reveal very similar robust learning behaviour in terms of training time steps.
Revised: Interestingly, RN and LRN reveal very similar robust learning behaviour in terms of training time steps.

Original: Summary.
Revised: Summary.

Original: The relation network reasoning modules and our LRN model represent a promising alternative over attention-based architectures commonly used in prior works for learning to manipulate an object on a multi-object train domain.
Revised: The relation network reasoning modules and our LRN model represent a promising alternative over attention-based architectures commonly used in prior works for learning to manipulate an object on a multi-object train domain.

Original: Our module outperforms these approaches in final success, sample efficiency, and variance across different seeds
Revised: Our module outperforms these approaches in final success, sample efficiency, and variance across different seeds

Original: 4. Z ERO - SHOT GENERALIZATION TO N- OBJECT ENVIRONMENTS .
Revised: 4. Z ERO - SHOT GENERALIZATION TO N- OBJECT ENVIRONMENTS .

Original: We now evaluate the trained models’ generalization ability when the number of objects is different from training.
Revised: We now evaluate the trained models’ generalization ability when the number of objects is different from training.

Original: Our main result to assess this question is summarized in Figure 3 (right).
Revised: Our main result to assess this question is summarized in Figure 3 (right).

Original: Remember that all agents have been trained for pushing and lifting an object in 3 cubes environments, synonymous with 2 distractors.
Revised: Remember that all agents have been trained for pushing and lifting an object in 3 cubes environments, synonymous with 2 distractors.

Original: Testing these agents’ generalization in N-object environments with N o ̸ = 3 reveals a clear separation between the relation network modules and the attention-based architecture.
Revised: Testing these agents’ generalization in N-object environments with N o ̸ = 3 reveals a clear separation between the relation network modules and the attention-based architecture.

Original: Our
Revised: Our

Original: RN and LRN modules can maintain strong train-time performance for any tested number of objects.
Revised: RN and LRN modules can maintain strong train-time performance for any tested number of objects.

Original: On the other hand, agents with the ATTN module cannot generalize to a changing number of objects and experience a strong and consistent fading ability to solve the task the more the number of objects deviates from the training environment.
Revised: On the other hand, agents with the ATTN module cannot generalize to a changing number of objects and experience a strong and consistent fading ability to solve the task the more the number of objects deviates from the training environment.

Original: ATTN-based agents fail to generalize under our test tasks completely when adding more than 2 additional objects.
Revised: ATTN-based agents fail to generalize under our test tasks completely when adding more than 2 additional objects.

Original: This result contrasts with the RN and
Revised: This result contrasts with the RN and

Original: LRN module-based agents, who can still reliably solve these tasks.
Revised: LRN module-based agents, who can still reliably solve these tasks.

Original: Because we see no relevant deteriorating trend, we suspect that this generalization capability could potentially even hold beyond the 9-distractor environments tested here.
Revised: Because we see no relevant deteriorating trend, we suspect that this generalization capability could potentially even hold beyond the 9-distractor environments tested here.

Original: Interestingly, the same trends can be seen for environments with less than 3 cubes.
Revised: Interestingly, the same trends can be seen for environments with less than 3 cubes.

Original: Even though it seems intuitive that an agent should be able to solve pushing or lifting an object when no further distractors are in the scene, the ATTN module seems to overfit on the specific 2-distractor input train distribution strongly and cannot solve these tasks without exactly two distractors.
Revised: Even though it seems intuitive that an agent should be able to solve pushing or lifting an object when no further distractors are in the scene, the ATTN module seems to overfit on the specific 2-distractor input train distribution strongly and cannot solve these tasks without exactly two distractors.

Original: On the other hand, RN and LRN module-based agents can robustly maintain their respective train-time performances, presumably due to their relation-based processing of the input objects.
Revised: On the other hand, RN and LRN module-based agents can robustly maintain their respective train-time performances, presumably due to their relation-based processing of the input objects.

Original: Summary.
Revised: Summary.

Original: Agents based on the relation network modules show a strong extrapolation ability of their learned manipulation skills to environments with many more objects.
Revised: Agents based on the relation network modules show a strong extrapolation ability of their learned manipulation skills to environments with many more objects.

Original: In contrast, we do not observe such zero-shot generalization with commonly employed attention-based architectures, which experience a substantial decrease in performance when deployed out-of-distribution.
Revised: In contrast, we do not observe such zero-shot generalization with commonly employed attention-based architectures, which experience a substantial decrease in performance when deployed out-of-distribution.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: In this work, we presented a new object encoder module based on relation networks, enabling reinforcement learning agents to learn challenging compositional manipulation tasks in multi-object robotics environments.
Revised: In this work we presented a new object encoder module based on relation networks, enabling reinforcement learning agents to learn challenging compositional manipulation tasks in multi-object robotics environments.

Original: We extensively studied a newly proposed linear relation network module on the generic task of transporting one particular cube out of N cubes to any specified goal position.
Revised: We extensively studied a newly proposed linear relation network module on the generic task of transporting one particular cube out of N cubes to any specified goal position.

Original: While previous approaches mostly build upon quadratically scaling attention-based architectures to process the set of input objects, our approach scales only linear with the number of objects.
Revised: While previous approaches mostly build upon quadratically scaling attention-based architectures to process the set of input objects, our approach scales only linear with the number of objects.

Original: At the same time, our method outperforms these previous approaches at train time.
Revised: At the same time, our method outperforms these previous approaches at train time.

Original: Most importantly, we can show that our agents can extrapolate and generalize zero-shot without any drop in performance when the number of objects changes which is a major limitation of most previous approaches. 
Revised: Most importantly, we can show that our agents can extrapolate and generalize zero-shot without any drop in performance when the number of objects changes which is a major limitation of most previous approaches. A CKNOWLEDGEMENTS

Revised: We thank the International Max Planck Research School for Intelligent Systems (IMPRS-IS) for supporting FT

Original: A I MPLEMENTATION DETAILS A. R EWARD D ESIGN
Revised: A I MPLEMENTATION DETAILS A. R EWARD D ESIGN

Original: We used the following same reward for both tasks:
Revised: We used the following same reward for both tasks:

Original: [Equation]
Revised: [Equation]

Original: where t denotes the time step, ρ t ∈ [0 , 1] is the fractional overlap with the goal cube at time t , e t ∈ R 3 is the end-effector position, o t ∈ R 3 the cube position of the object that is supposed to be used for the task, g t ∈ R 3 the goal position, and d ( · , · ) denotes the Euclidean distance.
Revised: where t denotes the time step, ρ t ∈ [0 , 1] is the fractional overlap with the goal cube at time t , e t ∈ R 3 is the end-effector position, o t ∈ R 3 the cube position of the object that is supposed to be used for the task, g t ∈ R 3 the goal position, and d ( · , · ) denotes the Euclidean distance.

Original: Across all experiments we set α =
Revised: Across all experiments we set α =

Original: [100 , 250 , 10] and β = 0 .
Revised: [100 , 250 , 10] and β = 0 .

Original: A. T RAINING SCHEME FOR AGENTS USING PPO
Revised: A. T RAINING SCHEME FOR AGENTS USING PPO

Original: We use the distributed PPO implementation from rlpyt (Stooke and Abbeel, 2019) with discount factor 0.98, entropy loss coefficient 0.01, learning rate 0.00025, value loss coefficient 0.5, gradient clipping norm 0.5, 40 minibatches, gae lambda 0.95, clip ratio 0.5 and 4 epochs.
Revised: We use the distributed PPO implementation from rlpyt (Stooke and Abbeel, 2019) with discount factor 0.98, entropy loss coefficient 0.01, learning rate 0.00025, value loss coefficient 0.5, gradient clipping norm 0.5, 40 minibatches, gae lambda 0.95, clip ratio 0.5 and 4 epochs.

Original: We applied a linear learning rate scheduling and Adam as optimizer (Kingma and Ba, 2015).
Revised: We applied a linear learning rate scheduling and Adam as optimizer (Kingma and Ba, 2015).

Original: A. S IMULATOR INPUT FEATURES
Revised: A. S IMULATOR INPUT FEATURES

Original: The ground truth representations for each object type are described in Table 1.
Revised: The ground truth representations for each object type are described in Table 1.

Original: Cubes are described using their position and orientation and each corresponding velocity.
Revised: Cubes are described using their position and orientation and each corresponding velocity.

Original: In addition, every cube is marked with a unique identifier which makes the set of cubes heterogeneous.
Revised: In addition, every cube is marked with a unique identifier which makes the set of cubes heterogeneous.

Original: Similarly, the goal is specified via its cartesian position, orientation, and an identifier variable used to specify which cube has to be moved.
Revised: Similarly, the goal is specified via its cartesian position, orientation, and an identifier variable used to specify which cube has to be moved.

Original: Both cubes and the goal include features describing their shape and size; however, we always use a cube shape and fixed size in this work.
Revised: Both cubes and the goal include features describing their shape and size; however, we always use a cube shape and fixed size in this work.

Original: Each finger of the robot hasdegrees of freedom, and it is described by a 9-dimensional vector composed by its joint positions, joint velocities, and end-effector position.
Revised: Each finger of the robot hasdegrees of freedom, and it is described by a 9-dimensional vector composed by its joint positions, joint velocities, and end-effector position.

Original: A. A GENT MODEL
Revised: A. A GENT MODEL

Original: Our agents will build upon the following standard architecture scaffold.
Revised: Our agents will build upon the following standard architecture scaffold.

Original: In a first step, ground-truth simulator state features represented by objects are encoded into fixed-size object tokens o i ∈ R d via object-type specific small MLPs.
Revised: In a first step, ground-truth simulator state features represented by objects are encoded into fixed-size object tokens o i ∈ R d via object-type specific small MLPs.

Original: In our case, these are three robotic states for each finger, goal objects, and physical objects for the cube.
Revised: In our case, these are three robotic states for each finger, goal objects, and physical objects for the cube.

Original: Next, an object-reasoning encoder module processes this token set and outputs a fixed-size representation z (highlighted in pink in Figure 2).
Revised: Next, an object-reasoning encoder module processes this token set and outputs a fixed-size representation z (highlighted in pink in Figure 2).

Original: Finally, this representation forms the input of any downstream predictor required to learn the optimal policy.
Revised: Finally, this representation forms the input of any downstream predictor required to learn the optimal policy.

Original: Here, we will assume a small MLP to predict the action (policy); however, we can likewise use the representation to predict the value (value network) or Q-Value (when combined with an action) when using Q-Learning or actor-critic approaches.
Revised: Here, we will assume a small MLP to predict the action (policy); however, we can likewise use the representation to predict the value (value network) or Q-Value (when combined with an action) when using Q-Learning or actor-critic approaches.

Original: The primary focus of this work is on proposing a simple and flexible object-encoder module that overcomes the limitations of previous approaches.
Revised: The primary focus of this work is on proposing a simple and flexible object-encoder module that overcomes the limitations of previous approaches.

Original: A. A TTENTION MODULE - B ASELINE
Revised: A. A TTENTION MODULE - B ASELINE

Original: The baseline we use in this work exploits a MHDPA module (Vaswani et al., 2017) operating on the set o ∈ R d × ( K +1) composed of common-size embedding of the K input objects o i ∈ R d , and the goal o g ∈ R d with[Equation]
Revised: The baseline we use in this work exploits a MHDPA module (Vaswani et al., 2017) operating on the set o ∈ R d × ( K +1) composed of common-size embedding of the K input objects o i ∈ R d , and the goal o g ∈ R d with

Original: [Equation]
Revised: [Equation]

Original: where o attn is the set of K + 1 embeddings scaled from the MHDPA module.
Revised: where o attn is the set of K + 1 embeddings scaled from the MHDPA module.

Original: In this module the goal-conditioning element is obtained by dealing with the goal as an object.
Revised: In this module the goal-conditioning element is obtained by dealing with the goal as an object.

Original: This set is then mappedUnder to the fixed-size embedding by an aggregation function and normalized to avoid changes in scale when the number of objects differ between training to testing.
Revised: This set is then mapped to the fixed-size embedding by an aggregation function and normalized to avoid changes in scale when the number of objects differ between training to testing.

Original: [Equation]
Revised: [Equation]

Original: This reasoning module resembles the one from Zambaldi et al.
Revised: This reasoning module resembles the one from Zambaldi et al.

Original: and it is used to introduce relational inductive biases into the learning process.
Revised: and it is used to introduce relational inductive biases into the learning process.

Original: The same self-attention layer is also used by
Revised: The same self-attention layer is also used by

Original: Zadaianchuk et al.
Revised: Zadaianchuk et al.

Original: (2021) to build a goal-conditioned policy that is able to handle set-based input representations.
Revised: (2021) to build a goal-conditioned policy that is able to handle set-based input representations.

Original: Similarly, it also describes approaches that rely on Graph Attention Networks (GAT) (Veliˇckovi´c et al., 2018) to build policies while using the common assumption in multi-object manipulation of a fully connected graph (Li et al., 2020; Wilson and Hermans, 2020).
Revised: Similarly, it also describes approaches that rely on Graph Attention Networks (GAT) (Veliˇckovi´c et al., 2018) to build policies while using the common assumption in multi-object manipulation of a fully connected graph (Li et al., 2020; Wilson and Hermans, 2020).

Original: Self-attention when computing the attention matrix takes into consideration each object-object pair, which makes it scale quadratically with the number of objects in the input set, as highlighted in Figure 2 (A).
Revised: Self-attention when computing the attention matrix takes into consideration each object-object pair, which makes it scale quadratically with the number of objects in the input set, as highlighted in Figure 2 (A).

Original: A. R ELATION N ETWORK M ODULE
Revised: A. R ELATION N ETWORK M ODULE

Original: To have a further element of comparison, we introduced Relation Networks as relational reasoning module in the context of reinforcement learning.
Revised: To have a further element of comparison, we introduced Relation Networks as relational reasoning module in the context of reinforcement learning.

Original: [Equation]
Revised: [Equation]

Original: Relation Networks have been proposed in Santoro et al.
Revised: Relation Networks have been proposed in Santoro et al.

Original: and they compute ( K + 1) 2 embeddings r i ∈ R d , referred to as relations, using the shared function g θ .
Revised: and they compute ( K + 1) 2 embeddings r i ∈ R d , referred to as relations, using the shared function g θ .

Original: Relations are aggregated using a sum and then normalized to obtain a fixed-size embedding used as input for the MLP policy head.
Revised: Relations are aggregated using a sum and then normalized to obtain a fixed-size embedding used as input for the MLP policy head.

Original: Similarly to self-attention, Relation Networks need to consider every object-object interaction, making them scale quadratically in the number of objects.
Revised: Similarly to self-attention, Relation Networks need to consider every object-object interaction, making them scale quadratically in the number of objects.

Original: A key aspect of this architecture is that g θ always takes as input a pair of objects, independent of the total number of objects in the input set.
Revised: A key aspect of this architecture is that g θ always takes as input a pair of objects, independent of the total number of objects in the input set.

Original: This property makes the architecture robust to changes in the number of objects as long as the objects’ statistics do not change.
Revised: This property makes the architecture robust to changes in the number of objects as long as the objects’ statistics do not change.

Original: A. E XPERIMENTAL SETUP .
Revised: A. E XPERIMENTAL SETUP .

Original: Both tasks are derived from the simulated robotic manipulation setup and tasks proposed in CausalWorld (Ahmed* et al., 2021) which builds upon the robotic TriFinger design (W¨uthrich et al., 2020;
Revised: Both tasks are derived from the simulated robotic manipulation setup and tasks proposed in CausalWorld (Ahmed* et al., 2021) which builds upon the robotic TriFinger design (W¨uthrich et al., 2020;

Original: Bauer et al., 2021).
Revised: Bauer et al., 2021).

Original: Compared to previous work on this setup, we attempt to solve tasks in scenes that contain more than 1 cube (Dittadi et al., 2021;
Revised: Compared to previous work on this setup, we attempt to solve tasks in scenes that contain more than 1 cube (Dittadi et al., 2021;

Original: Ahmed* et al., 2021; Allshire et al., 2021).
Revised: Ahmed* et al., 2021; Allshire et al., 2021).

Original: Specifically, the goal is to move any of the 3 available cubes to a freely specifiable target position.
Revised: Specifically, the goal is to move any of the 3 available cubes to a freely specifiable target position.

Original: We cover a simpler sub-task where the target position is constrained to the floor and the more general task where the target position is also allowed to be above the ground.
Revised: We cover a simpler sub-task where the target position is constrained to the floor and the more general task where the target position is also allowed to be above the ground.

Original: Such a lifting task involves learning to hold the cube stable in the air at a variable height.
Revised: Such a lifting task involves learning to hold the cube stable in the air at a variable height.

Original: We refer to these tasks as PushingCube and LiftingCube, respectively.
Revised: We refer to these tasks as PushingCube and LiftingCube, respectively.

Original: As discussed in Section 2, we specify the object to be rearranged through a matching object identifier in one of the goal objects’ features.
Revised: As discussed in Section 2, we specify the object to be rearranged through a matching object identifier in one of the goal objects’ features.

Original: For both tasks, we use a generic reward function designed to work for any single-goal task in this environment.
Revised: For both tasks, we use a generic reward function designed to work for any single-goal task in this environment.

Original: The reward is composed of three terms: (1) a reward term that directly measures the success at each time step, which is defined as the volumetric overlap of the cube with the goal cube; (2) a second reward term that encourages moving the target object closer to the goal; (3) a third curiosity-inspired reward term for efficient exploration that gives a positive reward signal for moving the target object when being further away from the goal.
Revised: The reward is composed of three terms: (1) a reward term that directly measures the success at each time step, which is defined as the volumetric overlap of the cube with the goal cube; (2) a second reward term that encourages moving the target object closer to the goal; (3) a third curiosity-inspired reward term for efficient exploration that gives a positive reward signal for moving the target object when being further away from the goal.

Original: See Appendix C for an analysis regarding the role of each term.
Revised: See Appendix C for an analysis regarding the role of each term.

Original: Across all tasks, we will report the factional success as the universal evaluation metric at the last timestep of an episode.
Revised: Across all tasks, we will report the factional success as the universal evaluation metric at the last timestep of an episode.

Original: Tasks can be visually considered solved with a score around 80%, and we can therefore ensure an objective and interpretable performance measure (Ahmed* et al., 2021; Dittadi et al., 2021).
Revised: Tasks can be visually considered solved with a score around 80%, and we can therefore ensure an objective and interpretable performance measure (Ahmed* et al., 2021; Dittadi et al., 2021).

Original: All models are being trained under the described reward with the help of PPO (Schulman et al., 2017) on distributed workers with 5 random seeds each.
Revised: All models are being trained under the described reward with the help of PPO (Schulman et al., 2017) on distributed workers with 5 random seeds each.

Original: Initial object poses and the goal pose are randomly sampled at the beginning of each episode, and the episodic task needs to be solved within 30 simulation seconds or 2500 frames (10 seconds per object).
Revised: Initial object poses and the goal pose are randomly sampled at the beginning of each episode, and the episodic task needs to be solved within 30 simulation seconds or 2500 frames (10 seconds per object).

Original: We compare the instantiation of our proposed RN and
Revised: We compare the instantiation of our proposed RN and

Original: LRN network architectures against a representative attention-based baseline derived from (Zambaldi et al., 2018), referred to as ATTN.
Revised: LRN network architectures against a representative attention-based baseline derived from (Zambaldi et al., 2018), referred to as ATTN.

Original: We ensured a fair comparison of this baseline through extensive hyperparameter optimization (see Appendix D.3).
Revised: We ensured a fair comparison of this baseline through extensive hyperparameter optimization (see Appendix D.3).

Original: We refer to the supplement for a comprehensive account of further implementation details.
Revised: We refer to the supplement for a comprehensive account of further implementation details.

Original: B D ISCUSSION
Revised: B D ISCUSSION

Original: Role of changing the train domains.
Revised: Role of changing the train domains.

Original: In the previous section, we always trained under exactlydistractors.
Revised: In the previous section, we always trained under exactlydistractors.

Original: In particular, we believed that a 2-distractor environment also includes edge cases that the agent must face during training to solve the same task with a large number of distractors.
Revised: In particular, we believed that a 2-distractor environment also includes edge cases that the agent must face during training to solve the same task with a large number of distractors.

Original: An example is when the selected cube is stuck between 2 distractor cubes blocking its grasping.
Revised: An example is when the selected cube is stuck between 2 distractor cubes blocking its grasping.

Original: We found this to be a frequent scenario when the arena gets crowded.
Revised: We found this to be a frequent scenario when the arena gets crowded.

Original: To validate this hypothesis, we repeated our experiments in environments with 3 distractors and additionally covered the limiting case of only one and no distractor at train time at all.
Revised: To validate this hypothesis, we repeated our experiments in environments with 3 distractors and additionally covered the limiting case of only one and no distractor at train time at all.

Original: Our results are presented in Figure 5, which shows how the zero-shot generalization under each architecture is affected by its training distribution.
Revised: Our results are presented in Figure 5, which shows how the zero-shot generalization under each architecture is affected by its training distribution.

Original: Concerning the ATTN module, we observe the effect of the train distribution on generalization capabilities to be generally limited.
Revised: Concerning the ATTN module, we observe the effect of the train distribution on generalization capabilities to be generally limited.

Original: Performances always peak in the training environment and monotonically decrease for N testo ̸
Revised: Performances always peak in the training environment and monotonically decrease for N testo ̸

Original: = N traino .
Revised: = N traino .

Original: Whereas, a steady improvement can be observed for the RN module when changing the training environment up until 2 distractors.
Revised: Whereas, a steady improvement can be observed for the RN module when changing the training environment up until 2 distractors.

Original: Note that the RN already achieves some level of generalization when training with a single distractor.
Revised: Note that the RN already achieves some level of generalization when training with a single distractor.

Original: It appears a phase change to robust performances happens when 2 distractors are used for training which yields strong generalization for any test environment.
Revised: It appears a phase change to robust performances happens when 2 distractors are used for training which yields strong generalization for any test environment.

Original: Finally, we can observe that our LRN module already achieves a strong generalization when trained with one distractor only, which outperforms both previous architectures.
Revised: Finally, we can observe that our LRN module already achieves a strong generalization when trained with one distractor only, which outperforms both previous architectures.

Original: We therefore conclude that our assumption on the minimum number of distractors is not necessary for our LRN module.
Revised: We therefore conclude that our assumption on the minimum number of distractors is not necessary for our LRN module.

Original: Assessing architectural choices of our LRN module.
Revised: Assessing architectural choices of our LRN module.

Original: We also tested the impact of various design choices for our proposed linear relation network module.
Revised: We also tested the impact of various design choices for our proposed linear relation network module.

Original: Results are summarized in
Revised: Results are summarized in

Original: Figure 6 (left).
Revised: Figure 6 (left).

Original: Specifically, we focus on 3 questions: First, do we need relations at all?
Revised: Specifically, we focus on 3 questions: First, do we need relations at all?

Original: To answer, we repeated the LRN experiments but replaced the relations with equivalent non-relational single-object feedforward MLPs, i.e. g θ ( o i ) .
Revised: To answer, we repeated the LRN experiments but replaced the relations with equivalent non-relational single-object feedforward MLPs, i.e. g θ ( o i ) .

Original: We found that the module cannot even learn the task on the training domain without relations.
Revised: We found that the module cannot even learn the task on the training domain without relations.

Original: Second, we want to understand if we could use a simple average instead of the LayerNorm after summing over all relations?
Revised: Second, we want to understand if we could use a simple average instead of the LayerNorm after summing over all relations?

Original: Our experiments show that even though we can learn the task on the train environments, the agents reveal deteriorating generalization the more objects we add to the task.
Revised: Our experiments show that even though we can learn the task on the train environments, the agents reveal deteriorating generalization the more objects we add to the task.

Original: We hypothesize this to be due to the changing variance of z which is kept constant when doing a layer normalization but not an average.
Revised: We hypothesize this to be due to the changing variance of z which is kept constant when doing a layer normalization but not an average.

Original: Lastly, we analyzed the role of the representations’ dimension, for which we used d z = 25 .
Revised: Lastly, we analyzed the role of the representations’ dimension, for which we used d z = 25 .

Original: When repeating the experiments with d z = 10 , the zero-shot generalization slowly decreases to about 60 % success with 9 objects.
Revised: When repeating the experiments with d z = 10 , the zero-shot generalization slowly decreases to about 60 % success with 9 objects.

Original: On the other hand, when allowing for d z = 50 or even d z = 100 , the agent achieves slightly elevated train time performance with a slight drop of about 5 % and a more pronounced drop of about% respectively when testing with up to 9 distractors.
Revised: On the other hand, when allowing for d z = 50 or even d z = 100 , the agent achieves slightly elevated train time performance with a slight drop of about 5 % and a more pronounced drop of about% respectively when testing with up to 9 distractors.

Original: This result suggests that the representation
Revised: This result suggests that the representation

Original: Role of the reward function.
Revised: Role of the reward function.

Original: We finally wanted to understand the impact of each term on the reward function we proposed.
Revised: We finally wanted to understand the impact of each term on the reward function we proposed.

Original: Figure 6 (right) shows the generalization performance of corresponding agents trained with each variation of the reward function.
Revised: Figure 6 (right) shows the generalization performance of corresponding agents trained with each variation of the reward function.

Original: We can distinguish reward terms based on their effect: sparse, and distance terms can be used to define the task; the curiosity term guides the exploration.
Revised: We can distinguish reward terms based on their effect: sparse, and distance terms can be used to define the task; the curiosity term guides the exploration.

Original: The first notable observation is that removing the distance term does not affect the overall performance.
Revised: The first notable observation is that removing the distance term does not affect the overall performance.

Original: Therefore, we can conclude that the sparse term is sufficient to solve the task successfully.
Revised: Therefore, we can conclude that the sparse term is sufficient to solve the task successfully.

Original: On the other hand, removing the sparse term while keeping the distance term stops the agent from learning to solve the task, underlying its importance.
Revised: On the other hand, removing the sparse term while keeping the distance term stops the agent from learning to solve the task, underlying its importance.

Original: C D ISCUSSION
Revised: C D ISCUSSION

Original: Role of changing the train domains.
Revised: Role of changing the train domains.

Original: In the previous section, we always trained under exactlydistractors.
Revised: In the previous section, we always trained under exactlydistractors.

Original: In particular, we believed that a 2-distractor environment also includes edge cases that the agent must face during training to solve the same task with a large number of distractors.
Revised: In particular, we believed that a 2-distractor environment also includes edge cases that the agent must face during training to solve the same task with a large number of distractors.

Original: An example is when the selected cube is stuck between 2 distractor cubes blocking its grasping.
Revised: An example is when the selected cube is stuck between 2 distractor cubes blocking its grasping.

Original: We found this to be a frequent scenario when the arena gets crowded.
Revised: We found this to be a frequent scenario when the arena gets crowded.

Original: To validate this hypothesis, we repeated ourUnder review at the ICLR workshop on and experiments in environments with 3 distractors and additionally covered the limiting case of only one and no distractor at train time at all.
Revised: To validate this hypothesis, we repeated our experiments in environments with 3 distractors and additionally covered the limiting case of only one and no distractor at train time at all.

Original: Our results are presented in Figure 5, which shows how the zero-shot generalization under each architecture is affected by its training distribution.
Revised: Our results are presented in Figure 5, which shows how the zero-shot generalization under each architecture is affected by its training distribution.

Original: Concerning the ATTN module, we observe the effect of the train distribution on generalization capabilities to be generally limited.
Revised: Concerning the ATTN module, we observe the effect of the train distribution on generalization capabilities to be generally limited.

Original: Performances always peak in the training environment and monotonically decrease for N testo ̸
Revised: Performances always peak in the training environment and monotonically decrease for N testo ̸

Original: = N traino .
Revised: = N traino .

Original: Whereas, a steady improvement can be observed for the RN module when changing the training environment up until 2 distractors.
Revised: Whereas, a steady improvement can be observed for the RN module when changing the training environment up until 2 distractors.

Original: Note that the RN already achieves some level of generalization when training with a single distractor.
Revised: Note that the RN already achieves some level of generalization when training with a single distractor.

Original: It appears a phase change to robust performances happens when 2 distractors are used for training which yields strong generalization for any test environment.
Revised: It appears a phase change to robust performances happens when 2 distractors are used for training which yields strong generalization for any test environment.

Original: Finally, we can observe that our LRN module already achieves a strong generalization when trained with one distractor only, which outperforms both previous architectures.
Revised: Finally, we can observe that our LRN module already achieves a strong generalization when trained with one distractor only, which outperforms both previous architectures.

Original: We therefore conclude that our assumption on the minimum number of distractors is not necessary for our LRN module.
Revised: We therefore conclude that our assumption on the minimum number of distractors is not necessary for our LRN module.

Original: Assessing architectural choices of our LRN module.
Revised: Assessing architectural choices of our LRN module.

Original: We also tested the impact of various design choices for our proposed linear relation network module.
Revised: We also tested the impact of various design choices for our proposed linear relation network module.

Original: Results are summarized in
Revised: Results are summarized in

Original: Figure 6 (left).
Revised: Figure 6 (left).

Original: Specifically, we focus on 3 questions: First, do we need relations at all?
Revised: Specifically, we focus on 3 questions: First, do we need relations at all?

Original: To answer, we repeated the LRN experiments but replaced the relations with equivalent non-relational single-object feedforward MLPs, i.e. g θ ( o i ) .
Revised: To answer, we repeated the LRN experiments but replaced the relations with equivalent non-relational single-object feedforward MLPs, i.e. g θ ( o i ) .

Original: We found that the module cannot even learn the task on the training domain without relations.
Revised: We found that the module cannot even learn the task on the training domain without relations.

Original: Second, we want to understand if we could use a simple average instead of the LayerNorm after summing over all relations?
Revised: Second, we want to understand if we could use a simple average instead of the LayerNorm after summing over all relations?

Original: Our experiments show that even though we can learn the task on the train environments, the agents reveal deteriorating generalization the more objects we add to the task.
Revised: Our experiments show that even though we can learn the task on the train environments, the agents reveal deteriorating generalization the more objects we add to the task.

Original: We hypothesize this to be due to the changing variance of z which is kept constant when doing a layer normalization but not an average.
Revised: We hypothesize this to be due to the changing variance of z which is kept constant when doing a layer normalization but not an average.

Original: Lastly, we analyzed the role of the representations’ dimension, for which we used d z = 25 .
Revised: Lastly, we analyzed the role of the representations’ dimension, for which we used d z = 25 .

Original: When repeating the experiments with d z = 10 , the zero-shot generalization slowly decreases to about 60 % success with 9 objects.
Revised: When repeating the experiments with d z = 10 , the zero-shot generalization slowly decreases to about 60 % success with 9 objects.

Original: On the other hand, when allowing for d z = 50 or even d z = 100 , the agent achieves slightly elevated train time performance with a slight drop of about 5 % and a more pronounced drop of about% respectively when testing with up to 9 distractors.
Revised: On the other hand, when allowing for d z = 50 or even d z = 100 , the agent achieves slightly elevated train time performance with a slight drop of about 5 % and a more pronounced drop of about% respectively when testing with up to 9 distractors.

Original: This result suggests that the representation dimension is optimally being adjusted for the right capacity requirements of the task.
Revised: This result suggests that the representation dimension is optimally being adjusted for the right capacity requirements of the task.

Original: Role of the reward function.
Revised: Role of the reward function.

Original: We finally wanted to understand the impact of each term on the reward function we proposed.
Revised: We finally wanted to understand the impact of each term on the reward function we proposed.

Original: Figure 6 (right) shows the generalization performance of corresponding agents trained with each variation of the reward function.
Revised: Figure 6 (right) shows the generalization performance of corresponding agents trained with each variation of the reward function.

Original: We can distinguish reward terms based on their effect: sparse, and distance terms can be used to define the task; the curiosity term guides the exploration.
Revised: We can distinguish reward terms based on their effect: sparse, and distance terms can be used to define the task; the curiosity term guides the exploration.

Original: The first notable observation is that removing the distance term does not affect the overall performance.
Revised: The first notable observation is that removing the distance term does not affect the overall performance.

Original: Therefore, we can conclude that the sparse term is sufficient to solve the task successfully.
Revised: Therefore, we can conclude that the sparse term is sufficient to solve the task successfully.

Original: On the other hand, removing the sparse term while keeping the distance term stops the agent from learning to solve the task, underlying its importance.
Revised: On the other hand, removing the sparse term while keeping the distance term stops the agent from learning to solve the task, underlying its importance.

Original: D A DDITIONAL R ESULTS D. D ISTRIBUTION S HIFT FOR OOD TASKS
Revised: D A DDITIONAL R ESULTS D. D ISTRIBUTION S HIFT FOR OOD TASKS

Original: To better understand the difference between the architectures, we visualized the representations of each reasoning module with the help of lower-dimensional embeddings.
Revised: To better understand the difference between the architectures, we visualized the representations of each reasoning module with the help of lower-dimensional embeddings.

Original: Each model has been tested for 200 episodes on environments with the number of distractors ranging from 0 up to 9.
Revised: Each model has been tested for 200 episodes on environments with the number of distractors ranging from 0 up to 9.

Original: The representations obtained as output have been mapped to a low-dimensional space using t-SNE (van der Maaten and Hinton, 2008).
Revised: The representations obtained as output have been mapped to a low-dimensional space using t-SNE (van der Maaten and Hinton, 2008).

Original: It is important to observe that these points represent the input of the MLP policy head; therefore, to have generalization capabilities, they should have similar statistics across different tasks.
Revised: It is important to observe that these points represent the input of the MLP policy head; therefore, to have generalization capabilities, they should have similar statistics across different tasks.

Original: In Figure 7(a) we can observe the result obtained for the attention-based module when visualized in a 2-dimensional space.
Revised: In Figure 7(a) we can observe the result obtained for the attention-based module when visualized in a 2-dimensional space.

Original: The embedding distribution changes with the task, which causes the MLP policy head to operate out-of-distribution justifying the drop in performances observed in Figure 3.
Revised: The embedding distribution changes with the task, which causes the MLP policy head to operate out-of-distribution justifying the drop in performances observed in Figure 3.

Original: In Figure 7(b) and Figure 7(c), respectively the RN and LRN output embeddings are visualized.
Revised: In Figure 7(b) and Figure 7(c), respectively the RN and LRN output embeddings are visualized.

Original: For both, we can say that the approximated support of the distribution appears to be invariant with respect to the task.
Revised: For both, we can say that the approximated support of the distribution appears to be invariant with respect to the task.

Original: D. R ELATIONS ANALYSIS
Revised: D. R ELATIONS ANALYSIS

Original: Relations Sub-Spaces We decided to visualize 2-dimensional embeddings computed using t-SNE color-coding different relations to understand whether relations with different objects are qualitatively different.
Revised: Relations Sub-Spaces We decided to visualize 2-dimensional embeddings computed using t-SNE color-coding different relations to understand whether relations with different objects are qualitatively different.

Original: We used relations generated from 200 episodes with a number of distractors varying from
Revised: We used relations generated from 200 episodes with a number of distractors varying from

Original: [Figure]
Revised: [Figure]

Original: Under review at the ICLR workshop on and

Original: 3 to 6.
Revised: 3 to 6.

Original: The obtained plot is presented in Figure 8, where cube-goal relations are shown in green, distractors-goal relations are visualized in blue, and fingers relations are colored respectively in black, brown, and red.
Revised: The obtained plot is presented in Figure 8, where cube-goal relations are shown in green, distractors-goal relations are visualized in blue, and fingers relations are colored respectively in black, brown, and red.

Original: Different relations are mapped to different sub-spaces of the representation space, suggesting that invariance to the number of distractors might be obtained via a gating property learned by the architecture to filter out relations embedded into certain sub-spaces when computing the action.
Revised: Different relations are mapped to different sub-spaces of the representation space, suggesting that invariance to the number of distractors might be obtained via a gating property learned by the architecture to filter out relations embedded into certain sub-spaces when computing the action.

Original: Relations over time We also analyzed how the norm of each relation type changes over time while the agent is solving the task.
Revised: Relations over time We also analyzed how the norm of each relation type changes over time while the agent is solving the task.

Original: In Figure 9 we show the result when testing on 5 distractors a LRN-based agent and normalizing each type of relation by the number of relations belonging to that type.
Revised: In Figure 9 we show the result when testing on 5 distractors a LRN-based agent and normalizing each type of relation by the number of relations belonging to that type.

Original: Since we used the sum as the aggregation function, the norm can be associated with each relation’s weight in the representation after the sum.
Revised: Since we used the sum as the aggregation function, the norm can be associated with each relation’s weight in the representation after the sum.

Original: Interestingly, the most important relations are those describing fingers and the target cube, whereas those linked to distractors have a much smaller norm.
Revised: Interestingly, the most important relations are those describing fingers and the target cube, whereas those linked to distractors have a much smaller norm.

Original: This qualitative result suggests that the cube-goal relation might be correlated with the cube-goal distance.
Revised: This qualitative result suggests that the cube-goal relation might be correlated with the cube-goal distance.

Original: Capturing higher-order relations The Linear Relation Network does not compute most of the possible relations explicitly, which is why it scales linearly with the number of objects.
Revised: Capturing higher-order relations The Linear Relation Network does not compute most of the possible relations explicitly, which is why it scales linearly with the number of objects.

Original: A question that naturally arises is what happens when a particular relation that the Linear Relation Network does not compute is strictly necessary to solve the task.
Revised: A question that naturally arises is what happens when a particular relation that the Linear Relation Network does not compute is strictly necessary to solve the task.

Original: An example is provided by the pushing task, where the agent must reason about where the fingers are with respect to the cube to touch it and then grasp it.
Revised: An example is provided by the pushing task, where the agent must reason about where the fingers are with respect to the cube to touch it and then grasp it.

Original: However, since the LRN only computes relations to the goal, the fingers-cube relation is not explicitly available to the agent.
Revised: However, since the LRN only computes relations to the goal, the fingers-cube relation is not explicitly available to the agent.

Original: The only way to manipulate the cube using available relations is to combine the cube-goal and fingers-goal relations using the goal as a reference.
Revised: The only way to manipulate the cube using available relations is to combine the cube-goal and fingers-goal relations using the goal as a reference.

Original: This capability means the agent can extrapolate and use higher-order relations.
Revised: This capability means the agent can extrapolate and use higher-order relations.

Original: The LRN successfully solves the pushing task, which shows its ability to infer essential relations for the downstream task.
Revised: The LRN successfully solves the pushing task, which shows its ability to infer essential relations for the downstream task.

Original: The higher-order reasoning is handled by by the subsequent layers of the architecture after computing the relations, which is the MLP policy head in our case.
Revised: The higher-order reasoning is handled by by the subsequent layers of the architecture after computing the relations, which is the MLP policy head in our case.

Original: D. L IMITATIONS OF THE A TTENTION -B ASED N ETWORK
Revised: D. L IMITATIONS OF THE A TTENTION -B ASED N ETWORK

Original: Given the observed poor generalization results and the very high variance during training for the attention-based baseline, we performed extensive hyperparameter tuning to ensure a fair comparison against our proposed RN and LRN modules.
Revised: Given the observed poor generalization results and the very high variance during training for the attention-based baseline, we performed extensive hyperparameter tuning to ensure a fair comparison against our proposed RN and LRN modules.

Original: We used Bayesian search (Biewald, 2020) across more than 550 hyperparameter configurations.
Revised: We used Bayesian search (Biewald, 2020) across more than 550 hyperparameter configurations.

Original: We found that none of those improves over our set of hyperparameters meaningfully.
Revised: We found that none of those improves over our set of hyperparameters meaningfully.

Original: To find potential bottlenecks in both the RL algorithm (PPO) and the attention modules, we used the hyperparameters search algorithm from Biewald (2020) to check elements of both.
Revised: To find potential bottlenecks in both the RL algorithm (PPO) and the attention modules, we used the hyperparameters search algorithm from Biewald (2020) to check elements of both.

Original: Regarding PPO, we tunned:
Revised: Regarding PPO, we tunned:

Original: • Gradient clipping norm • Discount factor • Entropy loss coefficient • Value loss coefficient • GAE lambda • Learning rate • Ratio clip
Revised: • Gradient clipping norm • Discount factor • Entropy loss coefficient • Value loss coefficient • GAE lambda • Learning rate • Ratio clip

Original: For the archicture of the module we explored different configurations of:
Revised: For the archicture of the module we explored different configurations of:

Original: • Embedding dimension • Number of heads • f ϕ output dimension
Revised: • Embedding dimension • Number of heads • f ϕ output dimension

Original: Due to the long training time, we evaluated each model after 50 million timesteps.
Revised: Due to the long training time, we evaluated each model after 50 million timesteps.

Original: This training length should be sufficiently long for the purpose of observing notable differences between models’fractional
Revised: This training length should be sufficiently long for the purpose of observing notable differences between models’

Original:  success and reward.
Revised: fractional success and reward.