End-to-end Invariance Learning with Relational Inductive Biases in Multi-Object Robotic Manipulation
Authors: Davide Mambelli
Frederik Träuble
Stefan Bauer
Bernhard Schölkopf
Francesco Locatello
Venue: Submitted to ICLR 2023 (2023)

Abstract: Although reinforcement learning has seen remarkable progress over the last years, solving robust dexterous object-manipulation tasks in multi-object settings remains a challenge. In this paper, we focus on models that can learn manipulation tasks in fixed multi-object settings \emph{and} extrapolate this skill zero-shot without any drop in performance when the number of objects changes. We consider the generic task of moving a single cube out of a set to a goal position. We find that previous approaches, which primarily leverage attention and graph neural network-based architectures, do not exhibit this invariance when the number of input objects changes while scaling as $K^2$. We analyse effects on generalization of different relational inductive biases and then propose an efficient plug-and-play module that overcomes these limitations. 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.

Under review as a conference paper at ICLR 2023
END-TO-END INVARIANCE LEARNING WITH RE- LATIONAL INDUCTIVE BIASES IN MULTI-OBJECT ROBOTIC MANIPULATION
Anonymous authors Paper under double-blind review
ABSTRACT
Although reinforcement leamnngg has seen remarkable progress over the last years, solving robust dexterous object-manipulatie tasks in multi- settings remains a challenge. 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. We consider the generic task of moving a single cube out of a set to goal position. We find that previous approaches, which primarily leverage attention and graph neural network -based architectures, do not exhibit this invariance when the number of input objects changes while scaling as K2, We analyse effects on generalization of different relational inductive biases and then propose an efficient plug-and-play module that overcomes these limitations. 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.
INTRODUCTION
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. [015 Silver et al.
2016; Vinyals et al., (219).. On the other hand, applying deep RL on real-world grounded robotic setups 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. 2018 Deisenroth et al.) 2011 Kroemer et al. 2018 Andrychowicz et al.
2020
Rajeswaran et al.
2017 Lee et al.) 2021 Funk et al. 2021).
The RL problem in robotics setups is much more challenging (Dulac-Amold al. 2019). Compared to discrete toy environments, state and action spaces are continuous, and solving tasks typically requires long-horizon time spans, where the agent needs to apply long sequences of precise low-level control actions. Accordingly, exploration under casy-to-define sparse rewards becomes only fcasible with horrendous amounts of data. This is usually impossible in the real world but has been likewise hard for computationally demanding realistic physics simulators. To alleviate this, manually designing task-specific dense reward functions is usually required, but this can often lead to undesirable or very narrow solutions. Numerous approaches exist to alleviate this even further by e.g. imitation learning from expert demonstrations (Abbeel & Ng 20044, curricula (Narvekar et al. 2020), or model-based learning (Kaelbling et al. 1996). Another promising path is to constrain the possible solution space of learning agent by encoding suitable inductive biases in their architecture (Geman et al. 1992).
Choosing inductive biases that leverage the underlying problem structure can help to learn solutions
that facilitate desired generalization capabilities Mitchell| (1980) Baxter (2000): Hessel et al. (019).
In robotics, multi-object manipulation tasks naturally suit compositional description of their current state in terms of symbol-like entities (such as physical objects, robot parts, etc.). These representations can be directly obtained in simulator settings and ultimately hoped to be inferred robustly from learned object perception modules Greff et al. (2020a): Kipf et al. (202I); Locatello et al. (2020). While such principle considered crucial for any systematic generalization
compositional understanding ability Greff et al. (2020b); Spelke ( (1990); Battagla et al. (2018); Garnelo et al. (2016) it remains an open question how to design agent that process this type of input data to leverage this promise.

---

Under review as a conference paper at ICLR 2023
Train distractors
Zero-shot generalization to any * distractors
=
T
==
= End
T End
T End
Figure I:
We consider the following problem setting: There are 2 identical cubes with unique identifiers in the arena. During episode, random cube (here green) is selected to be transported to a random episode-specific unused distractor during that episode but could
goal-location. The remaining cube (here dark red) acts be the cube to be transported the next episode. Can train RL agents that learn such manipulation tasks within a fixed object setting training with only one distractor and extrapolate this skill zero-shot when the number of distractors changes?
Consider transporting cube to specified target location and assume the agent has mastered solving this task with three available cubes (see Figure for this problem setting which forms the basis of our work). We refer to cubes different from the ones to transport as distractors. two additional cubes, the solution to solve this task would remain the same, but the input distribution changed. So why not just train with much more cubes and use a sufficiently large input state? First, the learning problem requires exponentially more data when solving such task with more available cubes (see Figure 2 left plot, red line), making infeasible already for only half a dozen cubes. Second, we do not want to put any constraint on the possible number of cubes test time but simply might not have this flexibility at train time. In such task, the agent must be able to learn end-to-end policy that is invariant the number of distractors, while distribution shift. Therefore, we demand to learn challenging manipulation tasks in fixed object settings and extrapolate this skill zero-shot without any drop in performance when the number of cubes changes. Achieving this objective will be the primary goal of this work.
we now introduced
observing at training time such input
essential prerequisite for achieving this is to endow the agents with the ability to process a variable number of input objects by choosing suitable inductive bias. One popular model class for these object encoding modules are graph neural networks (GNNs), and vast line of previous approaches builds upon the subclass of attentional GNNs to achieve this (Zambaldi et al. 2018 Li al. 2019)
Wilson & Hermans 2020 Zadaianchuk et al., 2020). Another possible model class although largely ignored for robotic manipulation tasks so far- relation networks that process input sets by computing relations for every object-object pair Santoro et al. (2017). As is the case for attentional GNNs, relation networks scale quadratically in the number of input objects (see Figure right plot), which can become computationally impractical for training and inference with many objects. In this work, are therefore primarily interested learning invariances regarding the manipulable objects of robotic agents' environment but not invariance regarding the objects' properties such as shape or weight which an orthogonal robustness objective in machine learning. Without properly accounting for such modularity, task complexities can otherwise grow exponentially.
Main contributions. As a first main contribution, we will demonstrate that utilizing popular attention- based GNN approaches does not achieve the desired invariance. Instead, uttention-bart GNN approaches fail to extrapolate any learned behavior changing number of objects. a solution, we will then present support to build agents upon relational reasoning inductive biases. In addition to a more traditional implementation of relational network, we introduce a linearized relation network module (LRN) which improves the computational complexity in the number of objects from quadratic to linear. Finally, we show that agents based on this proposed module can extrapolate and therefore generalize zero-shot without any drop in performance when the number of objects changes. We will present supporting evidence for this currently under-explored type of generalization and necessary requirements for multi-objes manipulation across two challenging robotics tasks.
2

---

Under review as a conference paper at ICLR 2023
Timestaps solve Pushinn with Distractors
Space Complexity of Reasoning Modules inference Time
sincttes 2.3
3063
Figure 2: Scaling challenges in multi-object manipulation. Left: Scaling behavior of training frames required for learning to push a specified object to a target pose with additional distractors. Training agents with a naive fixed-input size MLP encoder is constrained to predetermined object number and scales exponentially, Data points dashed lines are extrapolated from proposed models trained with only 2 distractors they achieve zero-shot generalization for any more distractors. Right: Required Memory for action prediction with diffcrent number of distractors using the ATTN object reasoning backbone vs. RN and LRN backbone.
2
PROBLEM SETTING
This work considers tasks where the goal is always to transport one of the N₂, movable objects, i.e. cubes, to target location, and all other No - cubes act distractors. 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. Specifically, we define the objective of any learning agent as follows: First, the agent needs to learn to solve the task at train time under constant and a minimal number of manipulable objects Ntrain Second, the agent needs to be able to zero-shot extrapolate solving the equivalent task under Niest + Nirain objects in the scene. Therefore, the agent needs to learn an invariant behavior with respect to the episode-specific distractors.
Note that our task formulation (i.e. moving/picking a selected object out of a variable number of other objects) can be seen as a generic subtask of more general manipulation tasks. If we can solve this task, more complex multi-object goals could be solved more easily by specifying one object-goal pose after another. As an example, stacking tasks can be trivially decomposed as sequence of lifting tasks where each goal is on top of the previous one. Crucially, constant number of available objects because varying the number of objects during training might limit in the real world. Possible real-world limitations include but are not limited the increased complexity of designing automatic resetting systems able to deal with a variable number of objects which are necessary train an agent over many episodes. We ultimately need adapting agents that can generalize when additional objects are added or removed from the environment scene.
want learn such tasks with
Background on compositional multi-object manipulation tasks. Given the compositional en- vironment structure of such real-world grounded manipulation tasks, we are concerned with input states s, E S that comprise a set of objects describing individual rigid objects in the scene such as movable cubes {sq, $2 .... SNN } or the individual parts of robot such as each finger in our setup {s} Sz, ..., sy,,, J. To facilitate a coal-conditioned formulation, we similarly feed additional eivironmeent-peccfiic information in the form of goal object 8 Concerning our tasks, the goal information can be provided via the object state 8 that describes the target pose of cube. To identify the object to be transported, we simply add unique identifier feature to each object state and specify the object to be moved through the same identifier in the goal object. Translated to robotic RL setting, the above state-space factorization requires a policy that is parameterized by a function a, T( (1$$...... sy,, Sy.... ..... syy. s }). Our demand is that this policy deal with variable number N, of input objects independent of their order. In other words, the policy needs to be permutation invariant, i.e. it produces the same output for any possible permutation of its input objects.
3

---

Under review as a conference paper at ICLR 2023
3 RELATED WORK
Flavors of this problem have a long history and are typically being approached by various suitable inductive biases that enable reasoning upon such object-oriented structured representations (Džeroski
et al. 1998| Džeroski et al.) 200T Kaelbling et al. 2001 Van Otterlo (002). A popular class of
neural architectures that offer these properties are various sub-types of graph neural networks (GNNs)
with attention- approaches or interaction networks being a special sub-case (Bronstein et al.
2021 Veličković et al. 2017: Vaswani et al. 2017 Battaglia et al. 2016 Santoro et al.] [207] & Welling 2016). Zambaldi et al. (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 constant number of features extracted from a CNN. 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. Interestingly, most of the previous work for learning similar manipulation tasks in multi-object settings uses similar object-encoding modules based on multi-head dot product attention (MHDPA) layers to process the object- based input representations, On related stacking tasks using robotic arm with ID gripper [i et al. (2019) propose model architecture that uses an MHDPA -quuivallen GNN. Compared to our work, their method heavily builds upon handcrafted curricula to achieve some degree of fading zero-shot adaptation to related multi-ohject tasks. Another recent approach SMORI. by Zadaianchuk al. (2020), which similarly passes object-oriented representations to an attention- and goal-conditioned policy that learns the rearrangement of objects. The method can adapt to related test tasks with fewer objects but needs to be fine-tuned on each of such tasks individually. Finally, Wil- son & Hermans (020) proposes another approach to learn to manipulate many objects with a focus on sim-to-real. The agent model is conditioned on state embedding obtained through GNN and builds upon an MHDPA- MHDPA-equivalent module, with a set of object poses as input. Due to slightly different task specification, multi-object generalization is achieved more easily as the movement to solve their tasks for changing number of objects remains the same. Recently, Zhou et al. (2022) studied similar gen- eralization setting as ours and again with focus on attention- based architectures. Importantly, their results validate the persisting limitation with attentional backhones as they also see decreasing gener- alization capabilities in terms of success rate with an increasing number of objects in the OOD task. In summary, previous approaches suffer from various fundamental limitations regarding our desider- ata. First, the compute of MHDPA-based architectures scales quadratically with the number objects (see Figure 2) which is especially limiting for training and inference with N. >> 1. Second, to achieve any meaningful extrapolation for different object numbers, the training domain must cover a variable and typically large number of objects. Third, the performance at test time typically decreases for a new number of objects, and fine-tuning is required to recover the original performance.
Related work on relation networks. Our method is motivated by means of an alternative perspective on this problem via relation networks Raposo et al.) (2017): y [Santoro et al. (2017). Relation networks process input sets by computing relations for every object-object pair, followed by a permutation invariant aggregation function. Relation networks scale quadratically in the number of input objects Santoro et al. (0117). They are predominantly used for various non-RL tasks such as visual and text-based question answering, few-shot image comparison, or object detection Santoro et al. (2017); Sung et al. (2018); Hu et al.] (2018). Importantly, we are not aware of RL applications upproaching this problem utiliz utilizing relation networks as relational inductive biases. In addition, standard relation networks typically operate on a fixed number of input objects and suffer from the same scalability issues as attention. Our linear relation network module alleviates this limitation.
4 AGENTS WITH OBJECT REASONING MODULES
We will now describe the agent model used in this work and then present the linear relation network (LRN) that we propose as an alternative object reasoning module for learning compositional multi- object manipulation skills (see Figure 3 for a summary). In Section we will benchmark this module against an attention-based module (ATTN) aiming to resemble the prevailing choice of architecture in previous approaches. Moreover, we will also compare our LRN module against a naive relation network implementation (RN) with inferior scaling in computational costs to further highlight the sole benefits of relations over naive attention. We refer to the appendix for further implementation details
4

---

Under review as a conference paper at ICLR 2023
Object Reasoning Modules
At Attention based Architecture (ATTN)
Groundaruthy Features
Ohjeche-wite
SD
Esyellom
Comparty OUX"
B: Pairwise Relations Architocture (RN)
Obient tyon
Objests-wise
FC Layerhlom
Gum
C: Linear Pairwise Relations (LRM)
Objests-wise
FC Layerhorm
Curpeurity CXRR)
Figure 3: Summary of object reasoning modules studied in this work. Left: In a first step. ground-truth feature states encoded into fixed-size object tokens via object-type specific small MLPs. Right: An object-encoder module processes this token set and outputs fixed-size representation (highlighted in pink). This representation forms the input of any downstream predictor required to learn the optimal policy. Besides a naive attention-baas approach chosen an object encoder module by most previous works (Box A), we propose to use architectures based relational reasoning and discuss two variants. The first presented approach is relation network architecture that combines all possible pairwise relations, and refer this module RN (Box B). The second main method our Linear Relation Network architecture which only uses relations involving goal object, thereby scaling linearly in the number of imput tokens (Box C). We refer to this model as LRN.
4.1
AGENT MODEL
Our agents will build upon the following standard architecture scaffold. In first step, ground-truth simulator state features represented by objects are encoded into fixed-size object tokens o, Rd via object-type specific small MLPs. In our case, there are three robotic states for each finger, goal objects, and physical objects for the cube. Next, an object-reasoning encoder module processes this token set and outputs fixed-size representation (highlighted in pink in Figure 3 Finally, this representation forms the input of any downstream predictor required to learn the optimal policy. Here, will assume small MLP to predict the action (policy); however, representation to predict the value (value network) or Q-Value (when combined with an action) when using Q-Learning or actor-critic approaches. The primary focus of this work on proposing simple and flexible object-encoder module that overcomes the limitations of previous approaches.
likewise the
4.2 LINEAR RELATION NETWORK REASONING MODULE
Similar to the original Relation Network from Santoro et al. (017)) used for visual question-answerte tasks, module builds upon the idea of reasoning about future actions from learned relations between the given objects. However, in the context of goal-conditioned RL, seems reasonable that the next optimal action is primarily governed by the relations of the available objects with the particular goal of interest. In contrast, most other relations might be only important in sparse and context-dependent way. We, therefore, argue focus on these relations, and propose the following adjusted composite representation function
LayerNorm fo E
(1)
where the O; E dRd are common-size embeddings of the K input objects. The objects are the scene's moveable elements in setting, which means cubes and the three robotic finger states. The goal information of the task is being embedded into separate object og E Rd. First, the shared learnable function ga computes the embeddings r; e Rd for each object-goal relation. For our experiment go is 4-layer MI.. Next, dimension-wie
operation the K relations is applied, which
5

---

Under review as a conference paper at ICLR 2023
Training With distractors
Zero-shot generalization with any distractors
1.0
0.8
Pushing
RN
0.0
success
Lifting
0.6
0.6
1.00
300 400
600
700 800
timestops [millions]
Figure 4: Training curves and zero-shot generalization from fixed multi-object train domains. Left: Training curves for PushingCube (top) and LiftingCube (bottom) tasks in one-distractor environments (5 random seeds per model). Our RN and LRN modules perform favorably against attention-base based implementation in terms final success, sample efficiency, and variance across different seeds. Right: Evaluating the corresponding zero-shot generalization on equivalent N-diitractor environments for both tasks. Our LRN and RN modules show strong extrapolation ability of their learned manipulation skills to environments with many more objects.
serves permutation invariant aggregation function. The final representation z Rd is obtained after processing this intermediate representation by another learnable function fo followed by layer normalization operation (Ba et al. (016)). The function f can serve the purpose of learning further necessary relations from the K object-goal relations. We will further discuss the models' empirical ability learn any other needed relations through the combination of this minimal set of goal-object relations Section 5.5 our experiments, we will use a simple linear transformation for f.. The layer normalization ensures that the final representation remains in distribution for changing number of input objects. The overall complexity of this goal-conditioned linear relation network rea- soning module is linear in K, i.c. O(K). See Figure 3 (Box C) for summary. To assess the impact of inferring the K object-goal relations only, we will also implement an equivalent variant that considers all possible K2 relations, similar to the naive implementation of relation networks. To the best of our knowledge, they have been largely ignored for any similar robotic manipulation tasks and general- ization far. Inspired by the original relation network module, will refer this module RN
EXPERIMENTS
We now aim to evaluate the previously discussed 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. We will start by outlining the experimental setup in Section 5. for two particularly generic tasks pushing and lifting objects. We will then discuss the learning behavior terms of sample efficiency and performance of proposed modules against the common attention-batn based approach from previous works Section 5.2 We close this section by demonstrating the models' zero-shot OOD extrapolation skills in Section 5.3 then by analyzing the impact of the train domain in section 5.4 and, finally, by looking at the ability of the model of reasoning about higher order relations section 5.5 Additional ablations on the architecture together with detailed analysis the models' OOD generalization and learned relations are discussed in the appendix (9.2. 9.2.4).
Note that the original relation network module for question answering tasks computes relations that are additionally conditioned on the question and operates on fixed number of input features such that normalization not required for these problem domains.

---

Under review as a conference paper at ICLR 2023
5.1 EXPERIMENTAL SETUP.
Both tasks are derived from the simulated robotic manipulation setup and tasks proposed in Causal- World (Ahmed et 2020) which builds upon the robotic trifinger design (Wüthrich et al. 2020 Bauer et al. 2021). Compared to previous work on this setup, we attempt to solve tasks in scenes that contain more than cube Dittadi et al. (202T); Ahmed et al. (2020); Allshire et al. (202T). Specifically, the goal is to move any of the 2 available cubes to a freely specifiable target position. We cover 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. Such lifting task involves learning to hold the cube stable in the air at variable height. We refer to these tasks as PushingCube and LiftingCube, respectively. As discussed in Section [ we specify the object to be rearranged through matching object identifier in one of the goal objects' features. For both tasks, we use a generic reward function designed work for any single-goal task in this environment. 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) second reward term that encourages moving reward term for efficient exploration that gives a positive reward signal for moving the target object when being further away from the goal. Across all tasks, we will report the factional success as the universal evaluation metric at the last timestep of an episode. Tasks can be visually considered solved with score around 80%, and we can therefore ensure an objective and interpretable performance measure (Ahmed et al. 2020 Dittadi let al. 202T). All models are being trained under the described reward with PPO (Schulman et al. 2017) distributed workers with 5 random seeds each. Initial object poses and the goal pose randomly sampled at the beginning of each episode, and the episodic task needs to be solved within 30 simulation seconds 2500 frames (10 seconds per object). We compare the instantiation of our proposed RN and LRN network architectures against a representative attention- from Zambaldi al. (0188), referred to ATTN. We ensured a fair comparison of this baseline through extensive hyperparameter optimization (see Section 9.2.5). We refer to the supplement for a comprehensive account of further implementation details.
the target object closer to the goal; (3) third puriosity-insping
baseline derived
5.2 LEARNING TO MANIPULATE OBJECTS FROM FIXED MULTI-OBJECT TRAIN DOMAINS.
Quantitative results for learning manipulate objects from the 2-cubes (i.e. 1 distractor) train domain are summarized in Figure for PushingCube in the top panel and LiftingCube in the lower panel. We refer the project-sitd" for additional videos of trained agents. Across tasks, we can see that our proposed backbone performs very favorably over the attention-based ATTN baseline. Specifically, both relation network variants (RN and LRN) reliably succeed in learning to solve the tasks almost perfectly with very high scores. On the other hand, ATTN exhibits much higher variance across different seeds and takes longer to train. Interestingly, RN and LRN reveal very similar robust learning behavior in terms of training time steps. A significant difference, however, is the required compute. For environments with 9 cubes, the LRN module requires computing only K = 12 relations (9 cube-goal and 3 finger-goal relations), whereas RN combines (K + 1)2 169 relations to process its observation set, highlighting the favorable properties of our LRN module.
Summary. The relation network reasoning modules and our LRN model represent promising alternative to attention-basse architectures commonly used in prior works for learning to manipulate an object on multi-object train domain. Our module outperforms these approaches in final success, sample efficiency, and variance across different seeds.
5.3 ZERO-SHOT GENERALIZATION TO N-OBIECT ENVIRONMENTS.
We now evaluate the trained models' generalization ability when the number of objects is different from training. Our main result to assess this question is summarized in Figure 4 (right column). Remember that all agents have been trained for pushing and lifting an object in 2-cube environments, synonymous of distractor. Testing these agents' generalization in N-object environments with No + 2 reveals a clear separation between the relation network modules and the attention-based architecture. Agents with the ATTN module cannot generalize to a changing number of objects and experience strong and consistent fading ability to solve the task the more the number of objects deviates from the training environment. ATTN-based agents fail to generalize under our test tasks
7

---

Under review as a conference paper at ICLR 2023
ATTN
RN
LRN
1.0
Distractors
0.6
Distractors
Distracto's
0.6
fractions
0.4
distractors
distractors
distractors
Figure 5: Zero-shot generalization across different numbers of training distractors. All three modules (left: ATTN: middle: RN: right: LRN) were retrained under baseline always peaks performance of cubes. Crucially, our RN module can generalize to any number of objects in the scene if trained on at least 2 distractors (3 cubes in total). The LRN module achieves this with as little distractor at train time.
3 distractors with 5 seeds each. The ATTN its train domain, quickly vanishing generalization to any new numher
completely when adding more than 2 additional objects. Next, a naive relation network module (RN) achieves significantly improved performances generalization over the ATTN baseline. Importantly, RN still suffers from fading generalization under more distractors, especially in the more complex task of LiftingCube. In contrast, our LRN module can maintain strong train-time performance for any tested number of objects. As we observe no relevant deteriorating trend, we suspect that this generalization capability could possibly even hold beyond the 9-distractor environments tested here. Interestingly, the same trends can be seen for environments with only cube and no further distractor. Even though it seems intuitive that agent should be able to solve pushing lifting object when no further distractors are in the scene, the ATTN module seems to overfit on the specific 1-distractor input train distribution strongly and cannot solve these tasks without exactly one distractor. On the other hand, RN and LRN module-based agents can robustly maintain their respective train-time performances, presumably due to their relation-basee processing that better facilitates invariance w.r.t. the input objects.
Summary. Agents based on the relation network modules show a strong extrapolation ability of their learned manipulation skills to environments with many more objects. In contrast, we do not observe such zero-shot generalization with commonly employed attention based architectures, which experience substantial decrease in performance when deployed ditto--diitribution.
5.4 ROLE CHANGING THE TRAIN DOMAINS.
Previous generalization results obtained by the LRN were unexpected given that they were obtained training with 1 distractor only. In particular, we believed that at least a 2-distractor environment would be necessary since also includes informative edge cases that the agent has to deal with to solve tasks with a large number of distractors. An example is when the selected cube is stuck between distractor cubes blocking its grasping. We found this to be a frequent scenario when the arena gets crowded. To verify if the training environment has effect on generalization performances, we repeated our experiments environments with up to distractors and additionally covered the limiting of distractors at train time at all. Our results presented in Figure 5 which shows how the zero-shot generalization under each architecture is affected by its training distribution. Concerning the ATTN module, observe the effect of the train distribution on generalization capabilities to be generally limited. Performances always peak in the training environment and monotonically decrease for Aitest # Ntraan, On the other hand, a steady improvement can be observed for the RN module when changing the training environment up until 2 distractors. Note that the RN already achieves some level of generalization when training with single distractor. appears phase change to robust performances happens when 2 distractors are used for training which yields a strong generalization for any test environment. Finally, we can observe that our LRN module while already achieving a strong generalization when trained with I distractor only can maintain the same level of performance when trained with 2 and 3 distractors. We, therefore, conclude that richer training distributions allow the RN to achieve generalization, whereas it does not actually change the overall behavior of the ATTN module. However, it is important to recall that, even if the RN achieves good generalization when trained on 2 distractors, the required number
8

---

Under review as a conference paper at ICLR 2023
of training steps improves significantly with respect to the I distractor environment highlighted in Figure [ (left plot). Summary. The training domain significantly affects generalization performances. Our LRN module appears to be the most robust architecture to changes in the training domain it is the only architecture that already achieves strong generalization when trained with single distractor.
5.5 CAPTURING HIGHER-ORDER RELATIONS
The Linear Relation Network (LRN) does not explicitly compute most of the possible relations. allowing it to scales only linearly with the number of objects. A question that naturally arises what happens when particular relation that the LRRN does not compute is strietly required to solve the task. In fact, example of this situation is even present in the pushing task: Here, the agent must reason about where the fingers are with respect to the cube to touch it and then grasp it. However, since the LRN only computes relations to the goal, the fingers-cube relation is not explicitly available to the agent. The only way to manipulate the cube using available relations is to combine the sube-goal and fingers-goal relations using the goal as a reference. This capability means the agent can extrapolate and use higher-order relations. The LRN successfully solves the pushing task, which therefore shows its ability infer essential relations for the downstream task. More precisely. the higher-order reasoning is handled by the subsequent layers of the architecture after computing the relations, which is the MLP policy head in our case.
LIMITATIONS
Choice of environments. With limited compute budget in mind and significant engineering effort necessary in robotics applications, we decided to specifically focus only on the TriFinger platform because successfully learning multi-object policies on this platform is still a long-standing open challenge given its highly flexible mechanical design which implicates large and factored action space at the same time. Importantly, the platform comprises multiple manipulation challenges like learning coordinated dexterity such that it is more challenging than many other robotic environments where gripping is defined as a Id space representing the finger distance Brockman et al. (2016).
Task formulation. The here proposed architecture is designed to handle single-goal environments. While being an important next step towards the long-term goal of multi-goal manipulation policies, we focused suitable architecture that can deal with variable number of input objects. As discussed in Section 2 our tasks mirror generic subtasks of general manipulation tasks such that future work can build upon such robust backbones and test how different (hierarchical) planning algorithms can master multi-goal tasks Yang et al. (2021): Kroemer et al. (2015).
RL algorithm. In our experiments, we decided to use PPO Schulman et al. 2017) due it being a popular choice in RL and favourable robustness to different hyper- parameter choices, Extensive results on pushing and lifting from ground-truth features on the same setup with single cube in Ahmed et al. (2020) indicate that methods like TD3 Fujimoto et al.] (2018) or SAC Haarnoja et al. (2018) perform very similarly to PPO under the same reward structure and observation space.
CONCLUSION
In this work, we presented new object encoder module based on relation networks, enabling reinforcement learning agents to learn challenging compositional manipulation tasks in multi-object robotics environments. 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. While previous approaches mostly build upon quadratically scaling attention based architectures to process the set of input objects, approach scales only linear with the number of objects. At the same time, our method outperforms these previous approaches at train time. Most importantly, we extrapolate and generalize shoo-shot without any drop in performance
show that our agents when the number of objects changes which is a major limitation of most previous approaches. As next steps we also envision using the proposed module for solving similar tasks directly from pixels, ie. combining our architecture with learned object perception modules. Locatello et al. (2020); Greff
al. (2020a):
Kipf et al. (202I).

---

Under review as a conference paper at ICLR 2023
8 REPRODUCIBILITY STATEMENT
The main results of this paper come from the implementation of object reasoning modules: LRN, RN ATTN. Each module is extensively described respectively in section 4.2 section 9.1.5 section 9.1.4 These modules, as described in section 4 are sub-components of an agent architecture which is detailed in section 4.1 These agent architectures are trained optimizing for the reward structure highlighted in section 9.1.1I which is unique across the 2 tasks (PushingCube, and LuftingCube) studied in this work. The robotic manipulation tasks we work on are easily reproducible using the Causal World environment (aaammeett/eeddddd AAhmed et al. [2020), which provides out-of-the-box support for multi- object environments. The details for the RI algorithm used in all experiments are described in section 9.1.2
10

---

Under review as a conference paper at ICLR 2023
REFERENCES
Pieter Abbeel and Andrew Y Ng. Apprenticeship learning via inverse reinforcement learning. In Proceedings of the twenty-first international conference on Machine learning. pp. 2004.
Ossama Ahmed, Frederik Träuble, Anirudh Goyal, Alexander Neitz, Yoshua Bengio, Bernhard Schölkopf, Manuel Wüthrich, and Stefan Bauer. Causalworld: A robotic manipulation benchmark for causal structure and transfer learning, 2020.
Arthur Allshire, Mayank Mittal, Varun Lodaya, Viktor Makoviychuk, Denys Makoviichuk, Felix Widmaier, Manuel Wüthrich, Stefan Bauer, Ankur Handa, and Animesh Garg. Transferring dexterous manipulation from gpu simulation to a remote real-world trifinger. arXiv preprint arXiv:2/08.09779, 2021.
OpenAl: Marcin Andrychowicz, Bowen Baker, Maciek Chociej, Rafal Jozefowicz, Bob McGrew, Jakub Pachocki, Arthur Petron, Matthias Plappert, Glenn Powell, Alex Ray, et al. Learning dexterous in-hand manipulation. The International Journal of Robotics Research, 39(1):3-25 2020.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint DrXiv:1607.060 2016.
Peter W Battaglia, Razvan Pascanu, Matthew Lai, Danilo Rezende, and Koray Kavukcuoglu. Interac- tion networks for learning about objects, relations and physics. arXiv preprint arXiv: arXiv:161..0022. 2016.
Peter W Battaglia, Jessica B Hamrick, Victor Bapst, Alvaro Sanzhez-Gonzalez, Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, et al. Relational inductive biases, deep learning, and graph networks. arXiv preprint arXiv: arXiv:1806.00 1806.01261, 2018.
Stefan Bauer, Felix Widmaier, Manuel Wüthrich, Niklas Funk, Julen Urain De Jesus, Jan Peters, Joe Watson, Claire Chen, Krishnan Srinivasan, Junwu Zhang, al. A robot cluster for reproducible research in dexterous manipulation. arXiv preprint arXiv:21099 1099522099.1055,, 2021.
Jonathan Baxter. A model of inductive bias learning. Journal of artificial intelligence research, 12: 199, 2000.
Lukas Biewald. Experiment tracking with weights and biases, 2020. URL httph::t//// /ttp://wwwww/ w/wd://d/wwwwwwbbbb..m com/ Software available from wandb.com.
Greg Brockman, Vicki Cheung. Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojeiech Zaremba. Openai gym. arXiv preprint arXiv: 1606.00.44 2016.
Michael M Bronstein, Joan Bruna, Taco Cohen, and Petar Veličković. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. arXiv preprint arXiv:2104.13478, 2021.
Mare Peter Deisenroth, Carl Edward Rasmussen, and Dieter Fox. Learning to control low-cost manipulator using data-efficient reinforcement learning. In Robotics: Science and Systems VII, volume 7, pp. 57-64, 2011.
Andrea Dittadi, Frederik Träuble, Manuel Wüthrich, Felix Widmaier, Peter Gehler, Ole Winther, Francesco Locatello, Olivier Bachem, Bernhard Schölkopf, and Stefan Bauer. The role of pretrained representations for the ood generalization of rl agents. arXiv preprint arXiv:2107.056a 2021.
Yan Duan, Marcin Andrychowicz. Bradly C Stadie, Jonathan Ho, Jonas Schneider, Ilya Sutskever, Pieter Abbeel, and Wojciech Zaremba. One-shot imitation learning. arXiv preprins arXiv:1703.07336, 2017.
Gabriel Dulac-Arold, Daniel Mankowitz, and Todd Hester. Challenges of real-world reinforcement learning. arXiv preprint arXiv:1904.12990 2019.
Saso Džeroski, Lue De Raedt, and Hendrik Blockeel. Relational reinforcement learning. In Interna- tional Conference on Inductive Logic Programming, pp. 11-22. Springer, 1998.

---

Under review as a conference paper at ICLR 2023
Sašo Džeroski, Luc De Raedt, and Kurt Driessens. Relational reinforcement learning. Machine Learning, 43(1):7-52, 2001. doi: 10.1023/A1000664411358888 URL htttp:://bi..rrg/// 1023/A: 100./32:1007664155553
Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor- critic methods. In International conference on machine learning, pp. 1587-1596. PMLR, 2018.
Niklas Funk, Charles Schaff, Rishabh Madan, Takuma Yoneda, Julen Urain De Jesus, Joe Watson, Ethan Gordon, Felix Widmaier, Stefan Bauer, Siddhartha S Srinivasa, et al. Benchmarking structured policies and policy optimization for real-world dexterous object manipulation. arXiv preprint 2 105.02087, 2021.
Marta Garnelo, Kai Arulkumaran, and Murray Shanahan. Towards deep symbolic reinforcement learning. arXiv preprint 1609.5550 2016.
Stuart Geman, Elie Bienenstock, and René Doursat. Neural networks and the bias/variance dilemma.
Neural computation, 4(1):1-58, 1992.
Klaus Greff, Raphaël Lopez Kaufman, Rishabh Kabra, Nick Watters, Chris Burgess, Daniel Zoran, Loic Matthey, Matthew Botvinick, and Alexander Lerchner. Multi-object representation learning with iterative variational inference, 2020a.
Klaus Greff, Sjoerd van Steenkiste, and Jürgen Schmidhuber. On the binding problem in artificial neural networks. arXiv preprint arXiv:2012.05200 2020b.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy
maximum entropy deep reinforcement learning with stochastic actor. In International conference on machine learning, pp. 1861-1870. PMLR, 2018.
Matteo Hessel, Hado van Hasselt, Joseph Modayil, and David Silver. On inductive biases in deep reinforcement learning, arXiv preprint N0ki::197.002000 2019.
Han Hu, Jiayuan Gu, Zheng Zhang, Jifeng Dai, and Yichen Wei. Relation networks for object detection, 2018
Michael Janner, Sergey Levine, William T Freeman, Joshua B Tenenbaum, Chelsea Finn, and Jiajun Wu. Reasoning about physical interactions with object-oriented prediction and planning. arXiv preprint 1XX.::18122.19722 2018
Leslie Pack Kaelbling, Michael L Littman, and Andrew W Moore. Reinforcement learning: A survey. Journal of artificial intelligence research, 4:237-285, 1996.
Leslie Pack Kaelbling, Tim Oates, Natalia Hernandez, and Sarah Finney. Learning in worlds with objects. In Working Notes of the AAAI Stanford Spring Symposium on Learning Grounded Representations, pp. 31-36, 2001.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 14X2:1412..690. 2014.
Thomas Kipf. Gamaleldin Elsayed, Aravindh Mahendran, Austin Stone, Sara Sabour, Georg Heigold, Rico Jonschkowski, Alexey Dosovitskiy, and Klaus Greff. Conditional object-centric learning from video. arXiv preprint arXiv:2111. 12594, 2021.
Thomas N Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. arXiv preprint arXiv:1609.00 029077 2016.
Oliver Kroemer, Christian Daniel, Gerhard Neumann, Herke van Hoof, and Jan Peters. Towards learning hierarchical skills for multi-phase manipulation tasks. In 2015 IEEE International Conference on Robotics and Automation (ICRA), pp. 1503-1510, 2015. doi: 10.1109//CRA.201
7139389.
Oliver Kroemer, Simon Leischnig, Stefan Luettgen, and Jan Peters. A kerdel-based approach to learning contact distributions for robot manipulation tasks. Autonomous Robots, ((33):8116000 2018.
12

---

Under review as a conference paper at ICLR 2023
Alex X Lee. Coline Manon Devin, Yuxiang Zhou, Thomas Lampe, Konstantinos Bousmalis, Jost To- bias Springenberg, Arunkumar Byravan, Abbas Abdolmalcki, Nimrod Gileadi, David Khosid, et al. Beyond pick-and-place: Tackling robotic stacking of diverse shapes. In 5th Annual Conference Robot Learning, 2021.
Richard Li, Allan Jabri, Trevor Darrell, and Pulkit Agrawal. Towards practical multi-object manipu- lation using relational reinforcement learning, 2019.
Francesco Locatello, Dirk Weissenborn, Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. Object-centric learning with slot attention, 2020.
Tom M Mitchell. The need for biases in learning generalizations. Department of Computer Science, Laboratory for Computer Science Research 1980.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, al. Human-level control through deep reinforcement learning. nature, 51 (7755))::529533 2015.
Sanmit Narvekar, Bei Peng, Matteo Leonetti, Jivko Sinapov, Matthew E Taylor, and Peter Stone. Curriculum learning for reinforcement learning domains: A framework and survey. arXiv preprint WrXiv:2003.04961 2020.
Aravind Rajeswaran, Vikash Kumar, Abhishek Gupta, Giulia Vezzani, John Schulman, Emanuel Todorov, and Sergey Levine. Learning complex dexterous manipulation with deep reinforcement learning and demonstrations. arXiv preprint arXiv: 1709. 700..7007777777 2017.
David Raposo, Adam Santoro, David Barrett, Razvan Pascanu, Timothy Lillicrap, and Peter Battaglia. Discovering objects and their relations from entangled scene representations. arXiv preprint arXiv:1702.05068; 2017.
Adam Santoro, David Raposo, David G. T. Barrett, Mateusz Malinowski, Razvan Pascanu, Peter Battaglia, and Timothy Lillicrap. A simple neural network module for relational reasoning, 2017.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms, 2017.
David Silver, Aja Huang, Chris Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schritwieser, Ioannis Antonoglou, Veda Panneershelvam,, Marc Lanctot, et al. Mastering the game of go with deep neural networks and tree search. nature, 529(7587):444-489, 2016.
Elizabeth S Spelke. Principles of object perception. Cognitive science, 14(1):29-56, 1990.
Adam Stooke and Pieter Abbeel. ripyt: A research code base for deep reinforcement learning in pytorch, 2019.
Flood Sung, Yongxin Yang. Li Zhang, Tao Xiang, Philip H. S. Torr, and Timothy M. Hospedales. Learning to compare: Relation network for few-shot learning, 2018.
Laurens van der Maaten and Geoffrey Hinton. Visualizing data using t-SNE. Journal of Ma- chine Learning Research, 9:2579 2605, 2008. URL.ttp::// jmlr www.//pprr//........
Haadermaatendor html
Martijn Van Otterlo. Relational represstttons reinforcement learning: Review and open problems. In Proceedings of the ICML, volume 2, 2002.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need, 2017.
Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, and Yoshua Bengio. Graph attention networks. arXiv preprint arXiv: 10903, 2017.
Oriol Vinyals, Igor Babuschkin, Wojciech M Czamecki, Michaël Mathieu, Andrew Dudzik, Junyoung Chung. David Choi, Richard Powell, Timo Ewalds, Petko Georgiev, et al. Grandmaster level in starcraft ii using multi-agent reinforcement learning. Nature, 575(7782)3300-3355 2019.
13

---

Under review as a conference paper at ICLR 2023
Matthew Wilson and Tucker Hermans. Learning manipulate object collections using grounded state representations, 2020.
Manuel Wüthrich, Felix Widmaier, Felix Grimminger, Joel Akpo, Shruti Joshi, Vaibhav Agrawal, Bilal Hammoud, Majid Khadiv, Miroslav Bogdanovic, Vincent Berenz, al. Trifinger: An open-source robot for learning dexterity. arXiv preprint TrXiv:2008.0359 2020.
Xintong Yang, Ze Ji, Jing Wu, Yu-Kun Lai, Changyun Wei, Guoliang Liu, and Rossitza Setchi. Hierarchical reinforcement learning with universal policies for multistep robotic manipulation. IEEE Transactions Neural Networks and Learning Systems, pp. 1-15, 2021.
Andrii Zadaianchuk, Maximilian Seitzer, and Georg Martius. Self-supervised visual reinforcement learning with object-centric representations, 2020.
Vinicius Zambaldi, David Raposo, Adam Santoro, Victor Bapst, Yujia Li, Igor Babuschkin, Karl Tuyls, David Reichert, Timothy Lillicrap, Edward Lockhart, Murray Shanahan, Victoria Langston, Razvan Pascanu, Matthew Botvinick, Oriol Vinyals, and Peter Battaglia. Relational deep reinforcement learning. 2018
Allan Zhou, Vikash Kumar, Chelsea Finn, and Aravind Rajeswaran. Policy architectures for compo- sitional generalization in control. 2022.
14

---

Under review as a conference paper at ICLR 2023
type size
position
orientation
linear velocity
angular velocity
identifier
end-effector position
Cube
Goal
3
3
Finger
Tahle Number variahles used to describe each element of the state for each object type
APPENDIX
9.1 IMPLEMENTATION DETAILS
9.1.1 REWARD DESIGN
We used the following same reward for both tasks:
R₂ = a1Pt
+ o2 (((or.g)) d(or-1, 9(-1)
(2)
+ #as(1 (c) loog [d(o,, 0--1) 1e 1e-
where & denotes the time step, Pt [0, the fractional overlap with the goal cube time e, E RR is the end-effector position, o E R the cube position of the object that is supposed to be used for the task, 9t R the goal position, and -(-,-) denotes the Euclidean distance, Across all experiments we set a [100, 250, 10] and _ 0.05.
9.1.2 TRAINING SCHEME FOR AGENTS USING PPO
We use the distributed PPO implementation from r rlyyt (Stooke & Abbeel, 2019) with discount factor 0.98, entropy loss coefficient 0.01, learning 0.00025, value loss coefficient 0.5, gradient clipping norm 0.5, 40 minibatches, gae lambda 0,95, clip ratio 0.5 and epochs. We applied linear learning rate scheduling and Adam optimizer Kingma Ba (2014).
9.1.3 SIMULATOR INPUT FEATURES
The ground truth representations for each object type described in Table T Cubes described using their position and orientation and each corresponding velocity. addition, every cube is marked with unique identifier which makes the set of cubes heterogeneous. Similarly, the goal is specified via its cartesian position, orientation, and an identifier variable used to specify which cube has be moved. 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. Each finger of the robot has degrees of freedom, and it is described by a 9-dimensional vector composed by its joint positions, joint velocities, and end-effector position.
9.1.4 ATTENTION MODULE BASELINE
The baseline we use in this work exploits a MHDPA module (Vaswani et al. 2017) operating on the set o E [dxx(K1) composed of common-sizi embedding of the input objects o, E Rd, and the goal Og E Rd with
g(())(())²
attn
softmax
v(o)
Vd
where atta is the of K embeddings scaled from the MHDPA module. In this module, the coal-conditionin element is obtained by dealing with the goal as an object. This set is then mapped to the fixed-size embedding by an aggregation function and normalized to avoid changes scale when the number of objects differs between training to testing.
K+1
LayerNorm fo
Σ
attn
This reasoning module resembles the one from Zambaldi et al. (2018), and it is used to introduce relational inductive biases into the learning process. The same self-attention layer is also used by
15

---

Under review as a conference paper at ICLR 2023
Linear RN
Name
Operation(Shape)
Output Shape
Cubes
Input(-)
Weubes x 18
Finger
Input(-)
N fingers 10
Goal
Input(-)
1 x 12
Doube
FC(18 - +444 ReI.U
N Woubes X
9 finger
FC(10 + +44) ReL.U
N fingers X 64
ggoot
FC(12 - +ReLU
X 64
g
4x FC(64 64 + 128) ReLU
N Neubes + N fingers X 128
sum
Sum pooling along dimension of relations
128
fout
FC(128 - 20)
20
LN
LayerNorm
20
MLP Policy Head
Name
Operation(Shape)
Output Shape
h
FC(20 + 256)+
256
l
FC(256 - 256)
256
la
FC(256 - 9)
9
Table 2: Architecture details Agent with LRN reasoning module
Zadaianchuk al. (2020) build goal-conditioned policy that is able to handle set-based input representations. Similarly, it also describes approaches that rely on Graph Attention Networks (GAT) /Veličković et al., (017) to build policies while using the common assumption in multi-object manipulation of fully connected graph (Li et al., 2019 Wilson & Hermans 2020). Self-attention when computing the attention matrix takes into consideration each object-object pair, which makes scale quadratically with the number of objects in the input set, as highlighted in Figure (A).
9.1.5 RELATION NETWORK MODULE
To have a further element of comparison we introduced Relation Networks as relational reasoning module in the context of reinforcement learning.
K+1 K+1 Σ
Σ ga(o;, o.))
z = LayerNorm
f
(5)
Relation Networks have been proposed in Santoro et al. (2017) and they compute (K + 1)² em- beddings r, Rd, referred to relations, using the shared function go. Relations aggregated using a sum and then normalized to obtain fixed-size embedding used as input for the MLP policy head. Similarly self-attention Relation Networks need consider every object-object interaction, making them scale quadratically the number of objects. A key aspect this architecture that go always takes input a pair of objects no matter the total number of objects in the input set. This property makes the architecture robust to changes in the number of objects as long as the objects' statisties do change.
9.1.6 LINEAR RELATION NETWORK ARCHITECTURE
Architectural details used in all our experiments for the Linear Relation Network are provided in T
9.2
ADDITIONAL RESULTS
9.2.1 ASSESSING ARCHITECTURAL CHOICES OF OUR LRN MODULE.
We also tested the impact of various design choices for proposed linear relation network module. Results are summarized in Figure 6 (left). Specifically, we focus on 3 questions: First, do we need relations at all? To answer. equivalent non-relational single-object feedforward MLPs, ie. go (o,). We found that the module cannot even learn the task on the training domain without relations. Second, want to understand
repeated the LRN experiments but replaced the relations with
16

---

Under review as a conference paper at ICLR 2023
0.-
LBN 20
LNN 2D curssey
sparoe)
0.4
0.4
datenced
Figure 6: Zero-shot generalization the linear relation network module on the PushingCube task concerning various architecture choices and the proposed reward structure. Left: We find that we cannot solve the task when dropping the object-goal relations (purple curve). We also observe that even though our agents can learn the task using average instead of the LayerNorm, our model experiences fäding extrapolation performance (orange curve). Finally, we find a much weaker fading in generalization ability when providing too few or too many representational dimensions (green and red curves). Right: When ablating on the individual terms of our reward function, we first find that removing the distance term does not affect the overall performance (red curve). Whereas the sparse term only combined with the curiosity term is sufficient to solve the task successfully.
if we could use a simple average instead of the LayerNorm after summing over all relations. Our experiments show that even though we can learn the task in the train environments, the agents reveal deteriorating generalization the more objects we add to the task. We hypothesize this to due the changing variance of z which is kept constant when doing a layer normalization but not an average. Lastly, we analyzed the role of the representations' dimension, for which we used de = 25. When repeating the experiments with d_ 10, the zero-shot generalization slowly decreases to about 60 % success with 9 objects. On the other hand, when allowing for d: = 50 or even d, = 100, the agent achieves slightly elevated train time performance with slight drop of about 5 % and more pronounced drop of about 20 % respectively when testing with up to 9 distractors. This result suggests that the representation dimension is optimally being adjusted for the right capacity requirements of the task.
9.2.2 ROLE OF THE REWARD FUNCTION.
We finally wanted to understand the impact of each term on the reward function we proposed. Figure (right) shows the generalization performance of corresponding agents trained with each variation of the reward function. 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. The first notable observation is that removing the distance term does not affect the overall performance. Therefore, we can conclude that the sparse term is sufficient to solve the task successfully. 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.
9.2.3 DISTRIBUTION SHIFT FOR OOD TASKS
To better understand the difference between the architectures, we visualized the representations of each Reasoning module with the help of lower-dimensional embeddings. Each model has been tested for 200 episodes on environments with the number of distractors ranging from 0 up to 9, The representations obtained as output have been mapped a low-dimensional space using t-SNE (van der Maaten & Hinton 2008). These low-dimensional points let us compare the support of the embedding distribution between different test domains. 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. In Figure (Ta) we can observe the result obtained for the attention- based module when visualized in 2-dimensional space. The embedding distribution changes with the task, which causes the MLP policy drop in performances observed in Figure 4 In F(b) and Figure 7(c) respectively the RN and
policy head to operate out-of-distribution justifying the
17

---

Under review as a conference paper at ICLR 2023
distractors
distractors
2 distractors
distractors
distractors
distractors
distrecctor
distractors
distractors
d'stractog
Distribution Shift Attention Network (training distribution: 2 distractors)
distractors
distractors
distractors
distractors
distractors
(b) Distribution Shift Relation Network (training distribution: distractors)
0 distractors
distrectors
2 distractors
distractors
d distractors
5 distractors
distractors
distractors
distractors
9 d distractors
(c) Distribution Shift Linear Relation Network (training distribution: distractors)
Figure 7: OOD Distribution shifts
18

---

Under review as a conference paper at ICLR 2023
LRN output embeddings visualized. For both, can say that the approximated support of the distribution appears to invariant with respect to the task.
9.2.4 RELATIONS ANALYSIS
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. We used relations generated from 200 episodes with a number of distractors varying from 3 to 6. The obtained plot is presented in Figure g where cube-goal relations are shown in green, distractors-gout relations are visualized in blue, and fingers relations are colored respectively in black, brown, and red. Different relations space, suggesting that invariance to the number of distractors might be obtained via gating property learned by the architecture to filter out relations embedded into certain sub-spaces when computing the action.
mapped different sub-spaces of the representation
Relations over time We also analyzed how the norm of each relation type changes over time while the agent is solving the task. In Figure we show the result when testing on 5 distractors LRN-based agent and normalizing each type of relation by the number of relations belonging that type. Since we used the sum as the aggregation function, the norm can be associated with each relation's 's weight in the representation after the sum. Interestingly, the most important relations are those describing fingers and the target cube, whereas those linked to distractors have much smaller norm. This qualitative result suggests that the cube-goal relation might correlated with the cube-goal distance.
9.2.5 LIMITATIONS OF THE ATTENTION BASED NETWORK
Given the observed poor generalization results and the very high variance during training for the attention-base based baseline, we performed extensive hyperparameter tuning to ensure a fair comparison against our proposed RN and LRN modules. We used Bayesian search (Biewald 2020) across more than 550 hyperparameter configurations. We found that none of those improves over our set of hyperparameters meaningfully.
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. Regarding PPO,
we tunned:
Gradient clipping norm
Discount factor
Entropy loss coefficient
Value loss coefficient
GAE lambda
Learning rate
Ratio clip
For the architecture of the module, we explored different configurations of:
Embedding dimension
Number of heads
fo output dimension
Due to the long training time, we evaluated each model after 50 million timesteps. This training length should be sufficiently long for the purpose of observing notable differences between models' fractional success and reward.
Additionally, we tried add skip connection to the self-artention layer following what is done by Vaswani et al. (0017). This change reduced the variance across seeds but did not affect the overall result in terms of generalization.
19

---

Under review as a conference paper at ICLR 2023
40-
-0
-40
Figure 8: Low-dimensional t-SNE embedding relations generated over episodes from the Linear Relation Network. Different relations are color-coded (Black: FingerO-Godd Brown: Fingerl-Ger -Goal. Red: Finger2-Goal. Green: Cube-Goal. Blue: Distractors-Goal), The architecture learns to map different relations in different sub-spaces of the representation space.
nibegggl
Snner goal
goal
10
8
100
150
200
250
300
dinestopo
Figure 9: Relations changing throughout the task in LRN when testing 4 distractors.
20


Reviews:
[{"review_id": "zzK9AV88w_", "reviewer": "Reviewer_T4pA", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper considers two block manipulation tasks, where the number of train and test block are different. They (1) demonstrates the computation and performance degradation of graph attention and relational networks as more blocks are introduced at test time and (2) introduce a linear relation network module that empirically mitigates the problems mentioned in the aforementioned baselines. ", "summary_of_the_review": "The main claims of the paper surrounding are validated (i.e., the proposed method can generalize to completing target manipulation tasks with more distractor objects than seen during training). Additionally the presentation in Fig. 4 and 5, which some some of the main empirical results, are well executed. \n\nHowever, I currently recommend rejection of the manuscript for the following reasons: key details necessary to understand the method are missing (*W3-5*). The experiments currently seem limited to simple settings, further experimentation seems necessary to more rigorously verify the method and determine its impact (*W6-8*). Additionally, Fig. 2 is misleading as it contains extrapolation, which may not be obvious to a reader who is glancing at the paper (*W2*).\n\nPOST REBUTTAL: \nchanging score from 5 to 3. See \"Post rebuttal\" for more details.", "strength_and_weaknesses": "**Strengths**\n\n*S1.* The method is simple to implement and performant in the two tasks considered.\n\n*S2.* Experiments generally back the main claims of the paper related to generalization to new numbers of distractor cubes.\n\n**Weaknesses**\n\n*W1.* I am not convinced of the fidelity of the simulation environment. In Fig. 1 it appears that the red and yellow blocks are overlapping, which is not physically realistic. This calls into question the results. It seems possible that some agents are able to exploit such artifacts in the simulation environment. I suggest doing a more thorough investigation of the simulation environment and increasing the collision checking in pybullet to prevent such cases.\n\n*W2.* My understanding is that the trends in Fig. 2 (left) are extrapolated for the blue and green lines after x=1 and x=2 respectively. This could be misleading to a reader who looks at the figure at a glance. It is unclear if there is enough data to extrapolate these trends. I recommend running these experiments rather than extrapolating the trends.\n\n*W3.* Clarity: some of the task/experimental details are not provided. For example, how are the states of each object $s_i$ represented? As 2D coordinates, 3D coordinates, poses? What are the inputs and outputs to the networks?\n\n*W4.* What is the intuition on why LRN is able to generalize to different numbers of objects but the baselines cannot? Making this more clear in the manuscript could help guide the readers towards the key takeaways.\n\n*W5.* What are the small object specific MLPs? Is there one for each cube? If so how do you decide which MLP to use for each cube at test time when there are more cubes? Is it based on color? Please consider including more details in this section.\n\n*W6.* The method uses GT state information.\n\n*W7.* The method, while effective for the downstream tasks, is potentially lacking in generality. How would the method perform if shapes that are not cubes are included, such as YCB objects? What about if state was inferred? Are there other tasks that this method could be applied to? \n\n*W8.* The paper argues that the network could be learning higher level interactions in deeper layers of the network. This suggests another baselines: only tokens are fed into $g(.)$ without any pairing. This may further elucidate the effects of the inductive bias considered in the paper.\n\n**Minor**\n\n*M1.* In the abstract, I recommend rewording the phrase: \"allows agents to extrapolate and generalize zero-shot to any new object number.\" This reads as an unsupported claim as it is not empirically verifiable and there is no proof included about why this is guaranteed. \n\n*M2.* I recommend removing phrases like \"horrendous amounts of data,\" which feel a bit editorial and imprecise. \n\n*M3.* I suggest removing the phrase \"we suspect that this generalization capability could possibly even hold beyond the 9-distractor environments.\" Alternatively, I suggest pushing the method with even more objects to see at what point it can no longer handle more distractor objects.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "forfh09aJo", "reviewer": "Reviewer_4jkK", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper presents a simple relational architecture that can learn and generalize on multi-object manipulation tasks. The main benefit of the proposed method is the linear-scale complexity in terms of the number of objects (i.e., distractors) in the single-goal tasks. The proposed method is tested on robotic manipulation tasks and compared against a pairwise relational model and an attention-based model. The empirical results demonstrate that the proposed linearized relation network (LRN) module improves the computational complexity and extrapolates to an unseen number of objects.", "summary_of_the_review": "Overall, I have three main concerns on the significance of this work.\n\n*1. Single-object goal task setting is limited*\n\nThis work assumes the task always involves a single object. In section 2, authors claim that the proposed task formulation can be seen as a generic subtask, so that composing such subtasks can complete a bigger task involving multiple objects. However, each subtask may still involve more than one object. Imagine a subtask of navigating to a certain location. While navigating,  the robot should not break any objects that may be useful for later subtasks or harm any person. Also if the robot is carrying a glass of water, it should try not to spill the water while navigating. The robot may need to choose a safer path depending on the state of the “glass of water” object (which is not the goal object). Given that the proposed task setting is quite limited, the strict dependence of the proposed method on the task setting further limits the significance of this work (continued in the next section).\n\n*2. The proposed method is limited to the proposed task setting*\n\nThe main contribution of this work comes from the linearized relational module. And such linearization would not harm the performance only on a single-goal task, where relations between other objects are not or less important. As discussed above, relation between other objects may be as important as goal objects. In section 5.5, authors discuss that the higher-order relations can be captured by the MLP policy head. However, this means the proposed LRN cannot handle the higher-order relations alone, which makes the main claim less convincing.\n\n*3. Baselines are limited*\n\nAs discussed in the previous point, the proposed method exploits the inductive bias (i.e., goal-conditioned task where goal is based on single object) in the task design. However, it seems the baselines do not benefit from such specific task design. The proposed method is compared with two baselines: attention-based model and pairwise relational model. However, both models do not have any inductive bias regarding the goal-based task, hence it is not a fair comparison. A more fair baseline can be appending the goal token to every paired input tokens in the pairwise relational model (i.e., goal + obj1 + obj2 instead of obj1 + obj2). Another important baseline to compare with is the set representation (e.g., [1]) with goal awareness (e.g., appending a goal token before/after the aggregation). On the other hand, there are many works on linearizing the quadratic architectural complexity such as LinFormer [2]. Comparing with the goal-aware version of linearized attention-based or relational models can strongly support the main claim.\n\n[1] Zaheer et al., \"Deep Sets\", 2017\n\n[2] Wang et al., \"Linformer: Self-attention with linear complexity\", 2020\n", "strength_and_weaknesses": "**Strength**\n* The paper is easy to read and well structured\n* The scope of problem is clearly defined\n* The proposed idea is simple and intuitive\n\n**Weaknesses**\n* Scalability: the proposed method seems limited to the proposed task\n* The baselines are limited\n* Task setting (or scope of the problem) is limited\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "_RF80xTVEu-", "reviewer": "Reviewer_E3aJ", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors propose a variant of relational networks to reason about interactions between objects and between objects and an agent in a manipulation task. They train policies with this architecture on a few simulated manipulation tasks, and demonstrate out-of-distribution generalization to variants of the task not seen during training.", "summary_of_the_review": "The approach seems sound, but the experiments feel a bit lackluster. I would not be upset if this paper were accepted, but I also don't feel particularly excited about it.", "strength_and_weaknesses": "Strengths:\n1. The proposed method is clearly motivated. The authors seek to encode strong inductive biases to enable out-of-distribution generalization.\n2. The experiments conducted are thorough. The authors compare their method against a few baselines on two manipulation tasks, and show that the trained policies extrapolate to situations beyond what they are trained on. \n\nWeaknesses:\n1. I can't help but feel that the empirical results are not very impressive. In both manipulation tasks, the relational reasoning boils down to \"ignore the distractor objects\", and the generalization amounts to increasing the number of distractor objects beyond what the policy was trained it. Are there other, more diverse tasks that better showcase the compositional generalization that this approach should enable?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "_MA4K1w_Tw", "reviewer": "Reviewer_Pmgk", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper is considering the problem of learning robotics manipulation in a multi-object setting. The specific problem considered is the successive repositioning of two cubes by a manipulator in the presence of a variable number of distractor cubes. \n\nThe authors argue that previous approaches, based on GNNs do not learn policies that generalize to a variable number of distractors. They proposed a representation that is based on a relation network, and show that this approach can extrapolate to a variable number of distractors. The authors also propose a \"linearized relation network\" module which works in a linear rather than a quadratic computational complexity. ", "summary_of_the_review": "The paper considers the case of a robotics manipulation of a cube in the presence of distractor cubes in a simulation setting. The contribution of the paper is a simplified relation network based representation where only the relations of the goal cube is considered. The paper shows that with this representation the learned policy generalizes better to variable number of distractors. \n\n---\n\nI read the authors answer to my and the other reviewers comments. The answers do not change my rating of the paper. ", "strength_and_weaknesses": "+ Good justification for the problem of generalizing to a variable number of distractors in robotics manipulation. \n- The proposed setting is extremely simple: manipulating individual objects, in simulation, under the assumption of perfect knowledge. The authors assume that the resulting policy will be so simple that it can be represented by a \"small MLP\" that doesn't even need specifying. \n- The contribution is essentially limited to formula (1) for the representation. This formula essentially amounts to ignoring all the relations except the binary relations ones between the goal cube and the individual distracting cubes. While this is indeed O(n), it is not a novel contribution, only a simplifying assumption. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "RfbzcOb05i", "reviewer": "Reviewer_vRaR", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper presents an RL-approach for multi-object manipulation. The main idea is to extrapolate a learned skill to manipulate one object in a zero-shot fashion onto a number of additional objects. To do this, the authors propose an exension of the recent Relation Network RN by Santtoro et al. In simulated experiments, the authors show that their approach improves over RN and an attention-based module. ", "summary_of_the_review": "Interesting approach, but some more background is needed for the non-expert reader to assess the novelty.  ", "strength_and_weaknesses": "Strengths: The paper is clearly written and the results show a clear improvement over the compared methods.\n\nWeaknesses: The RN framewor, which is used as a basis here, could have been explained in more detail. It is hard to reproduce the entire algorithm only based on equation (1). ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject