Original: Hierarchical Abstraction for Combinatorial Generalization in Object Rearrangement
Revised: Hierarchical Abstraction for Combinatorial Generalization in Object Rearrangement

Original: Anonymous Author(s)

Original: AfﬁliationAddress email

Original: Abstract
Revised: Abstract

Original: Object rearrangement is a challenge for embodied agents because solving thesetasks requires generalizing across a combinatorially large set of underlying entitiesthat take the value of object states.
Revised: Object rearrangement is a challenge for embodied agents because solving these tasks requires generalizing across a combinatorially large set of underlying entities that take the value of object states.

Original: Worse, these entities are often unknown andmust be inferred from sensory percepts.
Revised: Worse, these entities are often unknown and must be inferred from sensory percepts.

Original: We present a hierarchical abstractionapproach to uncover these underlying entities and achieve combinatorial generalization from unstructured inputs.
Revised: We present a hierarchical abstraction approach to uncover these underlying entities and achieve combinatorial generalization from unstructured inputs.

Original: By constructing a factorized transition graphover clusters of object representations inferred from pixels, we show how to learna correspondence between intervening on states of entities in the agent’s modeland acting on objects in the environment.
Revised: By constructing a factorized transition graph over clusters of object representations inferred from pixels, we show how to learn a correspondence between intervening on states of entities in the agent’s model and acting on objects in the environment.

Original: We use this correspondence to developa method for control that generalizes to different numbers and conﬁgurations ofobjects, which outperforms current ofﬂine deep RL methods when evaluated on aset of simulated rearrangement and stacking tasks
Revised: We use this correspondence to develop a method for control that generalizes to different numbers and configurations of objects, which outperforms current offline deep RL methods when evaluated on a set of simulated rearrangement and stacking tasks

Original:  Introduction
Revised:  Introduction

Original: A core property of intelligence is the ability to re-purpose previously acquired knowledge for solvingnew problems, but how to identify symmetries for factorizing modeling and control into independentcomponents has been a challenge for AI systems, which either assume human-deﬁned abstractions tobegin with [11, 13] or learn monolithic representations with no clear mechanism for reuse [19, 30].
Revised: A core property of intelligence is the ability to re-purpose previously acquired knowledge for solving new problems, but how to identify symmetries for factorizing modeling and control into independent components has been a challenge for AI systems, which either assume human-defined abstractions to begin with [11, 13] or learn monolithic representations with no clear mechanism for reuse [19, 30].

Original: The problem of object rearrangement offers an intuitive setting for studying this problem of knowledge reuse.
Revised: The problem of object rearrangement offers an intuitive setting for studying this problem of knowledge reuse.

Original: Because the space of object conﬁgurations is combinatorially large, solving novelrearrangement problems requires the agent to recognize that the same action for moving an objectfrom one location to another can be reused for different objects in different contexts.
Revised: Because the space of object configurations is combinatorially large, solving novel rearrangement problems requires the agent to recognize that the same action for moving an object from one location to another can be reused for different objects in different contexts.

Original: We can formulatethis problem as an ofﬂine goal-conditioned reinforcement learning (RL) problem, where the agent istrained on a dataset of sensorimotor interactions and is evaluated on rearranging objects to satisfyconstraints depicted in a goal image.
Revised: We can formulate this problem as an offline goal-conditioned reinforcement learning (RL) problem, where the agent is trained on a dataset of sensorimotor interactions and is evaluated on rearranging objects to satisfy constraints depicted in a goal image.

Original: The research questions are, therefore, how can we enable theagent to represent abstractions of objects in a way that is amenable for such reuse in control, and howcan we enable it to discover these abstractions on its own directly from the sensorimotor interface?a
Revised: The research questions are, therefore, how can we enable the agent to represent abstractions of objects in a way that is amenable for such reuse in control, and how can we enable it to discover these abstractions on its own directly from the sensorimotor interface?

Original:  Hierarchical Abstraction (HA) that offers an  to the  which  call entities and factorize each entity h k into two parts its identity z k and its state s k 
Revised: Our main contribution is a method called Hierarchical Abstraction (HA) that offers an answer to the above questions. The key idea is to represent objects as independent and symmetric latents, which we call entities, and factorize each entity h k into two parts, its identity z k and its state s k .

Original: The identity is the part that is not affected by action and the state is the part is component of  a  is  to the identity of the entity or the  sensorimotor interaction  control component of  of  solve new rearrangement problems  inferring  be taken  only the current and goal image observations e) Action Selection ! !
Revised: The identity is the part that is not affected by action and the state is the part that is. The modeling component of our approach abstracts the dataset of sensorimotor interactions into a graph over state transitions of individual entities. By construction, a state transition is agnostic to the identity of the entity or the context entities present in the sensorimotor interaction from which it was learned, enabling it to be reused across multiple entities and contexts. The control component of our approach then re-composes sequences of state transitions to solve new rearrangement problems by inferring what state transition can be taken given only the current and goal image observations.

Original:  slot attention slot attention align isolate bind bind ! ! !
Revised: e) Action Selection 𝑎 𝑎 slot attention slot attention align isolate bind bind 𝑜 !

Original:  " c) Representation Learning tensor([0.0921, 0.4875, 0.5775, 0.9429])  !  ! " ! " !"
Revised: 𝑜 " c) Representation Learning tensor([0.0921, 0.4875, 0.5775, 0.9429]) 𝑎 ! 𝑜 ! 𝑜 !"

Original: # " slot attention transformer slot attention prior ! ! ! ! ! d) Graph Construction concrete transitions 
Revised: # 𝑝 slot attention transformer slot attention prior 𝑎 𝑎 𝑎 𝑎 𝑎′ 𝑎′ 𝑎 𝑎′ 𝑎′′ 𝑎′′ 𝑎′′ 𝑎′′ d) Graph Construction concrete transitions abstract

Original: [Figure]
Revised: [Figure]

Original:  # ℎ !
Revised: 𝑜 !" # ℎ !

Original: ℎ !"
Revised: ℎ !"

Original: $ # !$ & ! $ # !
Revised: # 𝑠 𝑠′ 𝑎 ! 𝑎 ! 𝑎 𝑧 !

Original:  $% observation model
Revised: $ 𝑠 ! $ 𝑧 ! $ 𝑠 ! $% observation model

Original: [Figure]
Revised: [Figure]

Original: sensorimotor
Revised: sensorimotor

Original: [Figure]
Revised: [Figure]

Original: [Figure] [Figure] [Figure] [Figure]
Revised: environment imagined transformation

Original: [Figure]
Revised: [Figure]

Original: [Figure][Figure] [Figure]
Revised: [Figure]

Original: [Figure] [Figure] 
Revised: [Figure] [Figure] [Figure]

Original: correspondence
Revised: correspondence

Original: [Figure]
Revised: [Figure]

Original: ) Correspondence Problem ! " #$
Revised: ) Correspondence Problem 𝑜 "#$

Original: [Figure]
Revised: [Figure]

Original:  In contrast to other works in object-centric learning [14–16, 21, 24, 31, 34, 40] which evaluate thequality of inferred entities via segmentation metrics, we evaluate our inferred entities by how wellthey can reused for solving tasks , which more directly assesses what we want an agent’s objectrepresentations to do.
Revised: 𝑘 [Figure] In contrast to other works in object-centric learning [14–16, 21, 24, 31, 34, 40] which evaluate the quality of inferred entities via segmentation metrics, we evaluate our inferred entities by how well they can reused for solving tasks , which more directly assesses what we want an agent’s object representations to do.

Original: Kulkarni et al.
Revised: Kulkarni et al.

Original: [23], Veerapaneni et al.
Revised: [23], Veerapaneni et al.

Original: [36] also evaluate on control tasks.
Revised: [36] also evaluate on control tasks.

Original: Whereas Kulkarni et al.
Revised: Whereas Kulkarni et al.

Original: [23] considers how object representations improve exploration, we considerthe ofﬂine setting which requires zero-shot generalization.
Revised: [23] considers how object representations improve exploration, we consider the offline setting which requires zero-shot generalization.

Original: The shooting-based planning methodin Veerapaneni et al.
Revised: The shooting-based planning method in Veerapaneni et al.

Original: [36] suffers from compounding errors as other learned single-step modelsdo [20], but our non-parametric approach of composing previously seen transitions enables us to planfor longer horizons.
Revised: [36] suffers from compounding errors as other learned single-step models do [20], but our non-parametric approach of composing previously seen transitions enables us to plan for longer horizons.

Original: Indeed, our results show that HA outperforms both state-of-the-art ofﬂine RLmethods and slot-based planning methods across three simulated object rearrangement problems.
Revised: Indeed, our results show that HA outperforms both state-of-the-art offline RL methods and slot-based planning methods across three simulated object rearrangement problems.

Original:  Hierarchical Abstraction
Revised:  Hierarchical Abstraction

Original: We present our method for solving object rearrangement problems from pixels.
Revised: We present our method for solving object rearrangement problems from pixels.

Original: Given a transitiono, a ! o 0 from the training set in which a single object k has been moved, the modeling problem isto represent the state transition s k ! s k 0 as decoupled from the identity z k of the affected entity andas decoupled from the other context entities h [-
-]
Revised: Given a transition o, a → o ′ from the training set in which a single object k has been moved, the modeling problem is to represent the state transition s k → s k ′ as decoupled from the identity z k of the affected entity and as decoupled from the other context entities h [-
-]̸

Original: = k in the scene.
Revised: = k in the scene.

Original: The control problem is to composepreviously seen state transitions to satisfy goal constraints on new object conﬁgurations.
Revised: The control problem is to compose previously seen state transitions to satisfy goal constraints on new object configurations.

Original: Though wedo not assume the number of action primitives is ﬁnite, we assume the number of observed locationscan be clustered into a set of ﬁnite groups.
Revised: Though we do not assume the number of action primitives is finite, we assume the number of observed locations can be clustered into a set of finite groups.

Original: 2.1 Modeling
Revised: 2.1 Modeling

Original: Our modeling approach abstracts the training dataset of action-conditioned videos into a factorizedgraph over state transitions that can be reused across different rearrangement problems.
Revised: Our modeling approach abstracts the training dataset of action-conditioned videos into a factorized graph over state transitions that can be reused across different rearrangement problems.

Original: Constructingsuch a graph decomposes into two representation learning problems: the ﬁrst concerns how torepresent objects in a way that exposes their commonalities and the second concerns how to representactions in a way that enables them to be reused across these commonalities.
Revised: Constructing such a graph decomposes into two representation learning problems: the first concerns how to represent objects in a way that exposes their commonalities and the second concerns how to represent actions in a way that enables them to be reused across these commonalities.

Original: Our key insight is thatthese two problems are both different instances of the same fundamental problem of partitioning aspace into discrete groups, but at different levels of abstraction: learning how to cluster visual featuresinto entities (Fig. 1b) and learning how to cluster transitions over entity-sets into transitions overindividual states (Fig. 1c).
Revised: Our key insight is that these two problems are both different instances of the same fundamental problem of partitioning a space into discrete groups, but at different levels of abstraction: learning how to cluster visual features into entities (Fig. 1b) and learning how to cluster transitions over entity-sets into transitions over individual states (Fig. 1c).

Original: Our model thus implements a two-level abstraction hierarchy (Fig. 
Revised: Our model thus implements a two-level abstraction hierarchy (Fig. 1a). Level 1: abstracting visual features into sets of entities

Revised: The goal of the first level is to map a transition over two video frames o 1 , a 1 → o 2 into a transition over entity-sets ( h 11 , ..., h K 1 )

Original: 1a) different objects in the scene, i.e. they are independent and symmetric and (2) only the state s k ofentity k changes, not its identity z k , i.e. the entity is factorized .
Revised: → ( h 12 , ..., h K 2 ) (Fig. 1c). The criteria our entity-set transitions should satisfy are (1) entities represent different objects in the scene, i.e. they are independent and symmetric and (2) only the state s k of entity k changes, not its identity z k , i.e. the entity is factorized .

Original: Because each video transition depictsonly a single object being moved, we also want (3) the transition to be sparse , i.e. only one entity h kchanges in the transition and the rest h 6 = k remain unchanged.
Revised: Because each video transition depicts only a single object being moved, we also want (3) the transition to be sparse , i.e. only one entity h k changes in the transition and the rest h ̸ = k remain unchanged.

Original: Our approach manually enforces independence, symmetry, and factorization but lets sparsity emerge.
Revised: Our approach manually enforces independence, symmetry, and factorization but lets sparsity emerge.

Original: Speciﬁcally, we formulate the inference of entity-sets as the application of an equivariant actionconditioned sequential Bayesian ﬁlter with a mixture model as the latent state, where entity representations are the parameters of the mixture components.
Revised: Specifically, we formulate the inference of entity-sets as the application of an equivariant actionconditioned sequential Bayesian filter with a mixture model as the latent state, where entity representations are the parameters of the mixture components.

Original: Given the connection between the slots producedby slot attention (SA) [25] and mixture components Chang et al.
Revised: Given the connection between the slots produced by slot attention (SA) [25] and mixture components Chang et al.

Original: [6], we concretely implement thisﬁlter by temporally evolving the slots produced by the slot attention module of the state-of-the-art
Revised: [6], we concretely implement this filter by temporally evolving the slots produced by the slot attention module of the state-of-the-art

Original: SLATE architecture [32] with a transformer decoder (TD) [28, 35].
Revised: SLATE architecture [32] with a transformer decoder (TD) [28, 35].

Original: With h = ( h 1 , ..., h K ) ,h t +1 =
Revised: With h = ( h 1 , ..., h K ) , h t +1 =

Original: SA (ˆ h t , o t +1 )[Equation]
Revised: SA (ˆ h t , o t +1 )

Original:  queries[Equation] keys [Equation]
Revised: ˆ s t +1 = TD ( queries

Original: values [Equation] [ s [Equation]
Revised: = s t , keys/values =

Original:  a [Equation] ]) , where ˆ h t +1 =
Revised: [ s t , a t ]) , where ˆ h t +1 =

Original: [ z t , ˆ s t +1 ] and ˆ h 0 are independent samples from a Gaussian.
Revised: [ z t , ˆ s t +1 ] and ˆ h 0 are independent samples from a Gaussian.

Original: This approach enforcesindependence because ˆ h 0 are independent, symmetry because the transformer is equivariant, andfactorization because ˆ z t +1 is a copied from z t .
Revised: This approach enforces independence because ˆ h 0 are independent, symmetry because the transformer is equivariant, and factorization because ˆ z t +1 is a copied from z t .

Original: Empirically, we ﬁnd that this architecture modelssparse changes in the slots (Fig. 3).
Revised: Empirically, we find that this architecture models sparse changes in the slots (Fig. 3).

Original: We call this implementation dynamic SLATE (dSLATE), butour contribution of how structurally we enforce the above three criteria remains invariant to how thesequential Bayesian ﬁlter is implemented, which we expect will advance beyond SA and TD.
Revised: We call this implementation dynamic SLATE (dSLATE), but our contribution of how structurally we enforce the above three criteria remains invariant to how the sequential Bayesian filter is implemented, which we expect will advance beyond SA and TD.

Original: Level 2: abstracting transitions over sets of entities into individual state transitions
Revised: Level 2: abstracting transitions over sets of entities into individual state transitions

Original: The goalof the second level is to construct a state transition graph from a buffer of entity-set transitions{ ( h 1 , ... h T ) }
Revised: The goal of the second level is to construct a state transition graph from a buffer of entity-set transitions { ( h 1 , ... h T ) }

Original: Nn =1 produced from the ﬁrst level (Alg. 1, Fig. 1d).
Revised: Nn =1 produced from the first level (Alg. 1, Fig. 1d).

Original: Speciﬁcally, we seek each edgeof this graph to represent a state transition that has been shared across multiple entities in differentcontexts in the training dataset, such that we can reuse such state transitions when choosing actionsfor solving new rearrangement problems.
Revised: Specifically, we seek each edge of this graph to represent a state transition that has been shared across multiple entities in different contexts in the training dataset, such that we can reuse such state transitions when choosing actions for solving new rearrangement problems.

Original: Given our assumption that observed locations can begrouped into a ﬁnite set of clusters, this goal translates into a problem of clustering state transitions,and the problem of clustering state transitions reduces to clustering the states before and after thetransition.
Revised: Given our assumption that observed locations can be grouped into a finite set of clusters, this goal translates into a problem of clustering state transitions, and the problem of clustering state transitions reduces to clustering the states before and after the transition.

Original: We treat each cluster centroid as a node in the state transition graph, and an edge betweennodes is tagged with the single action that transforms one node’s state to another’s.[Table] [Table] [Table]
Revised: We treat each cluster centroid as a node in the state transition graph, and an edge between nodes is tagged with the single action that transforms one node’s state to another’s.

Original: [Table][Table] [Table] [Table]
Revised: [Table]

Original: [Table] j ] ⇤ ⌘ To create the nodes, the ﬁrststep is to identify the entity h kthat dSLATE predicted was affected by a t in each transition( h t , a t , h t +1 ) .
Revised: To create the nodes, the first step is to identify the entity h k that dSLATE predicted was affected by a t in each transition ( h t , a t , h t +1 ) .

Original: Our isolate procedure achieve this by solving
Revised: Our isolate procedure achieve this by solving

Original: [Equation]
Revised: [Equation]

Original: to identify the index of the entitywhose state has most changedduring the transition, where d ( · , · )104 is a distance function, detailedin the Appendix.
Revised: to identify the index of the entity whose state has most changed during the transition, where d ( · , · ) is a distance function, detailed in the Appendix.

Original: Having thebuffer of transitions over entitysets { ( h 1 , ... h T ) }
Revised: Having the buffer of transitions over entity sets { ( h 1 , ... h T ) }

Original: Nn =1 to a buffer108 of individual entity transitions{ [Equation] [Equation] +
Revised: Nn =1 to a buffer of individual entity transitions { ( h kt , a t , h kt +1 ) }

Original: [Equation] ) } Nn =1 , and thus abuffer of individual state transitions{ (
Revised: Nn =1 , and thus a buffer of individual state transitions

Original: [Equation]
Revised: [Equation]

Original: , a t , s

Original: [Equation] 
Revised: [Equation] [Equation]

Original: +1 ) } Nn =1 , the secondstep is to cluster similar transitionstogether by clustering the statesbefore and after each transition.
Revised: , the second step is to cluster similar transitions together by clustering the states before and after each transition.

Original: Our cluster procedure achieves this by clustering over all statestogether, which is sufﬁcient since the end state of transition is the starting state of another.
Revised: Our cluster procedure achieves this by clustering over all states together, which is sufficient since the end state of transition is the starting state of another.

Original: A properclustering should return centroids { s ⇤ } M m =1 that all correspond to actual states of entities, which we117 hypothesize will correspond to actual locations in the training set based on how dSLATE is trained.
Revised: A proper clustering should return centroids { s ∗ } Mm =1 that all correspond to actual states of entities, which we hypothesize will correspond to actual locations in the training set based on how dSLATE is trained.

Original: The purpose of edges between the nodes is to record an action that transforms an entity from onestate to another.
Revised: The purpose of edges between the nodes is to record an action that transforms an entity from one state to another.

Original: For each entity transition ( h kt , a t , h kt +1 )
Revised: For each entity transition ( h kt , a t , h kt +1 )

Original: we bind h kt and
Revised: we bind h kt and

Original: h kt +1 to their associatednodes s [ i ] ⇤ and s [ j ] ⇤ and create an edge between s [ i ] ⇤ and s [ j ] ⇤ tagged with action a , overwriting previous edges based on the assumption that with a proper clustering there should only be one action primitiveper pair of nodes.
Revised: h kt +1 to their associated nodes s [ i ] ∗ and s [ j ] ∗ and create an edge between s [ i ] ∗ and s [ j ] ∗ tagged with action a , overwriting previous edges based on the assumption that with a proper clustering there should only be one action primitive per pair of nodes.

Original: Applying the bind procedure to an entity h k returns the index of the cluster of thecentroid s ⇤ that is nearest to the entity’s state s k .
Revised: Applying the bind procedure to an entity h k returns the index of the cluster of the centroid s ∗ that is nearest to the entity’s state s k .

Original: For our experiments cluster is implemented with
Revised: For our experiments cluster is implemented with

Original: K-means and bind is implemented with nearest neighbors using the same distance metric used for
Revised: K-means and bind is implemented with nearest neighbors using the same distance metric used for

Original: K-means (see Appendix), but other clustering algorithms and distance metrics can also be used.
Revised: K-means (see Appendix), but other clustering algorithms and distance metrics can also be used.

Original: 2.2 Control
Revised: 2.2 Control

Original: Algorithm 1: given model , graph 2: input goal o g , observation o t 3: # infer goal constraints and current entities 4: h g , h t  model ( o g ) , model ( o t ) 5: align entity indices of h t with those of h g 6: ⇡ align ( h t , h g ) 7: permute indices of h t according to ⇡ 8: h t  ( h ⇡ [1] t , ..., h ⇡ [ K ] t ) 9: identify k th goal constraint to satisfy next 10: k  select-constraint ( h t , h g ) 11: infer cluster assignments 12:
Revised: 1: given model , graph 2: input goal o g , observation o t 3: # infer goal constraints and current entities 4: h g , h t ← model ( o g ) , model ( o t ) 5: align entity indices of h t with those of h g 6: π ← align ( h t , h g ) 7: permute indices of h t according to π 8: h t ← ( h π [1] t , ..., h π [ K ] t ) 9: identify k th goal constraint to satisfy next 10: k ← select-constraint ( h t , h g ) 11: infer cluster assignments 12:

Original: [ i ] , [ j ]  bind  h kt  , bind  h kg  13: action that transforms node [ i ] to node [ j ] 14: return graph.edges [ i, j ] .action
Revised: [ i ] , [ j ] ← bind (cid:0) h kt (cid:1) , bind (cid:0) h kg (cid:1) 13: action that transforms node [ i ] to node [ j ] 14: return graph.edges [ i, j ] .action

Original: Having constructed a graph of state transitions, ourapproach for control re-composes sequences of statetransitions to solve new rearrangement problems.
Revised: Having constructed a graph of state transitions, our approach for control re-composes sequences of state transitions to solve new rearrangement problems.

Original: Speciﬁcally, the agent decomposes the rearrangementproblem into a set of per-entity subproblems, searchesthe transition graph for a transition that transformsthe current entity’s state to its goal state, and executesthe action tagged with this transition in the environment.
Revised: Specifically, the agent decomposes the rearrangement problem into a set of per-entity subproblems, searches the transition graph for a transition that transforms the current entity’s state to its goal state, and executes the action tagged with this transition in the environment.

Original: This problem decomposition is possible because the transitions in our graph are constructed tobe agnostic to identity and context, enabling different rearrangement problems to share solutions to thesame subproblems.
Revised: This problem decomposition is possible because the transitions in our graph are constructed to be agnostic to identity and context, enabling different rearrangement problems to share solutions to the same subproblems.

Original: The core challenge in decidingwhich transitions to compose is in determining whichtransitions are possible to compose.
Revised: The core challenge in deciding which transitions to compose is in determining which transitions are possible to compose.

Original: That is, the agentmust determine which nodes in the graph correspondto the given goal constraints and which nodes correspond to the entities in the current observation, but the current entities h t and goal constraints h gmust themselves be inferred from the current and goal observations o t and o g , requiring the agent toinfer both what to do and how to do it purely from its sensorimotor interface.
Revised: That is, the agent must determine which nodes in the graph correspond to the given goal constraints and which nodes correspond to the entities in the current observation, but the current entities h t and goal constraints h g must themselves be inferred from the current and goal observations o t and o g , requiring the agent to infer both what to do and how to do it purely from its sensorimotor interface.

Original: Our approach takes four steps (Alg. 2, Fig. 3).
Revised: Our approach takes four steps (Alg. 2, Fig. 3).

Original: The ﬁrst step applies dSLATE to infer h t and h g fromo t and o g .
Revised: The first step applies dSLATE to infer h t and h g from o t and o g .

Original: The second step permutes the indices of h t and h g to align the identities of h kt with h kg .
Revised: The second step permutes the indices of h t and h g to align the identities of h kt with h kg .

Original: This step is necessary because the permutation invariance of the entities does not guarantee that theentities with the same identities will have the same index k , i.e. z kt might not correspond to z kg .
Revised: This step is necessary because the permutation invariance of the entities does not guarantee that the entities with the same identities will have the same index k , i.e. z kt might not correspond to z kg .

Original: Weimplement this using the Hungarian algorithm that minimizes the matching cost between z t and z g .
Revised: We implement this using the Hungarian algorithm that minimizes the matching cost between z t and z g .

Original: The third step selects which goal constraint h kg to satisfy next.
Revised: The third step selects which goal constraint h kg to satisfy next.

Original: In the case where objects can moveindependently, we implement this select-constraint procedure by determining which constrainth kg has the highest difference in state with its counterpart h kt , which reduces to solving the sameargxmax problem as in isolate .
Revised: In the case where objects can move independently, we implement this select-constraint procedure by determining which constraint h kg has the highest difference in state with its counterpart h kt , which reduces to solving the same argxmax problem as in isolate .

Original: We discuss how we handle more complex dependencies in the
Revised: We discuss how we handle more complex dependencies in the

Original: Appendix.
Revised: Appendix.

Original: Lastly, we bind h kt and h gt to the graph and return the action tagged to the edge betweentheir respective nodes.
Revised: Lastly, we bind h kt and h gt to the graph and return the action tagged to the edge between their respective nodes.

Original: If an edge does not exist between the nodes, we simply take a random action.
Revised: If an edge does not exist between the nodes, we simply take a random action.

Original:  Experiments
Revised:  Experiments

Original: The crucial test for knowledge reuse is to solve a set of tasks from a part of the combinatorial spaceof object rearrangement problems that are disjoint from the part of the space given for training, i.e.generalizing to different numbers of objects.
Revised: The crucial test for knowledge reuse is to solve a set of tasks from a part of the combinatorial space of object rearrangement problems that are disjoint from the part of the space given for training, i.e. generalizing to different numbers of objects.

Original: We measure the “fractional success rate:” the meanchange in the number of satisﬁed constraints divided by the number of initially unsatisﬁed constraintsfor each episode.
Revised: We measure the “fractional success rate:” the mean change in the number of satisfied constraints divided by the number of initially unsatisfied constraints for each episode.

Original: In the complete setting, all objects have an associated constraint, whereas in thepartial setting, only a subset of objects have associated constraints, and this is reﬂected in the goalimage (Fig. 2b).
Revised: In the complete setting, all objects have an associated constraint, whereas in the partial setting, only a subset of objects have associated constraints, and this is reflected in the goal image (Fig. 2b).

Original: We consider three challenging environments: block-rearrange , robogym-rearrange ,and block-stacking .
Revised: We consider three challenging environments: block-rearrange , robogym-rearrange , and block-stacking .

Original: We compare against ﬁve baselines.
Revised: We compare against five baselines.

Original: We implement a version of MPC [36] byreplacing our factorized graph search with CEM.
Revised: We implement a version of MPC [36] by replacing our factorized graph search with CEM.

Original: We investigate an ablation of HA that constructs amonolithic graph over transitions between sets of slot states rather than individual slot states (MGS).
Revised: We investigate an ablation of HA that constructs a monolithic graph over transitions between sets of slot states rather than individual slot states (MGS).

Original: We also compare with state-of-the-art pixel-based behavior cloning (BC) and implicit Q-learning(IQL) implementations based off of [22].
Revised: We also compare with state-of-the-art pixel-based behavior cloning (BC) and implicit Q-learning (IQL) implementations based off of [22].

Original: Our last baseline just takes random actions (Rand).
Revised: Our last baseline just takes random actions (Rand).

Original: Results Figure 4 shows that HA performs signiﬁcantly better than the baselines in combinatorialgeneralization (about a 5-10x improvement).
Revised: Results Figure 4 shows that HA performs significantly better than the baselines in combinatorial generalization (about a 5-10x improvement).

Original: Notably, the MPC baseline performs poorly because itsrollouts are poor, and it is signiﬁcantly more computationally expensive to run (11 hours instead of
Revised: Notably, the MPC baseline performs poorly because its rollouts are poor, and it is significantly more computationally expensive to run (11 hours instead of

Original: 20 minutes).
Revised: 20 minutes).

Original: Fig.
Revised: Fig.

Original: 3a shows a t-SNE plot of the factors over 100 batches in the training set, wheredifferent colors label different clusters, with several clusters annotated with the average slot attentionmask for the entities in the cluster.
Revised: 3a shows a t-SNE plot of the factors over 100 batches in the training set, where different colors label different clusters, with several clusters annotated with the average slot attention mask for the entities in the cluster.

Original: The right panel of Fig.
Revised: The right panel of Fig.

Original: 3 walks through the steps our algorithmtakes to select an action, visualizing the align , isolate , and bind procedures, showing that theentities inferred with dSLATE intuitively capture objects that can be independently acted upon.
Revised: 3 walks through the steps our algorithm takes to select an action, visualizing the align , isolate [Figure] and bind procedures, showing that the entities inferred with dSLATE intuitively capture objects that can be independently acted upon.

Original: [Figure] 
Revised: [Figure] [Figure]

Original: [Figure] [Table] [Table]
Revised: [Table]

Original: [Table] [Table] 
Revised: [Table] [Table] [Table]

Original:  Discussion
Revised:  Discussion

Original: We have demonstrated that HA outperforms various ofﬂine RL methods across three simulated objectrearrangement problems. 
Revised: We have demonstrated that HA outperforms various offline RL methods across three simulated object rearrangement problems. We hope HA will inspire future work on how abstractions of independent and symmetric entities can be learned and re-composed. Acknowledgments

Original: We hope HA will inspire future work on how abstractions of independentand symmetric entities can be learned and re-composed.
Revised: We would like to thank Leslie Kaelbling for valuable feedback and Yash Sharma and Yilun Du for valuable discussions. This material is supported in part by the Fannie and John Hertz Foundation, as well as with ONR grant #N00014-18-1- A Additional Related Work

Revised: Structured MDPs.

Revised: Many works have postulated that designing algorithms for structured MDPs will lead to improvements in sample efficiency and generalization over existing algorithms for the standard

Revised: MDP formulation.

Revised: The two closest types of MDP families to our work are factored MDPs [3, 4, 17], relational MDPs [37], and object-oriented MDPs [9].

Revised: Factored MDPs [3, 4, 17] assume that the environment can be represented by discrete attributes, and that transitions between these attributes can be modeled as a Bayesian network.

Revised: Our work differs from these in that we do not assume access to these attributes and the dependency graph is also not assumed to be known.

Revised: More importantly, the focus in this work is on using attributes to factorize a task into subcomponents; and in particular being able to generalize to new, more complex tasks at test time.

Revised: Our approach is also related to Relational

Revised: MDPs and Object-Oriented MDPs [1, 9, 12], where states are described as a set of objects, each of which is an instantiation of canonical classes, and each instantiated object has a set of attributes.

Revised: Our work is especially related to [18], where the aim is to show that by using a relational representation of an MDP, a policy from one domain can generalize to a new domain through planning.

Revised: However, we discover both objects (identities) and attributes (states), and achieve generalization through factorized planning, which grows linearly (rather than polynomially) in the number of factors.

Revised: B Implementation Details

Revised: This section details the implementation design decisions for each component of HA.

Revised: B.1 Level 1: abstracting visual features into sets of entities

Revised: Dynamic SLATE (dSLATE) maps a video demonstration to a sequence of transitions over sets of entities.

Revised: It consists of two main components: SLATE [32] and a transformer [28, 35] dynamics model.

Revised: Because SLATE itself also uses a transformer for generating image reconstructions, dSLATE uses two transformers: one to recover the observation model E and one to recover the dynamics model P of the structured MDP.

Revised: We use SLATE (cid:81)

Revised: k H k

Revised: × O → (cid:81)

Revised: k H k to infer entities h t from observation o t and initial guess ˆ h t .

Revised: The entity h k is also referred to as a slot in [25, 32] and is split in half as h k =

Revised: ( z k , s k ) .

Revised: The first guess for each entity ˆ h k 1 is sampled independently and identically distributed from a unit Gaussian, whose parameters are also trained.

Revised: The hyperparameters are given in Tab. 5.

Revised: SLATE preprocesses the image with a discrete variational autoencoder [29] into a grid of image features, encodes these features into a grid of tokens, infers slots from this token grid with Slot

Revised: Attention [25], which also produces an attention mask attn over the features each slot attends to.

Revised: These slots are trained using a transformer decoder [28, 35] to autoregressively reconstruct the tokens using the slots as keys/values.

Revised: We use the dynamics model (cid:81)

Revised: k S

Revised: k

Revised: ×

Revised: A → (cid:81)

Revised: k S k to predict a guess for ˆ s t +1 given the action a t and the inferred entities s t from the previous time-step.

Revised: This dynamics model is implemented also as a transformer decoder, taking the entity states as queries, and the entity states and action as keys/values.

Revised: This enables us the dynamics model to model the future state of the slots as an equivariant function of how the action affects it and of how it interacts with other entity states.

Revised: We trained dSLATE on an offline dataset of 5000 video demonstrations of length five, with each frame transition showing one of four objects being moved to a different location.

Revised: We used five slots, one more than the number of objects, following the convention used in Van Steenkiste et al.

Revised: [34], Veerapaneni et al.

Revised: B.2 Level 2: abstracting transitions over sets of entities into individual state transitions

Revised: HA constructs a graph of state transitions from a buffer of transitions over entity sets produced by dSLATE, as described in Alg. 1.

Revised: Each transition records the identity z , state s , and attn of each entity of the entity sets h t and h t +1 .

Revised: We also record the pre-condition of the transition, which we explain further in Appdx.

Revised: B.3.

Revised: Both s or attn can serve as the representation of the state.

Revised: We found that we obtained better clusterings when we used attn as the state for the block-* tasks and s as the state for the robogym-rearrange task.

Revised: We also empirically found that certain choices of distance metric used for K-means cluster ing and bind ing (implemented as nearest-neighbors) depended

Revised: [Table]

Revised: on which choice of state representation we used

Revised: [Table]

Revised: and this is summarized in Table 6.

Revised: The K-means implmentation is adapted from https://github.com/overshiki/kmeans_pytorch .

Revised: We found that increasing the number of slot attention iterations improved the entities representations especially when generalizing to more numbers of objects, so even though we dSLATE trained with slot attention three iterations, for inferring the slots from the buffer we used seven iterations.

Revised: Lastly, we found that the number of clusters used to for K-Means is the most important hyperparameter for creating a graph that reflected the state transitions.

Revised: We swept over 16 to 50 clusters and report the optimal number of clusters we found in Table

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: B.3 Action selection

Revised: Using dSLATE and the transition graph from HA, HA returns which action to execute in the environment given a goal and current observation, as described in Alg. 2.

Revised: It first infers goal constraints h g and current entities h t from the goal observation o g and current observation o t .

Revised: It then uses the align procedure to align the indices of the entities in h g and h t and uses the select-constraint to choose the index k of the entity to affect.

Revised: It binds h kt and h kg to the graph and returns the action associated with the edge between their respective nodes.

Revised: Because HA is a non-parameteric method, it could be the case the graph does not contain such an edge.

Revised: In this case, we sample a random action in the environment, but future work will replace this step with a more sophisticated method.

Revised: To implement align we use the scipy.optimize.linear_sum_assignment implementation of the Hungarian algorithm, with Euclidean distances between the z k ’s as the matching cost.

Revised: In select-constraint , we are given the set of current entities h t whose indices are aligned with the goal constraints h g and returns the index k of the goal constraint to satisfy next.

Revised: By HA’ construction, the edge between the nodes that h kt and h kg are bound to is the state transition that would be executed if the action associated to the edge were taken in the environment.

Revised: The select-constraint procedure consists of three steps: (1) filtering possible transitions from impossible transitions (2) ranking transitions (3) sampling a transition.

Revised: Filtering The filtering step implements HA’ model of possibility and impossibility.

Revised: In the filtering step, we consider, for each k , the transition between the nodes that h kt and h kg are bound to and mark the transition as possible or impossible.

Revised: It then returns the indices k over h t and h g whose associated transition from h kt and h kg is possible.

Revised: According to HA, a state transition between node [ i ] and node [ j ] is possible if its preconditions are met and there exists an edge for that state transition in the graph, and impossible otherwise.

Revised: An intuitive example of an impossible transition is to place a block midair at some intended height, but this transition becomes possible if prior to the transition there already exists a stack of blocks that would support the block if the block were to be placed at that intended height.

Revised: The existence of this supporting stack is thus the precondition for the transition to occur, rendering the stacked block dependent on the blocks supporting it.

Revised: When there are no dependencies among the entities, as in the *-rearrange tasks where any object can be moved to any open location without considering where other objects are, any transition present in the graph is possible.

Revised: When there are dependencies among the entities, as in block-stacking , we take the precondition of the transition into account.

Revised: Although a precondition of a transition from node [ i ] to node [ j ] could be a function of both the source node [ i ] and destination node [ j ] , for simplicity in this paper we consider preconditions as only a function of the destination node [ j ] , which rules out the possibility of placing a block in midair like the above example.

Revised: Because a precondition in general is a set of constraints that need to be satisfied for the transition to be possible, we represent the precondition of a transition into node [ j ] (denoting the state s [ j ] ∗ ) as the set of context states s [ j ′ ] ∗ that are always present whenever the state s

Revised: [ j ] ∗ is present.

Revised: Concretely, a block at height 3 at location x is always accompanied by the presence of some block at height 2 and some block at height 1, both at location x ; the states denoting (location x , height 2) and (location x , height 1) are the context states of the state (location x , height 3) and therefore serves as its precondition.

Revised: We thus implement the precondition of a transition into node [ j ] by recording the indices j ′ of the nodes of states that are always present when node [ j ] is a destination node.

Revised: To test whether a precondition is satisfied for a given scene, we check if all nodes in the precondition set have a corresponding concrete entity that can be bound to it.

Revised: Ranking The filtering step removes the indices from the entities h t and goal constraints h g whose transitions are impossible, yielding a possibly smaller set of entities ˜ h t and constraints ˜

Revised: h g .

Revised: That is, if | h t | = | h g | =

Revised: K , then | ˜ h t | = | ˜

Revised: h g | =

Revised: ˜ K ≤ K .

Revised: The goal of the ranking step is to compute a ranking among the indices of ˜ h t and ˜

Revised: h g for choosing which index k to actually select to affect with an action.

Revised: Intuitively, we should rank indices k according to how different s kt and s kg are because a difference indicates that the constraint h kg is not satisfied.

Revised: We reuse the distance metric d ( · , · ) used for isolate to implement this ranking.

Revised: Sampling The goal of the sampling step is to select a k ∈ { 1 , ..., ˜ K } whose associated entity we will affect with an action, given our ranking.

Revised: One way to do this is to simply choose k as k = arg max k ′ ∈{ 1 ,..., ˜ K } d ( s k ′ t , s k ′ t +1 ) as in isolate , but we empirically found that sampling k from a Categorical distribution whose pre-normalized probabilities are given by d ( s k ′ t , s k ′ t +1 ) resulted in better task performance so we used this stochastic sampling approach.

Revised: One explanation for why using the argmax may be worse is that it relies on the distance metric d ( · , · ) , and the state representation s , to be such that the distance metric flawlessly assigns high value to entities k that need to be moved and low value to entities k that do not need to be moved.

Revised: But because the state space S is learned through the dSLATE training process without explicit supervision on the geometry of the space, a pair of points that should be farther apart than another set of points may not be accurately reflected by using a fixed distance metric d ( · , · ) .

Revised: Future work will investigate imposing explicit supervision on the geometry of S .

Revised: C Baseline Implementation Details

Revised: Random (Rand)

Revised: The random policy takes actions using env.action_space.sample() .

Revised: Behavior cloning (BC)

Revised: This approach trains a policy to output the actions directly taken in the provided dataset.

Revised: We use an MSE loss to train the policy to imitate the actions.

Revised: Implicit Q-learning (IQL) IQL is a simple, offline RL approach that uses temporal difference (TD) learning with the dataset actions and trains a behavior policy value function.

Revised: To produce an optimal value function, IQL estimates the maximum of the Q-function using expectile regression with an asymmetric MSE using the following objectives:

Revised: [Equation]

Revised: ] where L τ

Revised: [Equation]

Revised: [Equation]

Revised: u <

Revised: [Equation]

Revised: The V ( s ) estimates are used for TD-backups and the optimal policy is extracted with advantageweighted behavioral cloning.

Revised: Model predictive control (MPC)

Revised: This approach uses model predictive control with the cross entropy method (CEM) to select actions, using the transformer dynamics model of dSLATE to perform rollouts in latent space.

Revised: This is similar to the approached used in OP3 [36], except that we use more recently proposed architectural components (slot attention [25] instead of IODINE [15], a transformer instead of a graph network [2, 7, 34]) so our MPC results are not directly comparable to that of OP3.

Revised: We use the same dSLATE checkpoint that was used for HA.

Revised: We implement this MPC baseline using the mbrl-lib library [27] with 10 CEM iterations, an elite ratio of 0.05, and a population size of 250 which was the best configuration we found that fit within a wall clock budget of two days for 8 objects and 100 test episodes.

Revised: We swept over CEM iterations of [5 , 10 , 20] , elite ratio of [0 . 05 , 0 . 1 , 0 . 2] , and population sizes of [250 , 500 , 1000] , and found that the elite ratio was the most important hyperparameter.

Revised: The cost function is computed by first aligning the predicted slots h T and goal constraints h g using the same align procedure in B.3, and then adding up the squared Euclidean distance between slots as cost =

Revised: (cid:80) k

Revised: ( h kT − h kg ) 2 .

Revised: Monolithic graph search (MGS)

Revised: This approach is an ablation to HA that does not construct a graph over state transitions of individual entities but instead constructs a graph over state transition over entity sets, i.e. each transition is ( s , a, s ′ ) rather than ( s k , a, s k ′ ) .

Revised: As with MPC, we use the same dSLATE checkpoint that was used for HA.

Revised: The purpose of this ablation is to elucidate the benefit of factorizing the transition graph over entities rather than entity sets.

Revised: Because nodes in this transition graph represent a set of entity states rather than individual entity states, we use Dijkstra’s algorithm, as in [10, 38, 39] to plan a unbroken path from the node the initial observation is bound to to the node a goal observation is bound to.

Revised: For each time-step, we plan a path along the nodes using Dijkstra’s algorithm, then return the action associated with the first edge along that path.

Revised: Like HA, MGS is a non-parametric model, which means that for a set of entities to be bound to a node in the graph, that node must contain the exact set of entity states corresponding to the states of the entities.

Revised: If we do not successfully bind to the graph, or if we do not find a path between the current node and the goal node, we sample a random action as HA does.

Revised: D Environment Details

Revised: Block-rearrange is our simplest environment, and robogym-rearrange and block-stack each add complexity along different axes.

Revised: States and identities correspond to object locations and appearance respectively.

Revised: In block-rearrange (Fig. 2a), all objects are the same size, shape, and orientation.

Revised: S covers 16 locations in a grid.

Revised: Z is the continuous space of red-green-blue values from 0 to 1 .

Revised: robogym-rearrange (fig. 2b) adapts the rearrange environment from OpenAI [26] and removes the assumption from block-rearrange that all objects are the same size, shape, orientation, and the assumption of predefined locations.

Revised: block-stack (fig. 2c) adds preconditions and postconditions on whether objects can be moved: blocks can only be picked if from the top of a stack, and blocks can only be placed at a given height if there is an object beneath to support it (otherwise it falls).

Revised: Block-rearrange and Block-stack are implemented in PyBullet [8] while Robogym-rearrange is implemented in Mujoco [33].

Revised: Environments In block-rearrange , all objects are the same size, shape, and orientation and can exist in any one of 16 locations in a grid.

Revised: Colors are sampled in a continuous space of red-green-blue values in [0 , 1] .

Revised: The robogym-rearrange environment (see figures 5 and 6) is adapted from the rearrange environment in OpenAI’s Robogym simulation framework [26] and removes the assumption from block-rearrange that all objects are the same size, shape, and orientation and the assumption of predefined locations.

Revised: Furthermore, due to 3D perspective, the objects can look slightly different in different locations.

Revised: The objects are uniformly sampled from a set of 94 meshes consisting of the YCB object set [5] and a set of basic geometric shapes, with colors sampled from a set of 13.

Revised: The camera angle is a bird’s eye view over the table, and the size of each object is normalized by its longest dimension, so tall thin objects appear smaller.

Revised: The objects’ target positions are randomly sampled such that they don’t overlap with each other or any of the initial positions, and the target orientation is set to be unchanged.

Revised: Due to the continuous nature of this environment, we define a match threshold of at most 0.05 for both the initial pick position and the goal placement (the table dimensions are 0.6 by 0.8).

Revised: The block-stack environment adds preconditions and postconditions on whether objects can be moved: blocks can only be picked if they are at the top of a stack, and blocks can only be placed at a given height if there is already an object beneath to support it (otherwise it falls).

Revised: Sensorimotor interface Each observation is a tuple of an initial image displaying the current observation and a goal image displaying constraints to be satisfied – the goal locations of the objects.

Revised: Each action is a tuple ( w, ∆ w ) , where w is a three-dimensional Cartesian coordinate ( x, y, z ) in the environment arena.

Revised: For the ∗ -rearrange tasks, objects are initialized at random non-overlapping locations that also do not overlap with their goal locations.

Revised: For these tasks the z (height) coordinate is always fixed.

Revised: For the block-stack task, the goal locations are generated by randomly picking objects from the tops of stacks and placing them on other stacks.

Revised: For this task the y (depth) coordinate is always fixed.

Revised: An object is picked if w is within a certain threshold of its location.

Revised: For block- ∗ where object locations are fixed points in a grid the object is snapped to the nearest grid location to w + ∆ w .

Revised: Constraints are considered satisfied if objects are placed within a certain threshold of their target location.

Revised: E Additional Results

Revised: This section presents additional results and analyses of HA.

Revised: We first analyze the sensitivity of task performance to several hyperparameters used in HA from creating the graph: the number of clusters, number of buffer size, and the number of slots used in slot attention.

Revised: We next tested whether giving our model-based baselines more computation time compared to HA would improve their performance to be comparable to HA’s.

Revised: E.1 Analysis of key hyperparameters

Revised: In this section, we analyze the sensitivity of our method to various hyperparameters, evaluated in the robogym environment with four objects in the complete goal specification.

Revised: As Fig.

Revised: 7 shows, performance depends on the number of initialized clusters and the number of batches from the training set used to construct the graph.

Revised: With too few clusters, the clusters are too coarse-grained to differentiate objects in significantly different positions, while with too many the performance deteriorates as the data is needlessly split into duplicate clusters.

Revised: Performance improves with more data, as the graph has better coverage.

Revised: Although our model performs worse when there are insufficient slots to represent all objects present in the environment plus one empty slot, performance is barely impacted by having double the number of necessary slots.

Revised: Our method can thus still work in environments with an unknown but upper-bounded number of objects.

Revised: E.2 More challenging evaluation settings

Revised: We analyzed HA in more challenging settings that reflect the noisy nature of real-world robotics.

Revised: As

Revised: Fig.

Revised: 8 (left) shows, HA is robust to the addition of Gaussian noise to the action at every time step, up until the noise variance is comparable to the maximum distance for successful picking and goal placements.

Revised: The performance also remains high given significantly fewer steps (Fig. 8, right).