Policy Architectures for Compositional Generalization in Control
Authors: Allan Zhou
Vikash Kumar
Chelsea Finn
Aravind Rajeswaran
Venue: Submitted to ICLR 2023 (2023)
Keywords: Reinforcement Learning
Imitation Learning
Compositionality

Abstract: Several tasks in control, robotics, and planning can be specified through desired goal configurations for entities in the environment. Learning goal-conditioned policies is a natural paradigm to solve such tasks. However, learning and generalizing on complex tasks can be challenging due to variations in number of entities or compositions of goals. To address this challenge, we introduce the Entity-Factored Markov Decision Process (EFMDP), a formal framework for modeling the entity-based compositional structure in control tasks. Geometrical properties of the EFMDP framework provide theoretical motivation for policy architecture design, particularly Deep Sets and popular relational mechanisms such as graphs and self attention. These structured policy architectures are flexible and can be trained end-to-end with standard reinforcement and imitation learning algorithms. We study and compare the learning and generalization properties of these architectures on a suite of simulated robot manipulation tasks, finding that they achieve significantly higher success rates with less data compared to standard multilayer perceptrons. Structured policies also enable broader and more compositional generalization, producing policies that extrapolate to different numbers of entities than seen in training, and stitch together (i.e. compose) learned skills in novel ways. Video results can be found at https://sites.google.com/view/comp-gen-anon.

Under review as a conference paper at ICLR 2023
POLICY ARCHITECTURES FOR COMPOSITIONAL GENERALIZATION IN CONTROL
Anonymous authors
Paper under double-bling review
ABSTRACT
Several tasks in control, robotics, and planning can be specified through desired goal configurations for entities in the environment. Learning goal-conditionel policies is a natural paradigm to solve such tasks. However, learning and generalizing on complex tasks can be challenging due to variations number of entities or compositions of goals. To address this challenge, we introduce the Markov Decision Process (EFMDP), a formal framework for modeling the entity- based compositional structure in control tasks. Geometrical properties of the EFMDP framework provide theoretical motivation for policy architecture design, particularly Deep Sets and popular relational mechanisms such graphs and self attention. These structured policy architectures are flexible and can be trained end- to-end with standard reinforcement and imitation learning algorithms. We study and compare the learning and generalization properties of these architectures on a suite of simulated robot manipulation tasks, finding that they achieve significantly higher success rates with less data compared to standard multilayer perceptrons. Structured policies also enable broader and more compositional generalization, producing policies that extrapolate to different numbers of entities than seen in training, and stitch together (i.e. compose) learned skills in novel ways. Video results can be found at https: //tp::::/siiee //stte.../ss//ese//// google. com /ttp/viiei//e/e//eee/eeeeeee/eee wwwww//oopp/oeeeeeotoet/ee//ot/e compp-ge--ee gen-anon
INTRODUCTION
Goal specification is powerful abstraction for training and deploying Al agents (Kaelbling, 1993) For instance, object reconfiguration (Batra al., 2020) tasks, like loading plates in dish- washer or arranging pieces on a chess board, can be described through spatial and semantic goals for various objects. In addition, the goal for scene can be described through compositions of goals for individual entities in it. Through this work, we introduce a new framework for modeling tasks with such entity-centric positional structure, applicable to domains like robotic manipulation, multi-agent systems, and strategic gameeplaying. Subsequently, we study policy architectures that can utilize structural properties unique to our framework for goal- conditioned reinforcement and imitation learn- ing. Through experiments in simulated robot manipulation tasks, we find that our policy ar- chitectures exhibit significantly improved Iearn- ing efficiency and generalization performance compared to standard multi-layer perceptrons (MLPs), as previewed in Figure 1. More importantly, our architectures are capable of learning near-optimal policies in complex table top manipulation tasks where MLP baselines completely fail.
Training: 3 cubes
Train
0.5
Success Rate
0
2
Samples
twr
Test: 3 cubes
1.0
Test
4
5
cubes
Sclf Attention (Ours)
MLP (baseline)
Figure 1. A family of tasks where the agent is trained re-arrange three cubes (top-left), but tested zero-shot cubes (bottom-lefit. RI. with stan- dard MLPs fails to even learn the 3-cube task. Our self- attention policy, on the other hand, successfully learns and extrapolates.
re-arrange
Consider the motivating task of arranging pieces on a chess board using a robot arm. A naive specification would provide goal locations for all 32 pieces simultaneously. However, can

---

Under review as a conference paper at ICLR 2023
immediately recognize that the task is composition of 32 sub-goals involving the rearrangement of individual pieces. This understanding of compositional structure can allow us to focus on one object at a time, dramatically reducing the size of effective state space and help combat the curse of dimensionality that plagues RL (Sutton and Barto, 1998; Bertsekas and Tsitsiklis, 1996). Moreover, such compositional understanding would make an agent invariant to the number of objects, enabling generalization to fewer or more objects. Most importantly, it can enable reusing shared skills like pick-and-place enhancing the learning efficiency. We finally note that successful policy cannot completely decouple the sub-tasks. For example, if piece must be moved to square currently occupied by another piece, the piece in the destination square must be moved first.
The generie Markov Decision Process (MDP) framework as well as policy architectures based on ML.ss lack the aforementioned compositional properties. To overcome this limitation, we turn the general field of "geometric deep learning" (Bronstein et al., 2021) which is concerned with the study of structures, symmetries, and invariances exhibited by function classes. We first introduce the Entity- Factored MDP (EFMDP), subclass of the generic MDP, environments with multiple entities (e.g. objects). We then characterize the geometric properties of EFDMP relative to the generic MDP. We subsequently show how set-based invariant architectures like Deep Sets (Zaheer et al., 2017) and relational architectures like Self- Self-Attention (Vaswani et al., 2017) and Graph Convolution (Kipf and Welling, 2016) are well suited to leverage the geometric properties of the EFMDP. Through experiments, we demonstrate that policies and critics parameterized by these architectures can be trained to solve complex tasks using standard RL and IL algorithms, without assuming access to any options or action primitives.
formal model for decision making in
Our Contributions. This paper organized into sections that present our three main contributions:
1. We develop the Entity-Factored MDP (EFMDP) framework, a formal model for decision making in tasks comprising of multiple entities (e.g. objects), and characterize its geometric properties.
2. We show how policies and critics parameterized by set-based invariance models (e.g. Deep Sets) or relational models (e.g. Self-Attention and Graph Convolution) can leverage the geometric properties of the EFMDP.
3. We empirically evaluate these structured architectures on a suite of simulated robot manipulation tasks (Figure 4), and find that they generalize more broadly while also learning more efficiently. Compared to MLPs, structured policies improve success rates by more than 50 on extrapolation tests which vary the numbers of entities in the environment, and by 10x stitching tests that require recombining learned skills in novel ways to solve new unsteen tasks.
2 PROBLEM FORMULATION AND ARCHITECTURES
In this section, we first formalize our problem setup by introducing the entity- MDP (EFMDP). This setting is capable of modeling many applications including table-top manipulation, scene reconfiguration, and muti-agent learning. Subsequently, we also introduce policy architectures that can enable efficient learning and generalization by utilizing the EFMDP's unique structural properties.
2.1 PROBLEM SETUP
We study learning paradigm where the agent interact with many entities in environment. The task for the agent is specified in the form of goals for some subset of entities (including the agent). We formalize this learning setup with the Entity-Factore Markov Decision Process (EFMDP).
Definition I (Entity- Factored MDP). An EFMDP with N entities is described through the tuple: M == (4, E 0, A, P, R, 7). Here l and È are the agent and entity state spaces, g is the subgoal space and A is the agent's action space. The overall state space S U EN has elements s (u,e1,-- eN) and the overall goal space G == g has elements g = (g .N). The reward and dynamics are described by:
R ( ()g) == R ((i(ei, Di, w))
(1)
P(s' (s)ssaa== = TP ((1l, (c))N)) (v. (e))
(2)
fors,s/ E, a E A. and g E g.
2

---

Under review as a conference paper at ICLR 2023
The EFMDP is goal-conditioned MDP (Kaelbling. 1993; Schaul et al., 2015) with additional structure. Each entity is associated with a specific reward F; = F(e,, 9s. u), which are aggregated follow various rules like requiring "all" entity
together to reward the agent. The aggregation subgoals to be satisfied or "any" entity subgoal be satisfied. We also note that the EFMDP does not force the entities be exchangeable identifying properties distinguishing each entity. The ultimate objective for the learning agent in is to learn policy x² S x G A that maximizes the long term rewards, given by:
indistinguishable, since the entity state space may contain
DO
arg max [(() [A))== = ES
Σ
Y RRis, g)
The EFMDP can model several applications including table-top manipulation, scene reconfiguration. multi-agent learning, and strategic game playing. At the same time, the EFMDP contains more structure and symmetry compared to the standard MDP model, which can enable more efficient learning and better generalization. The crucial symmetry exists in the reward and dynamics, which treat entity-subgoal pairs as unordered sets and are therefore invariant under permutations.
Property 1 (EFMDP Permutation Symmetry). For any permutation E SN (the group of all permutations of N items), the reward satisfies R(oss ag) = =R(s, and the transition dynamics satisfy (ass' os, a) = P(s'|s, (sss|ss,aee a) for any s, E S and a E where:
os (u, Cr(1): ... CC((N) and og (go(1)) 9o(())
This property captures the general intuition that the ordering of entity-subgoal pairs is arbitrary and not relevant to the actual environment. We also prove that any optimal policy and the optimal value function
Agent
permutation invariant.
Proposition 1 (Policy and Value Invariance). In any EFMDP with N entities, any optimal policy X² S optimal action-valie function Q* :S x Ax are both invariant permutations of the entity-subgoal pairs. That is, for any a E Sx:
Subgoals
A and
T (os, og) n (s, g) and (" (s,
Entities
This is a direct consequence of the permutation symmetry in reward and dynamics; we provide proof in Appendix A. Note that Proposition 1 only talks about the optimal policy and value function, the permulation symmetry does hold for every policy. In fact, the permutation symmetry does not hold for most commonly used architectures like MIPs. In the next subsections, use Proposition I to guide architecture design in reinforcement and imitation learning on EFMDPs. We show that certain "entity-centric" model classes achieve invariance for every policy and value in the class, readily utilizing the structure and symettries afforded by the EFMDP.
Figure In an EFMDP, an agent inter- acts with entities that have correspond- ing subgoals. This framework model rearrangement, strategic game playing, and multi-agent systems. In this "push and stack" example, the agent must move the green cube to its subgoal, indicated by the green sphere, and then stack the yellow cube on top of the green cube.
2.2 MULTILAYER PERCEPTRONS (MLPs)
Standard RL and IL approaches assume they are solving generic MDP, and do not use any additional structure. The generic approach is thus parameterize the learned policy by an MLP, which takes fixed size input vector and applies alternating layers of affine transforms and pointwise nonlinearities produce fixed size output vector. To implement क (s., g) with an MLP we arrange the contents of (s, g) into single long vector using concatenation:
vec(s. g) := Concatenate u1, eI, EN F1, .N
Denoting the action of the MLP as a vector-to-vectes vector function MLP(-), our policy is defined (s. g) := MLP((vc(s. g)). Since MLPs expect input vectors of a fixed dimension, testing tasks with more entities requires zero padding the inputs during training to ensure consistent input dimensionality across all tasks.
3

---

Under review as a conference paper at ICLR 2023
Deep Sets Policy
Self Attention Policy
Σ
AC
7
SAC
Σ
A
Agent Entities Goal
Agent Entities Goal
Figure 3. Visualizations of implementing an entity-based goal conditioned policy using either Deep Sets (left) or Self Attention (right). The policy т (s. g) 1-> a receives state s = (u, ci, exx containing agent state u and entity states e,. The goal (gr,- - 9N) contains subgoals for cach entity. Both policies arrange the input into N vectors y, (u, et, 9t). outputs, and maps the result to an action using MLP p(-). The self attention encoder SA(-) produces output 21," ... EN and uses self-attention model interactions between entities/subgoals. The 2; are mapped to an action by summation and an MLP p(-).
per entity. The Deep Set policy processes each yi independently with MLP
Φ(-), aggregates
2.3
DEEP SETS
The MLP policy represents a "black-box" approach generic MDPs that fails to guarantee permu- tation invariance (Prop. 1). As result, MLPs might require significant amount of data to learn the necessary permutation invariance. contrast, the Deep Sets (Zaheer et al., 2017) (DS) architecture can guarantee permutation invariance of subgoal-entity pairs by construction. Given a set of vectors x = (Ii,"", ;TN), constructs a model of the form:
DS(x) == [ Edlar))
(6)
where p and & are themselves typically MLPs. DS(.) is invariant to ordering of {z,). since E,)) is agnostic to the ordering of elements. More surprisingly, Zaheer al. (2017) showed that deep sets can represent any permutation invariant function of x, given that á are sufficiently expressive. Karch al. (2020) introduced Deep Sets for instruction following policies in 2D environment, though to our knowledge, Deep Sets remain underexplored in more complex environments.
We now present simple but general approach for implementing invariant policies using Deep Sets for any EFMDP. For this, we arrange the entity-subgoal pairs as a set {(e1,9) ........ (e. 9N)]. We also include the "global" or "shared" agent state u to every entity-subgoal pair. The Deep Set then produces produces an action from this set:
π (s, g) == DS (LW)))))
y/1 = Concatenate u. Cs. 9i
(8)
Eg
Figure 3 (left) visualizes how the input is arranged and processed by by the Deep Sets policy.
2.4 RELATIONAL. MECHANISMS: GRAPHS AND SELF ATTENTION
Although Deep Sets can represent any invariant policy in theory. its design aggregates representations for all the entities through single summation and then requires the MLP p to handle any interactions between them. For tasks involving complex entity-entity interactions, we might desire a stronger relational inductive bias. Recent relational RL (Džeroski et al., 2001) approaches often model the state as graph and use graph neural networks (GNNs) (Gori et al., 2005; Scarselli et al., 2008) to implement policies dynamics models. As GNNs invariant to permutations of their nodes (Bronstein et al., 2021). they can also satisfy our EFMDP invariance condition Prop. 1 if we construct the policy input as a graph of entity-subgoal pairs.
We implement and evaluate two simple relational policy architectures based on (1) graph convolutional networks (GCN) (Kipf and Welling, 2016) and (2) self attention (SA) (Vaswani et al., 2017). Although originally developed for sequence processing applications, previous relational RL architectures have already used self attention as graph message-passing mechanism (Zambaldi et al., 2018; Li et al., policy in general EFMDP, we consider the input
2020). To use either GNN-style architecture as a complete graph where each node corresponds to an entity and the corresponding node features are {yi,-. . yy J. where vector ys is defined in Eq. 8. The GNN component consists of either multiple GCN layers or multiple self attention layers and transforms the input graph into an output graph with node features (. , =2)
that now capture relationships between the nodes.
4

---

Under review as a conference paper at ICLR 2023
Finally, the policy pools the 2,'s together by summation and project the result to action A using a small MLP p(-). Figure (right) illustrates the self attention policy design. In addition to satisfying permutation invariance, the relational policies either graph convolution (GCN) self attention (SA) to produce intermediate representations 2; that include interactions between the inputs, which can be stronger inductive bias complex tasks.
EXPERIMENTS AND EVALUATION
In this section, we aim study the following questions through experimental evaluation.
How efficiently do structured policies learn given entity-centric task?
Can the structured policies extrapolate to more or fewer entities?
3. Can the structured policies solve tasks containing novel combinations of subtasks, by stitching together (i.e. composing) learned skills?
Extrapolation and stitching are particularly interesting as they require generalization to novel tasks with additional training. This is particularly useful when deploying agents in real world settings with enormous task diversity. Press et al. (2021) showed that self attention can achieve interesting sequence-lene lenghh extrapolation behaviors in natural language processing tasks, which suggests that these architecture classes may also display interesting generalization capabilities in control tasks.
Environment Description. We seek to answer our experimental questions suite simulated robotic manipulation environments, where the policy provides low level continuous actions control a Fetch robot and interact with any number of cubes and switches. There are three subtasks: to push cube to desired location the table, to flip switch to specified setting, to stack one cube on top of another. The higher level tasks can involve multiple cubes or switches and compose many subtasks together, framework: the robot is the agent, the cubes and switches are entities, and the goal specifies desired cube locations or switch settings.
shown in Figure 4, These environments fit naturally into the EFMDP
We organize the environments into families to test learning and generalization. Environments in the its corresponding subgoal. The
N-Push family require re-arranging N cubes by pushing each N- H-Switch family requires flipping each of N switches to its specified setting, and the N-Switch + N- Push family involves re-arranging N cubes and flipping N switches. We test extrapolation by varying N within a family at test time, which changes the number of entities: for example we train a policy in 3-Switch and evaluate in 6-Switch. As another example, we test stitching by training a single policy on 2-Switch and 2-Push, then evaluate it on 2-Switch 2-Push which requires combining the switch and pushing skills together a single trajectory. Note that entity-entity collisions are disabled in non-stacking tasks for the main results, but Appendix D.2 repeats the N-Push experiments with collisions enabled. Appendix B gives a full description of our environments.
Baselines and Comparisons. Our main comparisons task as regular MDP (Sec. 2.2), and (b) an "oracle" that manually coordinates solving one subtask at time. We construct subpolicies for the oracle by training one policy each distinct subtask (pushing, flipping switches, and stacking). The oracle chooses an initial entity and subgoal arbitrarily, and uses the corresponding subpolicy until that subtask is solved. The oracle then selects the appropriate
with: (a) baseline MLP that models the
Switch
Push
Switch Push
Stack
Figure 4. Illustrations of the robot manipulation environments we study. They consist of subtasks such as pushing a cube to its (spherical) target, flipping a switch to a specified position, and stacking one cube top of another. The overall task involve multiple entities and subtasks well their combinations (pushing cubes and flipping switches).

---

Under review as a conference paper at ICLR 2023
Reinforcement learning
Imitation learning
Success Rate Succoss Rate Success Rate
1.0
1 I-Push
2-P Push
3-Push
Success Rate
1.0
3 3-Push
0.5
0.5
0.0
0.0
0
2 166
0.0
0.5
1.0
0
2
2000
4000
1.0
1-Switch
2-Switch
Switch
Success Rate
3-Switch
0.5
0.5
0.0
50000 150000150000
1.0
0.0
2000
4000
1 S--Swich++1-pP PPPhh
2 2-Switch+2 2-witch+2-Push PPssh
1.0
Self Attention
Success Rate
2 Swwich++-Puss
Deep Set
1.0
0.5
Graph Convolution
MLP
0.0
Oracle
0.0
2.5
5.0
7.5
0
1
2
0.0
166
2000
4000
Samples
Num demos
Figure 5. Training on environments of varying complexity using either reinforcement or imitation learning. Each row corresponds to a single environment family (N-Push, N-Switch, and N-Switch + N-Push), where environments with larger N contain more entities and are more complex. For RL (left), each plot is training curve of success rate vs the number of steps taken in the environment. RL with standard MLPs can solve the simpler tasks such as I-Push, but structured policies (Self Attention, Deep Set, and Graph Convolution) are superior on the more complex environments. For IL (right), we show success rates of behavior cloning against number of expert demonstrations in the dataset. The structured policies far outperform the MLP even when given less data. Shaded regions indicate 95% CIs over 5 seeds.
subpolicy for the next sutity-subgoal pair and continues until the entire task is complete. The oracle is not guaranteed to achieve a 100% success rate since it does not consider entity-entity interactions. An example failure mode is pushing one cube into position but knocking another one off the table while doing so. Still, as the oracle represents a hand-crafted hierarchical approach using an entity-based task decomposition, will compare the RL and IL agents' performance against the oracle in the following experiments.
3.1 EFFICIENCY OF LEARNING
To evaluate the learning efficiency of different architectures, we consider the N-Switch, N-Push, and N-Switch + N-Push environment families. We try N = 2, 3 for the first two families and N 1, for the latter, with larger N corresponding to more entities and more complex tasks within a family. Evaluation eriteria: An episode in the environment is considered successful only if all the sub- goals in the environment are achieved.
We separately train policies each environment in each family, using either RL II. approaches. For RL training we use DDPG (Lillicrap et al., 2015) with Hindsight Experience Replay (HER) (Andrychowicz et al., 2017), where we use the same architecture (either MLP, Deep Set, or Self Attention) to implement both the policy and critic. For IL we use behavior cloning to train policies fit a dataset of expert trajectories using mean-squared error loss. For each environment, we use trained RL agent to generate the corresponding expert trajectory datasets. See Appendix C for full RL and training details. Additionally, Appendix D.3 contains further experimental comparisons of RL. learning efficiency when training with Proximal Policy Optimization (Schulman et al., 2017) instead of DDPG, which show that the following RL results are not broadly sensitive to the particular choice of algorithm.
RL results. Figure 5 (left) shows RL training curves as function of environment samples. In the simpler 1-Switch and 1-Push environments, all methods learn to solve the task fairly quickly. Once there is more than one entity, however, the structured policies learn faster than the MLP. In harder environments like 3-Push N-Stitch + N-Push, the MLP fails to achieve non-trivial success rate. Both Deep Set and Self Attention match Oracle performance in environments except 2-Switch 2- Push. Graph Convolution additionally struggles in I-Switch 1-Push, but still outperforms MI.P

---

Under review as a conference paper at ICLR 2023
Train
Test
Success Rate
1.0
0.0
# cubes
Success Rate
0.5
3
5
6
switches
Success Rate
Self Attention
1.0
Deep Set.
Graph Convolution
0.5
Oracle
0.0
2
3 cubes
Training
switches
Figure 6. Extrapolation capabilities of RL-trained policies with different architectures. Each row depicts environment family with varying number of entities. Policies trained single environment from each family before being tested all the others, with additional training. Bar charts show structured policies (Self
each environment, with the hatched bars corresponding training environments. Attention, Deep Set, and Graph Convolution) extrapolate beyond the training environment to solve tasks with more or fewer entities than seen in training. while MLP policies struggle on more complex testing environments. Error bars 95% CIs 5 seeds.
Although they achieve similar asymptotic performance on most tasks, the Deep Set policy tends to learns faster than the others, possibly because it is simpler and has fewer parameters.
IL Results. The imitation learning results appear in Figure (right), where the x-axis now indicates the size of the training dataset used for behavior cloning. Similar to the RL setting, structured policies learn far more efficiently than the MLP in all environments. For example, in 3-Push with 5000 demonstrations, the MLP's success rate still nearly zero while the Self Attention policy has a nearly 100% success rate.
that the
Conclusions. MLP policies struggle learn complex tasks with many entities with both RL and II, likely due to the lack of entity-centric processing that the structured policies employ. The Deep Set policy typically learns faster than the others RL, and matches outperforms Self Attention IL with 1000 trajectories. Although the asymptotic performance of the entity-centric methods is typically similar, the relational methods are superior to Deep Set on 3-Push for both RL and II. 3-Push is one of the more difficult tasks, and relational policies may benefit from greater relational expressivity through self attention mechanisms. Overall, this experiment suggests that architectures that utilize the structure and invariances in EFMDPs learn substantially faster than black box architectures.
3.2 ZERO-SHOT EXTRAPOLATION CAPABILITIES
To test whether trained policies can extrapolate and solve test tasks containing more or fewer entities than seen in training, we use the N-Switch, N-Push, and SwStitch + N-Push environment families. For N-Push and N-Switch we train policy with RL on N = 3 and test with N E f1...., 66. For 2 and test on N E (I....,3). For testing, we use the RL agent checkpoint with the highest success rate in its training environment.
N-Switch + N-Push we train a policy with RL on
Results and Observations Figure 6 shows the test performance of these policies on each environment family as the number of entities N varies. The MLP only successfully learns the training task in the N-Switch environments, and it generalizes decently to fewer than switches, but fails completely in environments with more than switches.
In contrast, the structured policies generalize well and achieve zero-shot success rates comparable to exceeding the Oracle in most test environments. Notably, these policies well exceed oracle performance on Swithh despite training in 3-Switch. Interestingly, Self Attention policies fare poorly single-entity environments, perhaps because the self attention mechanism relies critically
7

---

Under review as a conference paper at ICLR 2023
Train
Test
Success Rate Success Rate
1.0
0.5
0.0
Push
Switch
2 Push
2 Switch
2-Push
Swwitthh
2-Switch 2-Push
1.0
0.5
2-Push
Stack
Push + Stack
Onuja Conmabuton
Ocaci
2-Push
Stack
Push + Stack
Deep dad
Tralinng
Figure 7. Left: train/test setups that require solving test tasks stitching together training skills, with additional data. Top: train on 2-Switch and 2-Push, test 2-Switch+2-4 P-shhh. Bottom: trained on 2-Push and Stack, test on Push + Stack. Right: average success rates by architecture. Deep Set and Self Attention policies are moderately successful at solving the test tasks, and are comparable to the Oracle in Push + Stack. The MIP fails to achieve nontrivial success rates both test environments. Error bars indicate 95% Cls over seeds.
interactions between more than one entity during training. Despite its relative simplicity, the Deep Set architecture extrapolates as well as or better than the relational architectures most environments. crucial exception is in 3-Push with cube-cube collisions enabled (Appendix D.2). There, modeling entity-entity interactions is especially crucial and a relational method like Self Attention largely superior. Overall, find that geometric architectures can perform very effective extrapolation.
3.3 ZERO-SHIE SHOT STITCHING TO SOLVE NOVEL TASKS
When evaluating policies for stitching behavior, we use test tasks that combine subtasks from training in novel ways. In our first setting, we train a policy on 2-Push and 2-Switch, and then test this policy on 2-Switch 2-Push, which requires both pushing cubes and flipping switches. our second setting, we train single policy on 2-Push and Stack, which requires stacking one cube on top of another. The test environment Push + Stack, which requires pushing one cube into position and then stacking the other block on top. This setting is especially difficult because it requires zero-shot stitching of skills particular order (push, then stack). Figure (left) shows the train-test task to test stitching.
relationships
Results and Observations Since this experiment requires training a single policy on multiple training tasks, during each episode choose one of the training tasks uniformly random. Figure 7 (right) shows that the MLP policy fails to jointly learn the training tasks in the first setting, leading to poor performance in 2-Switch + 2-Push. However, the MLP averages above 35% success rate both training tasks in the second setting, but still only manages a 5% success rate on Push + Stack. This suggests that when MLP policies are capable of learning the training tasks, they unable to combine them to solve new ones.
The geometric architectures show substantially better (but not oracle-level) stitching capabilities compared to the MLP. Graph Convolution struggles with the switch component of 2-Switch + 2-Push, but still outperforms the MLP. It is particularly surprising that the relational architectures (Self Attention and Graph Convolution) achieve 60% zero-shot success rate on Push + Stack, which requires understanding that the push and stack subtasks must be executed in specific order. Poor performance in 2-Switch + 2-Push is again due to difficulties in training one policy on two different tasks, suggesting that better joint training could further improve stitching performance.
4 RELATED WORK
Compositionality and Hierarchy. Hierarchical approaches to solving long-horizon tasks explicitly maintain or learn subpolicies corresponding to useful skills, which can then be coordinated by a high- level policy (Dayan and Hinton, 1993; Parr and Russell, 1998; Dietterich, 2000). Variations of this approach include learning termination policies for each sub-policy or "option" (Bacon et al., 2017), training the high level policy propose subgoals for goal-conditioned low level skills (Nachum et al., 2018), or even using natural language as the interface between high and low level policies (Jiang al., 2019). Our approach also enables learning of long-horizon and compositional tasks, but simply
8

---

Under review as a conference paper at ICLR 2023
through architectural modifications to the policy in end-to-end learning, opposed to explicitly learning action representations or modifying the training process.
Entity-centric modeling. Recent works in relational RL (Džeroski et al., 2001) have investigated graph neural networks (GNNs) (Gori et al., 2005; Scarselli et al., 2008) for handling complex multi- entity tasks, where the relational or message passing mechanism may be implemented using self attention (Zambaldi et al., 2018; Li et al., 2020) or a variety of other means (Bapst et al., 2019; Lin al., 2022). Though we directly focus on policy (and critic) architecture for model free learning, related approaches have studied entity structured networks for dynamics models (Carvalho et al., 2021; Veerapaneni et al., 2020; Sanchez-Gonzaler et al., 2018). Further work has explored relational architectures for extracting entity- representations from high dimensional observations before doing control (Wilson and Hermans, 2020; Driess et al., 2022). In particular, Zadaianchuk et al. (2020) combine entity-centric representation learning with goal conditioned RI. approach that also demonstrated extrapolation, though not stitching. In settings that can be modeled as EFMDPs, our framework formally motivates using GNNs through the perspective of permutation invariance (Bron- stein et al., 2021). Tang and Ha (2021) study the permutation invariance of self attention policies in particular, mainly in the context of robustness to input corruptions. But the invariance properties of our EFMDP framework also suggest considering invariant architectures without self attention or other relational mechanisms, in which case GNNs reduce to simpler architectures like Deep Sets (Zaheer et al., 2017). These architectures remain relatively underexplored outside of basic 21) environments (Karch et al., 2020). Our experiments evaluate both relational and Deep Set approaches on a suite of complex entity-centric robot tasks.
Policy Architectures in RL. MLPs, LSTMs, and small CNNs remain the dominant architectures in continuous control (Lillicrap et al., 2015; Schulman et al., 2017; Haarnoja et al., 2018). Sinha et al. (2020) study deeper networks for continuous control with DenseNet-style (Huang al., 2017) connections. Recent work has also explored the use of self attention over the trajectory history rather than between entities (Chen et al., 2021; Janner al., 2021). Other approaches leverage inductive biases about the real world, e.g. by embedding learnable dynamical systems into the policy architecture (Bahl et al., 2020).
5 CONCLUSION
This work introduces the EFMDP framework for the learning paradigm where an agent can interact with many entities in an environment. We explore how structural properties of EFMDPs induce permutation symmetry in the optimal policy and value functions, motivating policy architectures that leverage symmetry: set-based invariant models (Deep Sets) and relational models (Self Attention and Graph Convolution). These policy architectures decompose goal coal-connitioned tasks into their constituent entities and subgoals. These architectures are flexible, do not require any manual task annotations or action primitives, and can be trained end-to-end with standard RL or IL algorithms.
We compare these architecture types with each other and standard MLPs in suite of complex entity- centric tasks. We find that geometric architectures can: (a) learn substantially faster than black-box architectures like the MLP; (b) perform zero-shot extrapolation new environments with more of fewer entities than observed in training; and (c) perform zero-shot stitching of learned behaviors to solve novel task combinations never seen in training. We find that the geometric architectures perform relatively similarly across most tasks, which can be surprising given the Deep Set's 's relative simplicity. Since many existing entity-centric approaches focus on graph neural networks or transformers, our results invite further investigation into simple invariant architectures like Deep Sets.
Limitations and Future Work: EFMDPs require eptity-ppecific subgoals, but some tasks may instead be specified in terms of entity relations ("place this block on top of that block"). In such cases, the relational subgoals must be first converted into an equivalent entity specific form. Additionally, EFMDP's invariance properties explain why structured policies perform well, but do not distinguish between them (e.g., DS vs SA). We make this comparison empirically here, but future work could provide more principled framework for choosing the right architecture for given entity-centric task. We also hope to analyze the how the geometric properties of EFMDPs interact with object-centric representation learning, vibrant area of research (Burgess al., 2019; Kipf et al., 2019; Locatello et al., 2020; Nanbo et al., 2020) which is important to enabling compositional generalization for policies that operate on high dimensional observations like images.
9

---

Under review as a conference paper at ICLR 2023
Reproducibility: Appendix B describes the environments, and Appendix C describes the training details. We have already open-sourced the code implementing the environments and experiments, and will include the link after the anonymous reviewing phase ends.
10

---

Under review as a conference paper at ICLR 2023
REFERENCES
Leslie Pack Kaelbling. Learning to achieve goals. In IJCAI, pages 1094-1099. Citeseer, 1993.
Dhruv Batra, Angel X Chang, Sonia Chernova, Andrew Davison, Jia Deng, Vladlen Koltun, Sergey Levine, Jitendra Malik, Igor Mordatch, Roozbeh Mottaghi, et al. Rearrangement: A challenge for embodied ai. arXiv preprint 2XXi:20011.01975 2020.
Richard Sutton and Andrew Barto. Reinforcement Learning: An Introduction. MIT Press, 1998.
Dimitri Bertsekas and John Tsitsiklis. Neuro-dynamic Programming. 1996.
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.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Ruslan Salakhutdinov, and Alexander Smola. Deep sets. arXiv preprint arXiv:1703.06114, 2017.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pages 5998-6008, 2017.
Thomas N Kipf and Max Welling. Semi-superviss classification with graph convolutional networks. arXiv preprint arXiv: 1609::/609.0000 (00.600099077. 2016.
Tom Schaul, Daniel Horgan, Karol Gregor, and David Silver. Universal value function approximators. In International conference machine learning, pages 1312-1320. PMLR, 2015.
Tristan Karch, Cédric Colas, Laetitia Teodorescu, Clément Moulin-Frier, and Pierre-Yves Oudeyer. Deep sets for generalization in rl. arXiv preprint arXiv:2003.09443, 2020.
Saão Džeroski, Luc De Raedt, and Kurt Driessens. Relational reinforcement learning. Machine learning, 43(1):7-52, 2001.
Marco Gori, Gabriele Monfardini, and Franco Scarselli. A new model for learning in graph domains. In Proceedings. 2005 IEEE international joint conference on neural networks, volume 2, pages 729-734, 2005.
Franco Scarselli, Marco Gori, Ah Chung Tsoi, Markus Hagenbuchner, and Gabriele Monfardini. The graph neural network model. IEEE transactions on neural networks, 20(1):61-80, 2008
Vinicius Zambaldi, David Raposo, Adam Santoro, Victor Bapst, Yujia Li, Igor Babuschkin, Karl Tuyls, David Reichert, Timothy Lillicrap, Edward Lockhart, et al. Relational deep reinforcement learning, arXiv preprint arXiv: 01830, 2018.
Richard Li, Allan Jabri, Trevor Darrell, and Pulkit Agrawal. Towards practical multi-object manipu- lation using relational reinforcement learning. In 2020 IEEE International Conference on Robotics and Automation (ICRA), pages 4051-4058. IEEE, 2020.
Ofr Press, Noah A Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. arXiv preprint arXiv:2108. 12409, 2021.
Timothy Lillicrap, Jonathan Hunt, Alexander Pritzel, Nicolas Heess, Tom Frez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. arXiv preprint *-Xivv1509.02977 2015.
Marcin Andrychowicz. Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, Pieter Abbeel, and Wojciech Zaremba. Hindsight experience replay, arXiv preprint arXiv: 1007:0177009999 2017.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint 1777007777. 2017.
P Dayan and GE Hinton. Feudal reinforcement learning. nips (pp. 271-278), 1993.

---

Under review as a conference paper at ICLR 2023
Ronald Parr and Stuart Russell. Reinforcement learning with hierarchies of machines. Advances in neural information processing systems, pages 1043-1049, 1998.
Thomas G Dietterich. Hierarchical reinforcement learning with the maxq value function decomposi- tion. Journal of artificial intelligence research, 13:227-303, 2000.
Pierre-Luc Bacon, Jean Harb, and Doina Precup. The option-critic architecture. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 31, 2017.
Ofir Nachum, Shixiang Gu, Honglak Lee, and Sergey Levine. Data-efficient hierarchical reinforce- ment learning. arXiv preprint arXiv: 180500826, 2018.
Yiding Jiang, Shixiang Gu, Kevin Murphy, and Chelsea Finn. Language as an abstraction for hierarchical deep reinforcement learning. arXiv preprint arXiv:1906.07343; 2019.
Victor Bapst, Alvaro Sonzaez-Gonzalez, Carl Doersch, Kimberly Stachenfeld, Pushmeet Kohli, Peter Battaglia, and Jessica Hamrick. Structured agents for physical construction. In International Conference on Machine Learning, pages 4644774 PMLR, 2019.
Yixin Lin, Austin S Wang, Eric Undersander, and Akshara Rai. Efficient and interpretable robot manipulation with graph neural networks. IEEE Robotics and Automation Letters, 2022.
Wilka Carvalho, Anthony Liang. Kimin Lee, Sungryull Sohn, Honglak Lee, Richard Lewis, and Satinder Singh. Roma: A relational, object-model learning agent for sample-efficient reinforcement learning. 2021.
Rishi Veerapaneni, John Co-Reyes, Michael Chang, Michael Janner, Chelsea Finn, Jiajun Wu, Joshua Tenenbaum, and Sergey Levine. Entity abstraction in visual model- based reinforcement learning. In Conference on Robot Learning, pages 1439-1455 PMLR, 2020.
Alvaro Sanchez-Gonzalez Nicolas Heess, Jost Tobias Springenberg, Josh Merel, Martin Riedmiller, Raia Hadsell, and Peter Battaglia. Graph networks as learnable physics engines for inference and control. In International Conference on Machine Learning, pages 1470-44799 PMLR, 2018.
Matthew Wilson and Tucker Hermans. Learning manipulate object collections using grounded state representations. In Conference on Robot Learning, pages 490-502. PMLR, 2020.
Danny Driess, Zhiao Huang, Yunzhu Li, Russ Tedrake, and Mare Toussaint. Learning multi-object dynamics with compositional neural radiance fields. arXiv preprint arXiv: 1855, 2022.
Andrii Zadaianchuk, Maximilian Seitzer, and Georg Martius. Self-supervised visual reinforcement learning with object-centric representations. arXiv preprint arXiv:2011.14381, 2020.
Yujin Tang and David Ha. The sensory neuron transformer: cermutation-inaariii neural networks for reinforcement learning. Advances in Neural Information Processing Systems, 34, 2021.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning. pages 1861-1870 PMLR, 2018
Samarth Sinha, Homanga Bharadhwaj, Aravind Srinivas, and Animesh Garg, D2rl: Deep dense architectures in reinforcement learning. arXiv preprint arXiv:2010.09163 2020.
Gao Huang, Zhuang Liu, Laurens Van Der Maaten, and Kilian Weinberger. Densely connected convolutional networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 4700-4708, 2017.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. arXiv preprint arXiv:2106.01345. 2021.
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. Advances in Neural Information Processing Systems, 34, 2021.
12

---

Under review as a conference paper at ICLR 2023
Shikhar Bahl, Mustafa Mukadam, Abhinav Gupta, and Deepak Pathak. Neural dynamic policies for end-to-end sensorimotor learning. Advances in Neural Information Processing Systems, 33: 5058-5069, 2020.
Christopher P Burgess, Loic Matthey, Nicholas Watters, Rishabh Kabra, Irina Higgins, Matt Botvinick, and Alexander Lerchner. Monet: Unsupervised scene decomposition and representation. arXiv preprint arXiv: 11390, 2019.
Thomas Kipf, Elise van der Pol, and Max Welling. Contrastive learning of structured world models. arXiv preprint arXiv: (et.::::991.127.... 2019.
Francesco Locatello, Dirk Weissenborn, Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. Object-centric learning with slot attention. arXiv preprint arXiv:2006. /5055, 2020.
Li Nanbo, Cian Eastwood, and Robert B Fisher. Learning object-centric representations of multi- object scenes from multiple views. In 34th Conference on Neural Information Processing Systems, 2020.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojeiech Zaremba. Openai gym, 2016.
John Banister Lanier. Curiosity-driven multi-criteria hindsight experience replay. University of California, Irvine, 2019.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 14Xiv:1412.6640 69800, 2014.
Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dormann. Stable-basclines: Reliable reinforcement learning implementations. Journal of Machine Learning Research, 1-2, 2021. URL http:/ htp::/fme //jrrr.....rrrr//////// oww.orrg/////.. wrggpaperss/222 10-1364-0 html.
13

---

Under review as a conference paper at ICLR 2023
A PERMUTATION INVARIANCE
We recall Proposition 1:
Proposition (Policy and Value Invariance). In any EFMDP with entities, any optimal policy π² S G + A and optimal action-value function Q" 8 A x G + IR are both invariant to permutations of the entity-subgoal pairs. That is, for any a E SN:
x (as, og) = x -s, g) and (as, a, ag (s,
We want to show that any optimal policy .t S g + A and the optimal action-valus function Q* : x A x G + R are both permutation invariant, that is for any a E Sx:
(as, og) (s.
('s, a, og) as.
(10)
Recall that in an EFMDP the reward and dynamics have permutation symmetry (Property 1):
R(s, a. g) = R a, og) P(s' js a) = P(os'(oss
where as and og are defined in Eq. 4. We assume for simplicity that the agent space U and entity space& & are discrete, so that the state space S = u x e is also discrete.
We begin with Q*, which can be obtained by value iteration, where Q denotes the k'th iterate. We initialize Q7 0, which is (trivially) permutation invariant. Permutation invariance is then preserved during each step of value iteration Qi Q2+1.
@7+1 a, og) = R(os, a og) +ymax max a'
Σ a)Qi
(11)
s'ES
(R(s, y max
Plo s'o, 11)00
(12)
EES
R(s, a. + max a'
P(s'ls, P(s'(s,a)Qtt 21((s,))
(13)
S'ES
= 1
(14)
Hence Q permutation invariant for all i 0, with Qi k-0
Q'. Line 12 follows from
the permutation invariance of the reward, transition probability, and the previous iterate Q Line 13 uses the fact that summing over a for all s' E S is the same as simply summing over all states E S. This can be seen more explicitly by expanding a sum over arbitrary function f(-):
=)
Σ f(, (c-1(())
Σ f())
SES
ErEE
eNEE
SES
The permutation invariance of Q' leads the permutation invariance of n²:
π (as, og) arg max Q" (as, a, og) arg max Q" (s, a, g) π (s, g)
14

---

Under review as a conference paper at ICLR 2023
B ENVIRONMENTS
Our environments are modified from OpenAl Gym's Fetch environments (Brockman et al., 2016) (MIT license), with our stacking environment in partieular being modified from the Fetch stacking environments of Lanier (2019). They have a 4D continuous action space with 3 values for end effector displacement and value for controlling the distance between the gripper fingers. The final action disabled when the neither the training or test tasks involve stacking, since gripping is not required for block pushing or switch flipping. Input actions are scaled and bounded to be between [-1, 1]. We set the environment episode length based on the number of entities and subtasks involved. Each switch added timesteps, and each cube pushing or stacking task added 50 timesteps. For example, 2-S 2-Swirch 2-Push had max episode length of 2 X 50 + 2 20 140 timesteps.
For non-stacking settings such as N-Push and N-Switch + N-Push, we disable cube-cube collision physics to make training easier for all methods. Note that subgoals may still interfere with each other since the gripper can interact with all cubes, so the agent may accidentally knock one cube away when manipulating another one. We repeat the extrapolation experiments for N-Push with collisions in Appendix D.2.
State and goals. The agent state describe the robot's end effector position and velocity the gripper finger's 's positions and velocities. The entity state for cubes include the cube's pose and velocity, and for switches include the switch setting 0 E [-0.7, 0.7] and the position of the switch base on the table. The switch entity state is padded with zeros to match the shape of the cube entity state, and all entity states include an extra bit to distinguish cubes from switches. Subgoals specify target position for cubes and target setting 0* {-0.7,0.1 0.7] for switches.
The dense reward is defined the average distance between each entity and its desired state as specified by the subgoal. For cubes, this is the L2 distance between current and desired position. For switches, this setting. The sparse reward is O if all entities are within a threshold distance of their subgoals, and -1 otherwise.
Reward.
f"), where 9 is the current angle of the switch and 0* is the desired
15

---

Under review as a conference paper at ICLR 2023
Table General shared RL hyperparameters
Hyperparameter
Value
Discount Y
0.98
Parallel envs
Replay buffer size
10"
Relabel prob
0.8
Ratio of episodes : updates
2:5 5
Optimizer
Adam
Learning rate
MLP, Deep Set: 0.001
Self Attention: 0.0001
Reward Scale
Sparse: 1; Dense: 5
Action noise 1Jo (initial)
0.2
Random action prob eo (initial)
0.3
C TRAINING DETAILS
C.I REINFORCEMENT LEARNING
We train RL agents using a publicly available (MIT license) implementation" of DDPG (Lillicrap al., 2015) and Hindsight Experience Replay (HER) (Andrychowicz al., 2017). Table 1 contains the default hyperparameters shared across all experiments. Our modified implementation collects experience from 16 environments in parallel into single replay buffer, and trains the policy and critic networks on single GPU. We used an internal cluster to parallelize experimentation across multiple random seeds and algorihmsshyperraaaeeesss and for HER we relabel the goals in 80% of sampled minibatches (the "relabel prob"). The reward scale is simply a multipler of the collected reward used during DDPG training. For exploration we use action noise 1) and random action probability e; the output action is:
We collect 2 episodes for every 5 gradient updates,
a (N 7). ā
with prob 1
(Iniform-11,11) with prob
Table 2 shows environment specific RL. hyperparameters. "Epochs" describes the total amount of RL. training done, with 1 epoch corresponding to 50 x PARALLEL ENVS episodes. Sparse reward is used for the simpler environments, and dense reward for the harder ones. For some environments we decay the exploration parameters 7), e by ratio computed per-epoch. Lin(.01, 100, 150) means that v,, are both decayed linearly from TJu and Ey to .01 X T/ and .01 Eg between epochs 100 and 150. The constant exploration decay schedule maintains the initial na. values throughout training. The target network parameters are updated as giarget + (1 - T) + grarget where 7 is the target update speed.
We use the same RL hyperparameters regardless of architecture type except that the learning rate is lower for Self Attention and the exploration decay schedule may vary. Where Table lists "Fast" and "Slow" decay schedules, works best. In each case, the Self Attention policy prefers the slower exploration schedule and Deep Sets prefers the faster one, while the MLP typically fails learn with either exploration schedule on the more complex environments.
sweep over both options for each architecture and use the schedule that
Architectures. The exact actor and critic architectures uses for each architecture family is shown in Table 3. Linear(256) represents an affine layer with 256 output units. ReLU activations follow every layer except the last. The final actor layer is followed by a Tanh nonlinearity, and the critic has no activation function after the final layer. A represents the action space dimension, and Block(N, M, H) represents a Transformer encoder block (Vaswani et al., 2017) with embedding size N, feedforward dimension M, and H heads. We disable dropout within the Transformer blocks for RL training.
http:///gthhhb..comm//mmmonoaaaa / nidddiiiii wew.gghteeeeeeeeeeeeee
16

---

Under review as a conference paper at ICLR 2023
Table 2. Environment specific RL hyperprrameters
Environment
Reward
Epochs
Exploration decay
Target update speed
1-Push
Sparse
50
Constant(
0.95
2-Push
Dense
150
Lin(.01, 75, 125)
0.99
3-Push
Dense
250
Fast: Lin(.O .01, 30, 80)
0.99
Slow:
Lin(.01, 100. 175)
[1,2.3)-Switekl
Sparse
[10, 50, 100)
Constant( I)
0.95
1-Switch I-Push
Dense
150
Lin(.01, 60, 100)
0.99
2-Switch + 2-Push
Dense
250
Fast:
Lin(.01, 75, 150)
0.99
Slow: Lin(.01, 100, 150)
Table 3. RI. architectures
Family
Actor
Critic
MLP
Lincar(256)x3, Lincar(A)
Linear(256) x3, Linear(1)
Deep Set
d: Linear(256) x3
di Linear(256) x2
Linear(A)
p: Linear(256). Linear(1)
Self Attention
SA: Linner(256),. Block(256, 256, 4)x2 p: Linear(A)
SA: Linear(256), Block(256, 256, 4)x2 p: Linear(1)
Graph Convolution
Linear(256) x3, GraphConn(256) x2 pt Linear(A)
Linear(256) x3, (raphCCnn(256)) x2
p: Linear(1)
Table 4. Parameter count (3-Push)
parameters
Actor
Critic
MLP
150,020
150,273
Deep Set
140,292
140,545
Self Attention
800,260
800,513
Graph Convolution
271,876
272,129
17

---

Under review as a conference paper at ICLR 2023
C.2 IMITATION LEARNING
The IL dataset is generated using the best performing RL agent in that environment-we record M E {1000, 2000, 3000, 4000, 5000) demonstration trajectories. This creates dataset of M x T transitions
{S.. a.) )=1 MMT for behavior cloning. However, in practice we filter the dataset slightly by discarding the transitions corresponding to trajectories that are not successful.
We use the same policy architectures shown in Table and optimize mean squared error loss over the dataset:
arg min J(m) = DD
|| all² (s,a)~
Σ
We use the Adam (Kingma and Ba, 2014) optimizer with learning rate 0.001 (MLP, Deep Sets) 0.0001 (Self Attention). Each policy is trained for 60, 000 gradient steps with a batch size of 128.
C.3 TRAINING AND INFERENCE SPEED
Here we consider the computational complexity of using different architecture classes (MLPs, Deep Sets, and Self Attention), as we scale the number of entities N. We consider the number of parameters, activation memory, and computation time (for forward pass). For MLPs with fixed hidden layer sizes, the number of parameters and computation time increase linearly with N while the memory required for activations stays fixed (due to fixed hidden layer sizes). In Deep Sets and Self Attention, the number of parameters does not depend on the number of entities N. The activation memory and computation time grow linearly in Deep Sets, and quadratically for the pairwise interactions of Self Attention. In practice, the number of entities N is modest in all our environments (e.g., fewer than 10), but computational complexity may be relevant in more complex scenes with lots of entities.
For more holistic real- world comparison of execution and training speed, Figure 8 shows both inference time and training time in the N-Push environments for N E (1, 2, 3). The inference time the number of milliseconds it takes actor do single forward pass (using GPU) a single input observation. The Self Attention policy involves more complex computations and is significantly slower than Deep Set and MLP policies. The RL training time is the actual number of hours required to run the reinforcement learning algorithms of Figure 5, for each architecture. Not surprisingly, we see that 3-Push takes significantly longer train than I-Push, since it harder environment. For a fixed environment, however, all three architecture types are comparable in speed, with the Self Attention version being slightly slower than the others. The surprising similarity in RL training time (despite much slower inference time for the Self Attention policy) suggests that most of the RL time is spent on environment simulation rather than policy or critic execution. Hence, the difference between architectures presented in this paper has only a minor effect reinforcement learning speeds in practice.
Inference and training speed N-Push
Actor inference time obsorvation)
RL Training Time
Milliseconds
1.0
Hours bn
Deep Set MLP
0.5
0.0
3
Number of cubes
Number of cubes
Figure 8. Left: the time (in milliseconds) it takes for each policy architecture to execute single forward pass on a single input observation from the N-Push environments, where E (1,2,3). The self attention policy significantly slower, while the Deep Set and MLP policies are comparable. Right: Real world reinforcement learning times (in hours) training each policy/critic architecture on the N-Push environments. Although the Self Attention policy is slightly slower, all policies train at comparable speeds in the same environment. This suggests that environment simulation, not policy execution, the dominant time consuming element.
18

---

Under review as a conference paper at ICLR 2023
D FURTHER COMPARISONS
D.1
DEEP SET ARCHITECTURE SIZE
Effect of Deep Set size on Extrapolation Stitching
N-Push
Push
Stack
Swoness Rate
1.0
Suoness Rate
0.5
Deep Set Sarpel Truntig
0.0
0.0
2.Push
Stack
Push Stack
cubes
Figure Comparison of N-Push extrapolation and Push + Stack stitching performance when using small and large variants of the Deep Set policy architecture. The small version implements p with l-layer linear map, while the large version implements p with 2-layer MLP. For N-Push, the larger network achieves greater success rates in the training environment (3 cubes) but is actually worse when extrapolating or cubes. On the other hand, the larger Deep Set displays superior stitching capability and achieves higher average success rate when generalizing to Push + Stack from 2-Push and Stack.
Recall that our Deep Set policy architecture involves two MLPs and p. where & produces inter- mediate representations for each entity, those intermediate representations are summed, and then p produces the final output (Eq. 6). In full generality, both ø and p may have two more layers with nonlinearities between. While our & is a 3-layer MLP, we use a linear throughout the main paper because found that it often works comparably or better than using larger 2-layer MLP p. Here we repeat the N-Push extrapolation and Push + Stack stitching experiments from the main paper using 2-layer p, which call "Deep Set (large)." The results from the main paper uses I-layer p which we refer to here as "Deep Set (small)."
Figure 9 shows the results. In N-Push, the larger Deep Set model achieves higher training success rates the 3-cube environment, but has worse extrapolation success rates for large numbers of cubes. For example, the smaller Deep Set model is significantly better at solving 6- 6-Push. Meanwhile, the large and small Deep Sets achieve very similar results in the pushing and stacking training environments. However, the larger Deep Set model achieves higher success rate the Push Stack environment, indicating superior stitching capability. This suggests that simpler Deep Set architectures may be better for extrapolating large number of entities, but more complex architectures may be superior for solving complex tasks with fixed number of entities.
D.2 N-PUSH WITH CUBE-CUBE COLLISIONS
N-Push Extrapolation with Collisions
Success Rate
1.0
0.5
0.0
cubes
Self Attention
Deep Set
MLP
Oracle
Figure 10. N-Push extrapolation with cube-cube collisions enabled. All methods observe some drop in performance relative to Figure 6, where N-Push has cube-cube collisions disabled. Self Attention tends to outperform Deep Sets when collisions enabled, likely because its relational inductive biases better suited to handling interactions between entities that arise from collisions.
As noted Appendix B, we disable cube-cube collisions in the N-Push and N-Switch+N--Pun experiments of the main paper (of course, the stacking settings require cube-cube collisions to be enabled). Here we repeat the N-Push extrapolation experiments with cube-cube collisions enabled. Figure 10 shows the results, which are gualitatively similar to when collisions disabled. All methods observe a decrease success rates of about 15%, with the Self Attention method often outperforming the Deep Set policy. This is likely because N-Push involves interaction between
19

---

Under review as a conference paper at ICLR 2023
entities once cube-cube collisions enabled, and Self Attention's 's relational inductive biases better suited for modeling these interactions.
D.3 N-PUSH WITH PPO INSTEAD OF DDPG
Our RL learning efficiency experiments all DDPG with Hindsight Experience Replay, which raises the question of whether the results (in particular, the inefficiency of the MI.P) are specific to particular training algorithm. To answer this, we repeated the 3-Push RI. experiment for the Deep Set and MLP architectures, but trained with Proximal Policy Optimization (PPO (Schulman al., 2017)) instead of DDPG. We use the Stable Baselines 3 (Raffin et al., 2021) (MIT License) implementation of PPO with default hyperparameters, trained for up to 1 x 10 environment steps. Figure 11 shows that the general trend we observed while using DDPG still holds: the invariant Deep Set policy learns the task far more efficiently than the MLP. In fact, as with DDPG the MLP policy fails learn 3-Push at all.
3-Switch PPO
Return
-10
-20
0
2
Step
1e6
Deep Set
MLP
Figure 11. Return throughout RL train- ing on 3-Push using Proximal Policy Op- timization (PPO) instead of DDPG.
20


Reviews:
[{"review_id": "xjS3bfn1gvb", "reviewer": "Reviewer_KcU5", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "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 architecture that leverages the symmetry of entity-centric goal specification. The author extends the MDP framework to modeling the symmetric structure of the goal specification. The framework motivates using architectures that are agnostic to permutation of the inputs like Deep Set, Self-Attention, and Graph convolution. Experiments are conducted in two robotic tasks, showing the advantages of the proposed architectures over the MLP architecture. \n\n", "summary_of_the_review": "The paper applies recent state-of-the-art architectures to a special class of robotic tasks. The author provides a theoretically-motivated justification but I don't find it providing significant value to the community, because the authors do not provide analyses to support the justification. I am leaning towards rejecting the paper.  \n\n=====After Rebuttal======\n\nI have read the responses from the authors and decided to keep my current score. My main ground for rejection is that the proposed framework lacks technical depth and it does not lead to novel methods, surprising findings, or useful insights. The framework justifies using architectures that are permutation-invariant for problems that require that property, but such a choice seems to be quite natural even without considering the framework. In addition, it is unclear from the experimental results whether the data provided to the agent enforces permutation invariance and the learned policies are verified to have actually learned that property. Showing improved performance on various task conditions is not sufficient; there could be various other properties of the experimented architectures that affect their performance. More ablation studies need to be conducted to identify the cause of the improvements and to demonstrate that the actual cause coincides with the anticipated cause. ", "strength_and_weaknesses": "Strengths:\n* The paper is clearly written. \n* The problem studied is interesting and important.\n* Results quantitatively show the effectiveness of the proposed method. \n\nWeaknesses:\n* The proposed architectures are not novel (according to the authors, the application of these architectures on the experimented robotic tasks is novel; I don't work on this space so I am not qualified to verify this claim). \n* The authors show mostly evaluation metrics without analyzing the behavior of the learned policies. How do the learned policies solve the entity tasks? Sequentially or simultaneously? Visualizing the change in the attention over the goals may offer such insights. \n* I am not sure why the \"oracle\" is called an \"oracle\" rather than a rule-based baseline. The results would be interpreted much differently if the \"oracle\" were considered as a baseline. \n* I am not convinced why the approach would work better than learning a policy separately for each entity. The authors mention that entity-entity relationships may not be considered in such an approach, but I can easily enforce some constraints to the policies through rewards or demonstrations (e.g. don't push other blocks while moving a block). The proposed method may resort to learning the \"oracle\" policy where it only attends to one goal at a time. Analyzing the behavior of the learned policy is important to show that the method learns a non-trivial policies. The current numerical results show that most of the time the method is on par with the \"oracle\" baseline, which hints that it may have learned trivial policies.  \n* The proposed architecture can outperform MLP because it simply has more learning capacity. The authors should show the number of parameters of the models in the comparisons. \n* When training the MLP, do you permute the input entity subgoals in each episode or do you always feed the subgoals in the same order? Doesn't permuting the subgoal order make the MLP somewhat invariant to permutation?\n* The problem studied, although interesting, is of limited practicality. In the real life, the decomposition into entity goal is usually inferred rather than given explicitly to the policy as in this work (e.g. \"switch all light off\" rather than \"switch light 1 off, switch light 2 off, ...\"). I understand that the authors are studying a sub-problem of this more difficult setting, but I just want to point out a limitation of the current task formulation. \n* The conclusion says \"These policy architectures decompose goal-conditioned tasks into their constituent entities and subgoals\" but this seems like a false claim, as the goal specification already consists of entity subgoals stitched together. Figure 3 illustrating the architectures also does not show any goal decomposition.  Am I missing something?", "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": "hAH1UTq-Ck2", "reviewer": "Reviewer_HR7M", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "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": "This paper proposes a framework for compositional structure modeling in control tasks, that is a modification of the typical Markov Decision Process (MDP). The framework, called Entity-Factored Markov Decision Process provides several insights into developing robot control policy architectures for compositional generalization. Results on simulated robot manipulation tasks show extrapolation to different number of test entities, and stitching of skills in novel ways. ", "summary_of_the_review": "I think the paper has interesting insights and opens up possibilities for future works to build upon, and I am leaning towards acceptance. Please look at the details of strengths/weaknesses above. \n\n------ AFTER AUTHOR RESPONSES ----\n\nThe authors have updated the limitations section of the paper to include more details, so the current paper is better scoped in terms of its contributions. My second point regarding comparisons to baselines is not adequately addressed and would likely require significant additional work: I was referring to comparisons with other goal-conditioned RL approaches (not just RL + HER) for example approaches described in this paper [A]. Currently, the baselines are not \"external\" goal-conditioned RL approaches in the sense that they do not correspond directly to a prior paper. \n\nIn addition, after going through the other reviews, some concerns regarding proper ablations to tease out the benefits in performance are unaddressed, as overall performance is not indicative of whether confounding factors were at play. \n\nBased on these, I am not able to strongly recommend accept, but would keep my rating of weak accept with lower confidence as I think the paper is interesting for future work to build on\n\n[A] Liu, Minghuan, Menghui Zhu, and Weinan Zhang. \"Goal-conditioned reinforcement learning: Problems and solutions.\" arXiv preprint arXiv:2201.08299 (2022).", "strength_and_weaknesses": "Strengths\n\n- The paper provides a formal study of understanding compositional tasks, which can be useful for several papers that target this domain, for developing future algorithms \n\n- Understanding deep RL policy architectures is generally not very well-studied, and this paper proposes a formal study of this, which can be useful for to several works in RL. \n\n- The experimental evaluations are very detailed and show the emergence of properties like extrapolation to out-of-domain targets, and stitching of skills for compositional generalization. \n\nWeaknesses\n\n- The limitations of the proposed framework are unclear. It will be good to be precise, clear, and transparent about exactly which king of control problems can be cast in this framework, and exactly when will the benefits be observed empirically, compared to casting the problem in the usual MDP framework, \n\n- No external baselines are compared against. In my understadning, a one-on-one comaprison with *any* RL algorithm that is goal-conditioned is possible, and should be performed so that the it is clear how the algorithms developed in this framework comapre to existing goal-conditioned RL algorithms under the MDP framework. ", "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": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "RuI1jStJOBg", "reviewer": "Reviewer_dhRF", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "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": "This paper presents a framework called Entity-Factorized Markov Decision Process (EFMDP), where the task is factorized across different entities, and the state and goal configurations of the entities are permutation-invariant. Based on this framework, the paper studies various permutation-invariant architecture designs, including Deep Sets, Graph Convolution, and Self Attention. Experiments on object pushing, switching, and stacking demonstrate that the method with permutation-invariant architectures learns faster than MLPs when there are multiple objects in the scene, and can do zero-shot extrapolation to an unseen number of objects.", "summary_of_the_review": "Overall, I think this is a solid work with good novelty and sufficient experiments. I believe the proposed EFMDP setting could be useful for many other problems as well. Therefore, I vote to accept this paper.\n\n\n--- After rebuttal ---\n\nI'm happy with the authors' response and decide to keep my score.", "strength_and_weaknesses": "**Strength:**\n\n- The paper presents a neat and intuitive idea to factorize the state and goal configurations across entities, which is beneficial to generalization and scaling.\n- Experiments are extensive, where multiple permutation-invariant architectures are studied in various different settings, e.g., zero-shot extrapolation and zero-shot stitching. The results demonstrate the effectiveness of the proposed method and empirically show how different architectures compare for these tasks.\n- The paper is well structured with clear problem definitions, and how different architectures fit in the big picture. The method is well-motivated and easy to follow.\n\n**Weakness:**\n\n- The method is only evaluated on robot manipulation, it would be nice to also evaluate on other settings such as multi-agent systems, and strategic game playing, which would make the experiments stronger.\n- While the paper shows different permutation-invariant architectures such as Deep Sets, Graph Convolution, and Self Attention can fit into EFMDP, the framework itself doesn’t provide any insights or guidelines as to which one is better in certain settings, which makes the theoretical analysis less useful.", "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": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "JJFeriOVLRH", "reviewer": "Reviewer_Et3S", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "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": "This paper introduced the Entity-Factored Markov Decision Process (EFMDP) for modeling the entity-based compositional structure in controlling tasks. The authors studied several structured policy architectures that can utilize the factorized discrete entities on a suite of manipulation tasks. Experimental results showed that structured policy architectures have faster learning speed in general when compared to MLP policy. Robust extrapolation and OOD generalization at the skill-level can also be observed.", "summary_of_the_review": "This paper introduced the Entity-Factored Markov Decision Process (EFMDP) for modeling the entity-based compositional structure in controlling tasks. Experiments showed that structured policy outperforms the MLP policy. However, similar results have been obtained from related work already.", "strength_and_weaknesses": "### Strength\n- The experiments are well designed, and the results support the author's statement well.\n\n### Weakness\n- Major:\n  - As the authors mentioned also, the current model takes vector state as input. However, a model that can handle visual input is more desired. A related work, SMORL[1], has shown promising results given visual input. Similar results can also be obtained from there.\n\n\n-------------After rebuttal-----------------------\n\nThanks for the author's responses. I decide to increase my score to 5. \n1)I still think the proposed work shares some common motivations with SMORL, such as the reduced size of effective state space. The authors did investigate more architecture than the SMORL work. However, the experiment's only finding is that utilizing the factorized entity-based structure improves the performance, which does not have much novelty. Similar findings have been demonstrated in prior work like SMORL. 2)Though given the ground-truth entity states and goal state as input, the proposed work is solving a much simpler task than SMORL, I'm convinced that the OOD experiments are meaningful, and I'd like to increase my score to 5. However, I am leaning toward rejection based on the current experimental results.\n\n\n[1]: Andrii Zadaianchuk, Maximilian Seitzer, and Georg Martius. Self-supervised visual reinforcement learning with object-centric representations.", "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": "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