Hierarchical Abstraction for Combinatorial Generalization in Object Rearrangement
Authors: Michael Chang
Alyssa Li Dayan
Franziska Meier
Thomas L. Griffiths
Sergey Levine
Amy Zhang
Venue: ICLR 2023 poster (2023)
Keywords: objects
combinatorial generalization
abstraction
rearrangement
slots
binding
hierarchy
compositionality
symmetry
independence
graph search

Abstract: Object rearrangement is a challenge for embodied agents because solving these tasks requires generalizing across a combinatorially large set of configurations of entities and their locations. Worse, the representations of these entities are unknown and must be inferred from sensory percepts. We present a hierarchical abstraction approach to uncover these underlying entities and achieve combinatorial generalization from unstructured visual inputs. By constructing a factorized transition graph over clusters of entity 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. 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 simulated rearrangement tasks.

Published as a conference paper at ICLR 2023
HIERARCHICAL ABSTRACTION FOR COMBINATORIAL GENERALIZATION IN OBJECT REARRANGEMENT
Michael Chang, Alyssa L. Dayan, Franziska Meier, Thomas L. Grifliths, Sergey Levine, Amy Zhang
ABSTRACT
Object rearrangement is challenge for embodied agents because solving these tasks requires generalizing across a combinatorially large set of configurations entities and their locations. Worse, the representations of these entities un- known and must be inferred from sensory percepts. We present a hierarchical abstraction approach to uncover these underlying entities and achieve combinato- rial generalization from unstructured visual inputs. By constructing a factorized transition graph over clusters of entity representations inferred from pixels, we show how to learn a correspondence between intervening on states of entities the agent's model and acting on objects in the environment. We use this corre- spondence 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 simulated rearrangement tasks.
INTRODUCTION
The power of an abstraction depends on its usefulness for solving new problems. Object rearrangement (Batra et al. 2020) offers intuitive setting for studying the problem of learning reusable ab- stractions. Solving novel rearrangement problems requires agent to not only infer object representations without supervision, but also recognize that the same action for moving an object between two locations can be reused for different objects in different contexts.
S2
The second lovel groupe entiliee that share the sume state transition to produce a graph over entily- agnostic state transitions
state (testuure)
entity-oet
We study the simplest setting in simulation with pick-and-move ac- tion primitives that move one object at a time. Even such a simple setting is challenging because the space of object configurations is combinatorially large, resulting in long-horizon combinatorial task
hg
ideatity (soldd cakor
hic
spaces. We formulate rearrangement as an offline goal-conditioned reinforcement learning (RL) problem, where the agent is pretrained on a experience buffer of sensorimotor interactions and is evaluated on producing actions for rearranging objects specified in the input image to satisfy constraints depicted in a goal image.
groupe
features from seturorimoto Interaction produce transitions between sets of entities
Offline RL methods (Levine et al. 2020) that do not infer factorized representations of entities struggle to generalize to problems with more objects. But planning with object-centric methods that do in- fer entities (Veerapaneni et al., 2020 is also not easy because the difficulties of long-horizon planning with learned parametric mod-
objects
Rg
tanget object
els (Janner et al. 2019) are exacerbated in combinatorial spaces.
Figure I: NCS uses a two-level hierarchy to abstract sensorimo- tor interactions into a graph of learned state transitions. The af- fected entity is in black.
Instead of planning with parametric models, our work takes inspi- ration from -ar-marretttri planning methods that have shown suc- cess in combining neural networks with graph search to generate long-horizon plans. These methods (Yang et all 2020 Zhang et al.] 2018 Lippi et al., 2020 Emmons et al. 2020 explicitly construct a transition graph from the perience buffer and plan by searching through the actions recorded in this transition graph with a learned distance metric. The advantage of such approaches is the ability to stitch different path seg- ments from offline data to solve new problems. The disadvantage is that the non-maramttric nature
"work done as an intern at Mcta AI. Correspondence to: mbchang@ berkeley.edu and amyzhang meta.com A step-by-step explanatory video of our method can he found in the supplementary material.

---

Published as a conference paper at ICLR 2023
of such methods requires transitions that will be used for solving new problems to have already been recorded in the buffer, making conventional methods, which store entire observations mono- lithically, ill-suited for combinatorial generalization. Fig. P shows that the same state transition can manifest for different objects and in different contexts, but monolithic non-parametric methods are not constrained recognize that all scenarios exhibit the same state transition an abstract level. This induces an blowup in the number of nodes of the search graph. To overcome this problem, we devise a method that explicitly exploits the similarity among state transitions in different contexts.
Our method, Neural Constraint Satisfaction (NCS), marries the strengths of pon-parametric plan- ning with those of object-centric representations. Our main contribution is show that factoriz- ing the traditionally monolithic entity representation into action-invariant features (its type) and action-dependent features (its state) makes possible during planning and control reuse action representations for different objects in different contexts, thereby addressing the core combinato- rial challenge in object rearrangement. To implement this factorization, NCS constructs a two-level hierarchy (Fig- T to abstract the experience buffer into a graph over state transitions of individual en- tities, separated from other contextual entities (Fig- Sp. To solve new rearrangement problems, NCS infers what state transitions can be taken given the current and goal image observations, ce-oomossss sequences of state transitions from the graph, and translates these transitions into actions.
In S we introduce a problem formulation that exposes the combinatorial structure of object re- arrangement tasks by explicitly modeling the independence, symmetry, and factorization of latent entities. This reveals two challenges in object rearrangement which we call the correspondence problem and combinatorial problem. In $ we present NCS, a method for controlling an agent that plans over and acts with emergent learned entity representations, ling both challenges. We show in $ that NCS outperforms both state-of-the-art offline RL methods and object-centric shooting- -aased planning methods in simulated rearrangement problems.
unified method for tack-
RELATED WORK
The problem of discovering re-composable representations is generally motivated by combinatorial task spaces. The traditional approach to enforcing this compositional inductive bias is to compactly
represent the task space with MDPs that cuuann-defieed abstractions of entities, such as factored
MDPs Boutilier et al. (1955; 2000); Guestrin et al. (003a), relational MDPs Waag et al.
2008);
Guestrin et al. (003b) : Gardiol Kaelbling (2003). and object-oriented MDPs Diuk et al.
2008);
Abel et al. (015). 50.
Approaches building off of such symbolic abstractions (Chang et al.
2016
Battaglia et al. 2018; & Zadaianchuk et al. 2022] Bapst et al. 2019 Zhang et al. 2018) do not address
the problem of how such entity abstractions arise from raw data. Our work is one of the first to learn compact representations of combinatorial task spaces directly from raw sensorimotor data.
Recent object-centric methods (Greff et al. 2017, Van Steenkiste et al.] 2018) Greff et al. 2019
2020 Locatello et al. 2020a; Kipf et al. 2021| Zoran ct al. [022] Singh et al. 2021) do learn entity representations, as well as their transformations (Goyal al. 202 2020), from sensorimotor data, but only do so for modeling images and video, rather than for taking actions. Instead, we study how well rnteettrepresentatiions can reused for solving tasks. Kulkarni et al. 2019) considers how object representations improve exploration, but we consider the offline setting which requires zero-shot generalization. |Veerapaneni et al. (200)) also considers on control tasks, but their shooting-based planning method in suffers from compounding errors as other learned single-step models do (Janner ct al. 2019) while our hierarchical non- -arramemtt approach enables us to plan for longer horizons.
Non-parametric approaches have recently become popular for long horizon planning (Yang el al., 2020 Zhang et al. 2018: Lippi al. 2020 Emmons al. 2020 Zhang al. 202TT. p. but the draw- back of these approaches is they represent the entire scenes monolithically, which causes a blowup of nodes in combinatorial task spaces, making it infeasible for these methods be applied in rear- rangement tasks that require generalizing to novel object configurations with different numbers of objects. Similar to our work, |Huang et al. 2019 also tackles rearrangement problems by search- ing over a constructed latent task graph, but they require demonstration during deployment time, whereas NCS does not because it reuses context-agnostic state transitions that were constructed dur- ing training. Zhang et al. (2021) conducts non-parametric planning directly on abstract subgoals rather than object-centric states - while similar, the downside of using subgoals rather than ab- stract states is that those subgoals are not used to represent equivalent states and therefore cannot generalize to new states at test time. Our method, NCS, captures both reachability between known states and new, unseen states that can be mapped to the same abstract state.
2

---

Published as a conference paper at ICLR 2023
Coreepondetoo Problem
(b) Combinatotial Probiem
eame type (echl coloe), difeent state (textum)
tiuggire iearo for
sheerved insettie
dibert tagget otject
Alfimma puelis
modd cillis
cornepondence problum
obecrottion
retios (jrkkglaee)
envertmante lste fare
dilieee ddpe
(ifieret lirght oljes netnst
Seealled daas insation
mail tranffrmatiii
object
curinsanaul
suuse state traakiku cus munifes fox difereat oljects in diliveat ocpriats
Figure 2: Solving object rearrangement requires solving two challenges. (a) The correspondence prob- lem is problem of abstracting raw sensorimotor signal into representations of entities such that there is a correspondence between how an agent intervenes on an entity and how its action affects an object in the envi- romment. & denotes the index of the entity, denotes its type (shown with solid colors), and denotes its state (shown with textures). The entity representing the moved object is in black. (b) The combinatorial problem is the problem of representing the combinatorial task space in a way that enables an agent transfer knowledge of a given state transition (indicated by the dotted cirele) to different contexts.
3
GOAL-CONDITIONS
REINFOR LEARNING WITH ENTITIES
This section introduces a set of modifications to the standard goal-conditioned partially observed Markov decision process (POMDP) problem formulation that explicitly expose the combinatorial structure of object rearrangement tasks of the following kind: "SSquentiallyy move a subset (or all) of the objects depicted in the current observation on to satisfy the constraints depicted in the goal image Og." We assume an offline RL setting, where the agent is trained on a buffer of transitions {(01, .1. 1r))) ) and evaluated on tasks specified as (01,0g
The standard POMDP problem formulation assumes an observation space O, action space A, latent space H. goal space G, observation function E : H + O, transition function P H x A + H. and reward function R H. x + R. Monolithically modeling the latent space this way does not expose commonalities among different scenes, such as scenes that contain objects in the same location or scenes with multiple instances of the same type of object, which prevents us from designing control algorithms that exploit these commonalities to collapse the combinatorial task space.
To overcome this issue, we introduce structural assumptions of independence, symmetry, and fac- torization to the standard formulation. The independence assumption encodes the intuitive property that objects can be acted upon without affecting other objects. This is implemented by decompos- ing the latent space into independent subspaces as 7 = HH' x x HK one for each independent degree of freedom (e.g. object) in the scene. The symmetry assumption encodes the property that the the same physical laws apply to all objects. This is implemented by constraining the observation function E, transition function and reward function to be shared across all subspaces, thereby treating 7²
H We define an entity h E H* as a member of such subspace, and an entity-set as the set of entities h = (h' ..... h that explain an observation, similar to Diuk et al. 2008); Weld Lastly, the factorization assumption encodes that each subspace can be decomposed as H Z x S. where the = E Z represents the entity's 's action-invariant features like appearance, and E S represents its action-dependent features like location. We call = the type and the state. Introducing these assumptions solves the problem of modeling the commonalities among different scenes stated above. It allows us to describe scenes that contain objects in the same location by assigning entities in different scenes to share the same state s. It allows us to describe a scene with multiple instances of the same type of object by assigning multiple entities in the scene share the same type z. This formulation also makes it natural to express goals as a set of constraints h, = (h)--- -.... h.),,, h.) To solve task is to take actions that transform the subset of entities in the initial observation whose types are given by 2, to new states specified by Sg.
Fl
Exposing this structure in problem formulation enables
exploit it by designing methods that represent entities in an independent, symmetric, and factorized way and that use these three properties to collapse the combinatorial task space. To do involves solving two problems: the cor- respondence problem of learning to represent entities in this way and the combinatorial problem
We use "object" refer to an independent degree of freedom in the environment, and "entity" to refer to the agent's representation of the object.
3

---

Published as a conference paper at ICLR 2023
Repersentation Learning
Graph Constraction
samee tope bubi colo), diliseat state (lodiur))
cnrrett estry trassiioe
Ake
poice
plane)
abstrutt stabe tranettion
Figure 3: Modeling. NCS constructs a two-level abstraction hierarchy to model transitions the experience buffer. (a) Level 1: NCS learns to infer a set of entities from sensorimotor transitions with pick-and-move actions, in which one entity is moved per transition. We enforce that the type z (shown with solid colors) of an entity remains unchanged between steps. The GPT dynamics model learns sparsely predict the states entities the next time-step. This addresses the correspondence problem by forcing the network to use predict and reconstruct observations through the entity bottleneck. (b) Level 2: NCS abstracts transitions over entity-scts into transitions over states of individual entities, constructing a graph where states are nodes and transitions between them edges. This is done by clustering entity transitions that share similar initial states and final states. This addresses the combinatorial problem by making it possible for state fransitions to reused for different entity types and with different context entities.
(shown with textures)
of using these properties to make planning tractable. The correspondence problem is hard because it assumes no human supervision of what the entities are. It also goes beyond problems solved by exist- ing object-centric methods for images and videos because it involves action: it requires representing entities such that there is a correspondence between how the agent models how its actions affect entities and how its actions actually affect objects in the environment. The combinatorial problem goes beyond problems solved by methods for solving object-oriented MDPs, relational MDPs, and factorized MDPs because it requires the agent to recognize whether and how previously observed state transitions can be used for new problems, using learned, not human-defined, entity represen- tations. The natural evaluation criterion for both problems is to test to what extent agent can pero-shot-generalle to solve rearrangement tasks involving new sets of object configurations that aree disjoint from the configurations observed in training, assuming that the training configurations have collectively covered Z and S. Our experiments in $ test exactly this.
Simplifying assumptions To focus on the combinatorial nature of rearrangement, we are not in- terested in low-level manipulation, so we represent each action as (w, Aw), where are Cartesian coordinates w = (x, y, 2). We assume actions sparsely affect one entity at a time and how an action affects object's 's state does not depend its identity. We are not interested in handling occlusion, so we assume that objects are constrained to the zy plane or xe plane and are directly visible to the camera. Following prior work Chansen-Estruch lation assumption that the state space can be partitioned into a finite set of equivalence classes, and that there is one action primitive that transitions between each pair of equivalence classes. Lastly, we assume objects can be moved independently. Preliminary experiments suggest that NCS can be augmented to support tasks like block-stacking that involve dependencies among objects, but how to handle these dependencies would warrant a standalone treatment in future work.
al. 2022 Castro al., 2009). we make bisimu-
4 NEURAL CONSTRAINT SATISFACTION
In 3 we introduced a structured problem formulation for object rearrangement and reduced it to solving the correspondence and combinatorial problems. We now present our method, Neural Con- straint Satisfaction (NCS) as a method for controlling an agent that plans over and acts with a state transition graph constructed from learned entity representations. This section divided into two parts: modeling and control. The modeling part is further divided into two parts: representation learning and graph construction. The representation learning part addresses the correspondence problem, while the graph construction and control parts address the combinatorial problem.
4.1 MODELING
The modeling component of NCS abstracts the experience buffer into a factorized state transition graph that can be reused across different rearrangement problems. Below we describe how we first
4

---

Published as a conference paper at ICLR 2023
train object-centric world model to infer entities that are independent, symmetric, and factorized and then construct the state transition graph by clustering entities with similar state transitions. These two steps comprise two-level abstraction hierarchy over the raw sensorimotor transitions.
Level 1: representation learning The first level concerns the unsupervised learning of entity rep- resentations that factorizes into their action-invariant features (their type) and their action-dependen features (their state). Concretely our goal is model a video transition Ot, dg 0t++1 as transition over entity-sets h,, dg -> h where each entity h is factorized as a pair h = (2* s*). Given setting where action moves only single object in the environment representation learning implies three criteria: (1) the world model properly identifies the individual entity h corresponding to the moved object. (2) only the state 8* of that entity should change, while its type z should remain unaffected, and (3) other entity representations 1i*h should also remain unaffected. Criteria (1) and (3) rule out standard approaches that represent an entire scene with a monolithic representation. so we need an eejject-centric world model instead of monolithic world model. But criterion (2) rules out standard object-centric world models (e-g. (Veerapaneni et all
time, successful
2020 Elsayed et al. 2022, [Singh et al,, 2022b)), which do not decompose entity representations into action-invariant and action-dependent features.
Because the parameters of mixture model are independent and symmetric by construction, we pro- pose to construct our factorized object-centric world model as an equivariant sequential Bayesian filter with mixture model as the latent state, where entity representations are the parameters of the mixture components. Recall that a filter consists of two major components, latent estimation and latent prediction. We implement latent estimation with the state-of-the-ae slot attention (SA) (Lo- catello et al. 2020b), based on the connection Chang et al. (2022) between mixture components and SA slots. We implement latent prediction with the transformer decoder (TFD) architecture (Vaswani et al. 2017) because TFD is equivariant with respect to its inputs. We denote the SA sl slots as X and SA a attn masks as a. We split each slot À E R" into two halves X E IR 2 and X R Given observations o and actions a, we embed the actions as ā with an feedforward network and implement the filter as:
\ Gaussian
X t+1 = TFD (queries = Xi, keys/values = [X*, azl)
Ae, C2 = SA ()).
d = X S+11
where [. is the concatenation operator, Å is the output of the latent prediction step, and X is the output of the latent estimation step. We embed this filter inside the SLATE backbone (Singh et al. 2022a) and call this implementation dynamic SLATE (dSLATE).
By constructing A++41 copy of X; dSLATE enforces the information contained X to be action- invariant, hence we treat A as dSLATE's ' representation of the entities' types. As for the entities' states, either the action-dependent part of the slots X or the attention masks may be sufficient and more intuitive to analyze all objects looks similar and there is no occlusion, while A may be more suitable other cases, and we provide an example of each in the experiments. To simplify notation going forward and connect with the notation in 3 we use h to refer to (A, a),, use to refer to X², and use s to refer to X or a. Thus by construction dSLATE satisfies criterion (2). Empirically we observe that it satisfies criterion (1) as well as SLATE does, and that TFD learns to sparsely edit Xj thereby satisfying criterion (3).
be used. Using
Level 2: graph construction Having produced from the first level a buffer of entity-set transitions (h,, aj - h } the goal of the second level (Fig. 3D) is to use this buffer to construct factor- ized state transition graph. The key to solving the combinatorial problem is to construct the edges of this graph to represent not state transitions of entire entity-sets (i.e. St, Gr + 51+1) as prior work does (Zhang al. 2018), but state transitions of individual entities (i.e. st az -> s)41). Construct- ing edges over transitions for individual entities rather than entity sets enables the same transition to be reused with different context entities present. Constructing edges over state transitions instead of entity transitions enables the same transition to be reused across entities with different types. This would enable the agent to recompose sequences of previously encountered state transitions for solv- ing new rearrangement problems with different entities in different contexts. Henceforth our use of "state" refers specifically to the state of individual entity unless otherwise stated.
Given our bisimulation assumption that states can be partitioned into a finite number of groups, construct our graph such that nodes represent equivalence classes among individual states and the edges represent actions that transform state from one equivalence class to another. To implement
5

---

Published as a conference paper at ICLR 2023
this we cluster state transitions of individual entities in the buffer, which reduces clustering the states of individual entities before and after the transition. We treat each cluster centroid as node in the graph, and edge between nodes is tagged with the single action that transforms one node's 's state to another's. The algorithm for constructing the graph is shown in Alg. | and involves three steps: (1) isolating the state transition of individual entity from the state transition of the entity- set, (2) creating graph nodes from state clusters, and (3) tagging graph edges with actions.
The first step is to identify which object was moved in each transition, i.e. identifying the entity h* that dSLATE predicted was affected by at in the transition (h,, oe 11) We implement a function isolate that achieves this by solving arg (1...e d(s)"" stt to identify the index of the entity whose state has most changed during the transition, where -(;;-) is distance function, detailed in Table B of the Appendix. This converts the buffer of transitions over entity-sets he, 0r - h1+1 into a buffer of transitions over individual entities
Algorithm Building the Graph
1: imput mode buffer
2: for {(or, a., 0c41)) in butfer do
3:
infer entities from transition
4: {(hr, at, h(+1) t- model (for, dt, 0t+1)a).
5: # iidettfy which entity changed in transition
6:
{(hā, ae, h241))
isolate (f ae, bet1)ba)
7: end for
8: # partition transitions by clustering entities
9: {s.)m=1 +- cluster ([(st. de =
10: transitions between clusters edges
11: initialize graph with nodes !mm for m E [1 [M:M]
12: for each {(hō ax, hf+1))= do
13:
# infer cluster assignments , (hi+)) [i). bind (h²) bind
14:
15:
tag edge with action ar
hū, a. + + h
16: graph.edgesf.e edges[i, + create- (3!) 2 slt))
The second step is to cluster the states before and after each transition. We implement a function è cluster that uses K-means to returns graph nodes as the centroids {s.) of these state clusters.
17:
end for
18: return graph
The third step is to connect the nodes with edges that record actions that actually were taken in the buffer to transform one state the next. We implement function bind that, given entity h* returns the index [i] of the centroid s, that is the nearest neighbor to the entity's state s For each entity transition (hh az, 11+1) we bind entity hy and h to their associated nodes s and s!
and create an edge between s! and s!il tagged with action a, overwriting previous edges based on the assumption that with proper clustering there should only be one action per pair of nodes.
In our experiments both cluster and bind use the same distance metric (see Table g in the Appendix), but other clustering algorithms and distance metrics can also be used. Our experiments (Fig. M also show that it is also possible to have more than one action primitive per pair of nodes as long as these actions all map between states bound to the same pair of nodes.
4.2 CONTROL
To solve new rearrangement problems, we re- compose sequences of state transitions from the graph. Specifically, the agent decomposes the rear- rangement problem into a set of per-entity subprob- lems (e.g. initial and goal positions for individual objects), searches the transition graph for a transition that transforms the current entity's state its goal state, and executes the action tagged with this tran- sition in the environment. This problem decomposi- tion is possible because the transitions in our graph are constructed to be agnostic to type and context, enabling different rearrangement problems to share solutions to the same subproblems. The core chal- lenge in deciding which transitions to compose is in determining which transitions are possible to com- pose. 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, and
Algorithm 2
Action Selection
given model, graph
2: input goal 0g. observation
# infer goal constraints and current entities
4: hg, hr <- model (0,), model (ox)
5: align entity indices of he with those of hg
6: x <- al i iggn (he, hg)
7: permute indices of h, according to m
8: h + [A ...., ArlK])
identify kth goal constraint to satisfy next
10: k +- select-cooe contt-constraaa iaaistrrnnnn (hr, hg)
11: infer cluster assignments
12: [i]. [i] +- bind (ht) bind (ht)
13: action that transforms node [i] to node [i]
14: return graph. edges [i, action

---

Published as a conference paper at ICLR 2023
(a) Action Selcction
(b) align
Og
Og
goal
gool
goal
gool
bind
blnd
saabe (hstar))
trje bolld
consomint
caty
sktt rstestion
400: atteotion
( bind
Figure 4: Planning and control. Given rearrangement problem specified only by the current and goal observations Oy), NCS decomposes the rearrangement problem into one subproblem 0g) per entity. (a) shows the computations NCS uses to solve each subproblem and (b-d) show these steps context. For each subproblem (ot, og). NCS infers entities from both the current and goal observations. The states of the goal entities indicate constraints on the desired locations of the current entities. (b) NCS aligns the indices of the current entities to those of the goal entities with corresponding types. (c) It selects the index lc of the next goal constraint s to satisfy, as indicated by the red box. The selected goal constraint and current entity are also colored black in (a), and note that their types are the same but states are different; we want to choose the action to transform the state of the current entity to the state of the goal constraint. (d) It binds the selected goal constraint and its corresponding current entity to nodes and s. in the transition graph. Lastly, it identifies the edge connecting those two nodes and executes the action tagged to that edge in the environment.
goal constraints hy must themselves be inferred from the current and goal observations or and Oy. requiring the agent to infer both what to do and how to do it purely from its sensorimotor interface. Our approach takes four steps, summarized in Alg. P and Fig. In the first step, we use dSLATE to infer h, and h, from og and o, (e.g. the positions and types of all objects in the initial and goal images). In the second step (Fig. Hp). because of the permutation symmetry among entities, we find bipartite matching that matches each entities in h2 with corresponding entity in Ae that shares the same type and permute the indices l of h, to match those of hg. We implement a function align that uses the Hungarian algorithm to perform this matching over (..,) and (2)... ..... K ) with Euclidean distance as the matching cost. The third step selects which goal constraint h to satisfy next (Fig. D)) W implement this select-constrant constraint hn has the highest difference in state with its counterpart htf which reduces to solving the same argmax problem as isolate with the same distance function used in sol lale. The last step chooses action given the chosen goal constraint hy and its counterpart ht, by binding hy and h to the graph based on their state components and returning the action tagged the edge between their respective nodes (Fig, 4d). If an edge does not exist between the inferred nodes, then we simply take random action.
procedure by determining which
EXPERIMENTS
We have proposed NCS solution the object rearrangement problem that addresses two chal- lenges: NCS addresses the correspondence problem by learning a factorized object-centric world model with dSLATE and it addresses the combinatorial problem by abstracting entity representa- tions into a queryable state transition graph. Now we test NCS's efficacy in solving both problems.
The key question is whether NCS is better than sfte-of-theeart offline RL algorithms in generalizing over betiaatoriilly-stttturrre task spaces from perceptual input. As stated in S the crucial test for answering this question is to evaluate all methods solving new rearrangement problems with a disjoint set of object configurations from those encountered during training. The most straightfor- ward way to find a disjoint subset of the combinatorial space is evaluate with novel number of

---

Published as a conference paper at ICLR 2023
objects. We compare NCS several offline RL baselines and ablations two rearrangement envi- ronments and find a significant gap in performance between our method and the next best method.
Environments. In block-rearrange (Fig. 5) all objects are the same size, shape, and orientation. S covers 16 locations in grid. Z is the con- tinuous space of red-green-blue values from O to 1. robogym-rearrangs (Fig. pD) is adapte from the OpenAT (2020) rearrange environment and removes the assumptions from block-rearrane same size, shape, and orientation. The objects are uniformly sampled from of 94 meshes consisting of the YCB object set Calli et al. (213)) and a set of basic geometric shapes, with colors sampled from set of 13. Although locations not pre-defined in aubogym-rearrang as in block-rearrange, in practice there is a limit to the number of ways to arrange objects on the table still be visible to the camera, which makes the bisimulation still reasonable assumption here. For block- we use the SA attention mask as the state s, and for aabogym--earrange we use the action-dependent part of the SA slot X as the state
that all objects have the
initial obs
goad obs
(a) rearkaree
initial
goal abs %
(h) robogyn-rearrange
Our environ- ments are block-reurrange and robogym-rearrange Fig. 5a shows complete specification of goal con- straints; Fig. [5b] shows a partial specification that only specifies the desired locations for two objects.
Figure
Experimental setup. We evaluate settings: complete and partial. In the complete setting, the goal image shows all objects in new locations. The partial setting is underspecified: only a subset of objects have asso- ciated goal constraints (Fig. 5p). In block- rearranea all constraints are unsatisfied in the start state. In robogym--earrangee rearrange, four constraints are unsatisfied in the start state. Our metric is the fractional success rate, the average change in the number of satisfied constraints divided by the number of initially unsatisfied constraints.
The experiences buffer consists 5000 trajectories showing objects. We evaluate on 4-7 objects for 100 episodes across 10 seeds. Even if we assume full access to the underlying state space, the task spaces are enormous: with | object locations and objects. the number of possible trajectories over object configurations of timesteps is (xx) x (A (IS) -,)), which amounts to searching over more than 10 possible trajectories for the complete specification setting of block-rearrange with k 7 objects (see Appdx. E for derivation). Our setting of assuming access only pixels makes the problem even harder.
Baselines. The claim of this paper that, for object rearrangement, (1) object-centric methods fare better than monolithically-struutu better than parametric planning for object rearrangement and (3) factorized graph search over state transitions of individual entities fares better than a nontfactorized graph search over state transitions entire entity-sets. To test (1), compare with state-of-the-art pixel-based behavior cloning (BC) and implicit Q-learning (IQL) implementations based off of Kostrikov (2021). To test (2), compare against version of object-centric model predictive control (MPC) (Veerapaneni al. 2020) that uses the cross entropy method over dSLATE rollouts. To test (3), we compare against an ablation (abbrv. NE, for (ron-factorize") that constructs a graph with state transitions of entity-sets than of individual states. Our last baseline just takes random actions (Rand).
offline RL methods (2) non non-parametra graph search fares
5.1 RESULTS
Figure shows that NCS performs significantly better than all baselines (about 5-10x improve- ment), thereby refuting the alternatives to our claims. Most of the baselines perform no better or only slightly better than random. We observe that it is indeed difficult to perform shooting-bart based planning with an entity-centric world model trained to predict a single step forward (Janner et al., 2019:: the MPC baseline performs poorly because its rollouts are poor, and it is significantly more computationally expensive run (11 hours instead of 20 minutes). We also observe that the NF ablation performs poorly, showing the importance of factorizing the non-paramttric graph search. Additional results are in the Appendix.
5.2 ANALYSIS
Having quantitatively shown the relative strength of NCS in combinatorial generalization from pix- els, examine how our key design choices of (1) factorizing entity representations into action- invariant and action-dependent features and (2) querying a state transition graph constructed from action-dependent features contribute NCS's behavior and performance. Is copying the entity type
8

---

Published as a conference paper at ICLR 2023
Table 1: This table compares NCS with various baselines in the complete and partial evaluation settings of block-rearrange and rebrgym-rearraage. The methods trained on 4 objects and evaluated on generalizing 5, 6, and 7 objects. We report the fractional success rate, with a standard error computed over 10 seeds.
(a) reack-re, complete specification.
(b) block-rearrange, complete specification.
Method
7
Method
(ours)
0.94 2000
1
1
1:000
(ours)
0.89
0.56 1.0
0.78
0.70 1.
Rand
0.06 am
0.07 1.00
0.07 1 1
0.08 1010
Rand
0.06 1ove
0.08 1.00
0.08
1001
0.08 om
0.16 0.0
100
1010
MPC
0.13
0.10
0.08
NF
0.07 1.00
0.06 1000
0.07 1 boz
0.08 1
NF
0.06 LUG
0.07 1.00
0.08 000
0.07 o
IQL
0.07 1 0m
w.0
1
IQL.
0.01
0.07
4.000
0.05
0.05
BC
0.03 0.00
0.02 0.00
0.01 + 0000
0.01 0:000
BC
0.05 5m0
0.04
0.000
0.03 + 0:00
0.03 000
(c) robogym-rearrange, complete specification.
(d) rabogym-reerrange. partial specification.
Method
4
5
Method
7
NCS (ours)
cm
0.49
0.41
NCS (omrs)
0.47
0.33 1:10
0..27 1:00
0.22
Rand
0.01 000
0.01 600
0.00 0.00
0.00 0.0
Rand
0.005
0.001
0.002
0.001
MPC
000
0.00
0.00
MPC
NF
0.01 1,000
0.01 1.0m
0.00 1:000
0.00 0000
0.001
EX0L
0.00
IQL.
0.00
0.00
0.00
0.00
IQ1.
0.00 0.0
0.00 1:10
1:00
2:000
BC
0.00 100000
0.00 1.000
0.00 0.00
0.00 1,000
BC
0.00 0.0
0.00
1:000
0.00
0.0
0.00 0.000
during latent prediction dSLATE does sufficient for disentangling the location and appearance of objects into the state and type respectively? Does dSLATE learn to sparsely modify only the entity that corresponds to the moved object in the sensorimotor transition, such that the nodes of the state transition graph meaningfully can be reused across entities? These are nontrivial capabilities because NCS self-supervised on only the experience buffer.
Fig. 1. which visualizes the align, and bind functions of NCS on robogym-rearrange, suggests that, at least for the simplified setting consider, the answer to both questions is yes. NCS has learned represent different objects in different slots and construct a graph whose nodes capture location information. Fig. shows t-SNE (Van der Maaten & Hinton 2008 plot that clusters entities inferred from the robogym environment. Because we have not provided supervision on what states should represent, we observe there are multiple cluster indices that map onto similar groups of points. This reveals that multiple different regions appear to be modeling similar states. We also tried merging redundant clusters, but found that this did not improve quantitative performance.
Partitioning states from the training equivalance classes
select-constraiing
30
20
-10
-20
-10-
40
-50
-20
-10
10
Figure 6: Nodes as equivalent classes over states. We show a clustering of states inferred for rearrange, where each cluster centroid is treated node in our transition graph. A sub- of clusters are labeled with an attention mask computed by averaging the slot attention masks for the entities associated with the cluster.
6 DISCUSSION
Object rearrangement offers an intuitive setting for studying how an agent can learn reusable abstractions from its sensorimotor experience. This paper takes a first step toward connecting the world of symbolic planning with human-defined abstrac- tions and the world of representation learning with deep networks by introducing NCS. NCS is method for controlling an agent that plans over and acts with state transition graph constructed with entity representations learned from raw sensorimotor transitions, without any other supervision. We showed that factorizing the entity representation into action-invariant and action-dependent features are important for solving the correspondence and combinatorial problems that make the object rear- rangement difficult, and enable NCS to significantly outperform existing methods on combinatorial generalization in object rearrangement. The implementation of NCS provides proof-of-concept for how learning reusable abstractions might be done, which we hope inspires future work to engineer methods like NCS for real-world settings.
9

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGMENTS
This work was done while MC was an intern at Meta AI. We would like to thank Leslie Kaelbling for valuable feedback and Yash Sharma and Yilun Du for valuable discussions. This material supported in part by the Fannie and John Hertz Foundation, as well as with ONR grant #N00014- 18-1-2873.
REFERENCES
David Abel, D. Ellis Hershkowitz, Gabriel Barth-Maron, Stephen Brawner, Kevin O'Farrell, James MacGlashan, and Stefanic Tellex. Goal-based action priors. In ICAPS, pp. 306-314. AAAI Press, 2015.
Victor Bapst, Alvaro Sanchez-Gonzales Carl Doersch, Kimberly Stachenfeld, Pushmeet Kohli, Pe- ter Battaglia, and Jessica Hamrick. Structured agents for physical construction. In International conference on machine learning, pp. 464-474. PMLR, 2019.
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 arXiv:2011.01975 2020.
Peter W Battaglia, Jessica B Hamrick, Victor Bapst, Alvaro Sanchez-Gonzalez, Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, et al. Relational inductive biases, deep learning, and graph networks. arXiv preprint 1rXiv:1806.01220 2018.
Craig Boutilier, Richard Dearden, and Moises Goldszmidt. Exploiting structure in policy construc- tion. In Proceedings of the 14th International Joint Conference on Artificial Intelligence Volume 2, IJCAI'95, pp. 1104-1111, 1. San Francisco, CA, USA, 1995. Morgan Kaufmann Publishers Inc.
Craig Boutilier, Richard Dearden, and Moisés Goldszmidt. Stochastic dynamic programming with factored representations. Artif. Intell., 121(1-2):99107, 2000.
Berk Calli, Arjun Singh, Aaron Walsman, Siddhartha Srinivasa, Pieter Abbeel, and Aaron M Dollar. The yeb object and model set: Towards common benchmarks for manipulation research. In 2015 international conference advanced robotics (ICAR). pp. 510-517. IEEE, 2015.
Pablo Samuel Castro, Prakash Panangaden, and Doina Precup. Equivalence relations in fully and partially observable markov decision processes. In Twenty-First International Joint Conference on Artificial Intelligence. Citeseer, 2009,
Chien-Yo Yi Chang, De-An Huang, Danfei Xu, Ehsan Adeli, Li Fei-Fei, and Juan Carlos Niebles. Procedure planning in instructional videos. In European Conference on Computer Vision, pp- 334-350 Springer, 2020.
Michael Chang. Thomas L Griffiths, and Sergey Levine. Object representations as fixed points: Training iterative refinement algorithms with implicit differentiation. arXiv preprint arXiv:2207.0072 2022.
Michael B Chang, Tomer Ullman, Antonio Torralba, and Joshua B Tenenbaum. A compositional object-based approach learning physical dynamics. arXiv preprint arXiv:1612.0034 2016.
Erwin Coumans and Yunfei Bai. Pybullet, python module for physics simulation for games, robotics and machine learning. 2016.
Coline Devin, Payam Rowghanian, Chris Vigorito, Will Richards, and Khashayar Rohanimanesh. Self-supervised goal-conditioned pick and place. arXiv preprint arXiv:2008. //466, 2020.
Carlos Diuk, Andre Cohen, and Michael L. Littman. An object-oriented representation for efficient reinforcement learning. In ICML, volume 307 of ACM International Conference Proceeding Series, pp. 240-247. ACM, 2008.
10

---

Published as a conference paper at ICLR 2023
Gamaleldin F Elsayed, Aravindh Mahendran, Sjoerd van Steenkiste, Klaus Greff, Michael C Mozer, and Thomas Kipf. Savi++: Towards end-to-end object-centric learning from real-world videos. arXiv preprint arXhv:2206.07766 2022.
Scott Emmons, Ajay Jain, Misha Laskin, Thanard Kurutach, Pieter Abbeel, and Deepak Pathak. Sparse graphical memory for robust planning, Advances in Neural Information Processing Sys- tems, 33:5251-5262. 2020.
Ben Eysenbach, Russ R Salakhutdinov, and Sergey Levine. Search on the Replay Buffer: Bridging Planning and Reinforcement Learning. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' Alché- Buché-Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL http::tst::ppcceeaiiiiaae
wwwrips....... ec/
Natalia Gardiol and Leslie Kaelbling, Envelope-based Planning in Relational MDPs. In S. Thrun, L. Saul,, and B. Schölkopf (edss), in Neural Information Processing Systems, volume 16. MIT Press, 2003. URL http::/spproeeeiisss...e..s
Tejas Gokhale, Shailaja Sampat, Zhiyuan Fang, Yezhou Yang, and Chitta Baral. Cooking with blocks: A recipe for visual reasoning image-pairs. In Proceedings of the IEEE/CVF Confer- ence on Computer Vision and Pattern Recognition Workshops, pp. 5-8, 2019.
Anirudh Goyal, Alex Lamb, Phanideep Gampa, Philippe Beaudoin, Sergey Levine, Charles Blun- dell, Yoshua Bengio, and Michael Mozer. Object files and schemata: Factorizing declarative and procedural knowledge in dynamical systems. arXiv preprint 1iXX::2006.1622,3 2020.
Anirudh Goyal, Aniket Didolkar, Nan Rosemary Ke, Charles Blundell, Philippe Beaudoin, Nicolas Heess, Michael C Mozer, and Yoshua Bengio. Neural production systems. Advances in Neural Information Processing Systems, 34:25673-25663 2021.
Klaus Greff, Sjoerd Van Steenkiste, and Jürgen Schmidhuber. Neural expectation maximization. arXiv preprint arviv: / 7077700..33... 2017.
Klaus Greff, Raphaël Lopez Kaufman, Rishabh Kabra, Nick Watters, Christopher Burgess, Daniel Zoran, Loic Matthey, Matthew Botvinick, and Alexander Lerchner. Multi-object representation learning with iterative variational inference. In International Conference on Machine Learning, pp. 2424 2433. PMLR, 2019.
Klaus Greff, Sjoerd van Steenkiste, and Jürgen Schmidhuber. On the binding problem in artificial neural networks. arXiv preprint arXiv:2012.052088 2020.
C. Guestrin, D. Koller, R. Parr, and S. Venkataraman. Efficient Solution Algorithms for Factored MDPs. Journal of Artificial Intelligence Research, 19:399-466 October 2003a. ISSN 1076-9757. doi: 10.1613/jain 1000. URL htttp:://jair...... org/ in//xxidxxxx////// http://///aeeeepppgaa//a////a/aa 10341 tex.ids: guestin2OOElhccenSsuuroagggggii
Carlos Guestrin, Daphne Koller, Chris Gearhart, and Neal Kanodia. Generalizing plans to new environments in relational mdps. Proceedings of the 18th International Joint Conference on Artificial Intelligence, IJCAI'03, pp. 1003-1010, San Francisco, CA, USA, 2003b. Morgan Kauf- mann Publishers Inc.
Philippe Hansen-Fistruch, Amy Zhang, Ashvin Nair, Patrick Yin, and Sergey Levine. Bisimulation makes analogies in goal-conditioned reinforcement learning. arXiv preprint arXiv:2204. 2022.
De-An Huang, Suraj Nair, Danfei Xu, Yuke Zhu, Animesh Garg, Li Fei-Fei, Silvio Savarese, and Juan Carlos Niebles. Neural task graphs: Generalizing to unseen tasks from a single video demon- stration. In Proceedings of the IEEF/CVF conference on computer vision and pattern recognition, pp. 8565-8574, 2019.
Michael Janner. Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model- based policy optimization. Advances in Neural Information Processing Systems, 32, 2019.
11

---

Published as a conference paper at ICLR 2023
Thomas Kipf, Gamaleldin F Elsayed, Aravindh Mahendran, Austin Stone, Sara Sabour, Georg Heigold, Rico Jonschkowski, Alexey Dosovitskiy, and Klaus Greff. Conditional object-centric learning from video. arXiv preprint arXtv:2111.12594, 2021.
Iyya Kostrikov. JAXRL: Implementationn of Reinforcement Learning algorithms in JAX, 10 2021.
URL httphtts:///
Tejas D Kulkarni, Ankush Gupta, Catalin lonescu, Sebastian Borgeaud, Malcolm Reynolds, Andrew Zisserman, and Volodymyr Mnih. Unsupervised learning of object keypoints for perception and control. NeurIPS, 2019.
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tuto- rial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643 2020.
Martina Lippi, Petra Poklukar, Michael C Welle, Anastasiia Varava, Hang Yin, Alessandro Marino, and Danica Kragic. Latent space roadmap for visual action planning of deformable and rigid ob- ject manipulation. In 2020 /EEE/RS International Conference on Intelligent Robots and Systems (IROS), pp. 5619-5626. IEEE, 2020.
Francesco Locatello, Dirk Weissenborn, Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. Object-centric learning with slot atten- tion. arXiv preprint arXiv:2006. 15055, 2020a.
Francesco Locatello, Dirk Weissenborn, Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. Object-centric learning with slot atten- tion. In NeurIPS, 2020b. URL /ttp:://rrceeddng wwwpkinngnneurrpppp wwwripe. cttpp/ppppccpppp//////////
html]
OpenAl. Robogym. httpps//giihhbb..cmm/errr//////p
2020.
Luis Pineda, Brandon Amos, Amy Zhang, Nathan O. Lambert, and Roberto Calandra. Mbrl-lib: A modular library for model-based reinforcement learning. Arxiv, 2021. URL bttps://axxiv
org/abs/ 20000./214...0 10159
Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language under- standing by generative pre-training, 2018.
Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. arXiv preprint arXiv:2102. /2092, 2021.
Gautam Singh, Fei Deng, and Sungjin Ahn. Illiterate DALL-E learns to compose. arXiv preprint arXiv:2110.11055, 2021.
Gautam Singh, Fei Deng, and Sungjin Ahn. Illiterate DALL-e learns compose. In Interna- tional Conference on Learning Representations, 2022a. URL https://ppenreviey net/
"orum??d=hooyyv/0ee
Gautam Singh, Yi-Fu Wu, and Sungjin Ahn. Simple unsupervised object-centric learning for com- plex and naturalistic videos. arXiv preprint 120.0:22055.4065, 2022b.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ international conference on intelligent robots and systems, pp. 5026-5033. IEEE, 2012.
Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine learning research, 9(11), 2008
Sjoerd Van Steenkiste, Michael Chang, Klaus Greff, and Jürgen Schmidhuber. Relational neural ex- pectation maximization: Unsupervised discovery of objects and their interactions. arXiv preprint 10X2::3802.11333. 2018.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural informa- fion processing systems, 30, 2017.
12

---

Published as a conference paper at ICLR 2023
Rishi Veerapaneni, John D 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 Robot Learning, pp. 1439-1456. PMLR, 2020.
C. Wang, S. Joshi, and R. Khardon. First order decision diagrams for relational mdps. Journal of Artificial Intelligence Research, 31:431-472, Mar 2008. ISSN 1076-9757. doi: 10.16 3/jai23// URL http: doi org/ 10. 1et.161333313 2489
Daniel Weld. Solving Relational MDPs with First-Order Machine Learning. pp. 8.
Ge Yang, Amy Zhang, Ari S. Morcos, Joelle Pincau, Pieter Abbeel, and Roberto Calandra. Plan2vec: Unsupervised Representation Learning by Latent Plans. In Proceedings The 2nd Annual Con- ference on Learning for Dynamics and Control, volume 120 of Proceedings of Machine Learning Research, pp. 1-12, 2020.
Shuo Yang, Wei Zhang, Ran Song, Jiyu Cheng, and Yibin Li. Learning multi-object dense descriptor for autonomous soal-conditioned grasping. IEEE Robotics and Automation Letters, 6(2):4109- 4116, 2021.
Andrii Zadaianchuk, Georg Martius, and Fanny Yang. Self-supervised reinforcement learning with independently controllable subgoals. In Conference on Robot Learning, pp. 384-394. PMLR, 2022.
Amy Zhang, Adam Lerer, Sainbayar Sukhbaatar, Rob Fergus, and Arthur Szlam. Composable Plan- ning with Attributes. In Proceedings of the 35th International Conference Machine Learning, ICML 2018, Stockholmsmüssan, Stockholm, Sweden, July 10-15, 2018, volume 80, pp. 5837- 5846. JMLR.org, 2018.
Lunjun Zhang, Ge Yang, and Bradly Stadie. World model as graph: Learning latent landmarks for planning. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, m.t..rress
pp. 12611-12620 PMLR, 18-24 Jul 2021. URL
html
Daniel Zoran, Rishabh Kabra, Alexander Lerchner, and Danilo J Rezende. Parts: Unsupervised seg- mentation with slots, attention and independence maximization. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 10439-10447, 2021.
13


Reviews:
[{"review_id": "fII_EWcp_pj", "reviewer": "Reviewer_wsHP", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper aims to tackle the combinatorial generalization problem in the offline sensorimotor object rearrangement task. By homogenizing the dynamics of different objects, the authors propose a framework named NCS that trains the correspondence matching module and plans a path from the initial state to the goal state in a non-parametric manner. Results on the proposed environments demonstrate that the proposed method significantly outperforms the offline RL baselines.", "summary_of_the_review": "Overall, the authors touch on an important problem in machine learning. Improving the scalability of the non-parametric methods is also appreciated. However, the assumptions made in the paper are strong and there is much room to improve the writing of this paper.\nSo I lean towards rejecting this paper.", "strength_and_weaknesses": "Strengths:\n1. This paper studies an important problem, combinatorial generalization.\n2. The authors aim to improve the scalability issue of the previous neural symbolic approaches.\n\nWeaknesses:\n1. This paper makes strong assumptions that limit their scope. For example, in the heterogeneous scenario where objects are significantly different in scale (e.g. side length), we cannot directly transfer the knowledge (e.g., offline transition) of a large object to a smaller one.\n2. The planning and control algorithms are not clear. For instance, how to search for a reasonable plan without conflicts. I also cannot figure out how to align the entities concretely from only one sentence minimizes the matching cost between z_t and z_g.'. I believe this section is more important than the training details of the submodules. Thus, I recommend that the authors re-organize the paper to emphasize this section more.\n3. There are many typos in this paper. I recommend that authors improve their writing, especially in the method section.", "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": "Wmpu1Ptqo-", "reviewer": "Reviewer_UuMC", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper studies the problem of object rearrangement and focuses particularly on the design of an object-centric planner capable of solving rearrangement tasks requiring multi-step reasoning and systematic generalization. A method, NCS, for combining object-centric representation learning with non-parametric planning is introduced. NCS uses a two-level abstract hierarchy for planning, where a set-structured first level corresponds to sets of entities, and the graph-structured second level corresponds to transitions between “abstracted” entities. Experiments on a few synthetic object rearrangement tasks demonstrate promising performance compared to offline MBRL and object-centric shooting-based planners.\n", "summary_of_the_review": "While the main underlying idea is solid and the general research problem is important and relevant to the community, the technical contribution of the paper appears to be of limited significance. Further improvements to the clarity and quality of the paper may reveal otherwise, but this could require significant revisions.\n\n---\nUpdate: After reading the author's response, they have addressed many of the comments and concerns that were brought up. I have increased my score from 5 to 6.", "strength_and_weaknesses": "Strengths\n=======\nThe main high-level idea, which I take to be the proposition of building a higher level of abstraction on top of object-centric representations upon which it is easier to conduct long-term planning --- is solid. NCS appears to be a promising first step based on the decent empirical results on synthetic tasks.\n\nWeaknesses\n=======\nGiven that the key contribution of this paper appears, to me, to be the NCS planner as a whole (specifically, as a demonstration of how object-centric representations can be leveraged for object rearrangement tasks), NCS’s design appears to rely too heavily on certain strong assumptions that cater to the synthetic rearrangement environments it is demonstrated on. I believe this makes it difficult to see how NCS’s technical designs can inform future frameworks applied on more challenging environments. In particular, the treatment of object identity is highly simplistic; NCS copies the identity of each object discovered in the first step forward through time. This requires an assumption of no occlusion and no objects being introduced or leaving during an episode. \n\nI think it's important to discuss and compare against this missing related work [1] --- the question of why non-parametric planning should be conducted over abstract object-centric states versus planning directly on abstract subgoals (which may be less restrictive of an assumption?) is an interesting one.\n\nReferences\n=======\n1. Zhang, Lunjun, Ge Yang, and Bradly C. Stadie. \"World model as a graph: Learning latent landmarks for planning.\" International Conference on Machine Learning. PMLR, 2021.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "NyU9sU3lb05", "reviewer": "Reviewer_vW5u", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper is about the task of object rearrangement. The paper proposes a neurosymbolic approach for solving this problem.\nThe method is described as a hierarchical abstraction approach-- object representations are learned, and a factorized transition graph is constructed over these representations.  The method, \"Neural Constraint Satisfaction\" disentangles representation learning from planning and control, therefore leading to a modularized approach.  Experiments are performed on three block-style datasets.\n", "summary_of_the_review": "An effective neurosymbolic method for object re-arrangement, effective across many domains.  Real-world applicability (i.e. with complex images) has not been explored -- all the experimental settings are in synthetic environments, which in my opinion, limits the impact of this paper", "strength_and_weaknesses": "*Strengths*\n- Object rearrangement using image inputs is a very interesting problem and I can see a lot of domains where it could be useful: robotics, planning, visual reasoning, navigation / RL to name a few.\n- The paper is well written and does a good job of establishing the preliminaries and the method (algo 1, 2, and the figures are especially useful in this regard)\n- experiments appear to be extensive on three datasets/environments (although, note that I am not an expert in RL/planning), with recent baselines\n- The neuro-symbolic system proposed has the potential to be applicable in various domains, by simply swapping the visual module.\n\n*Weaknesses*\n- One weakness (which is fixable) is that IMO it would be better to introduce the problem statement very early on so that readers get a taste of what's to come later.  The introduction is dedicated to mostly explaining the importance and difficulty of object rearrangement, without specifying what input-output spaces are being considered -- this comes very late.\n- It would be interesting to see the method being applied to real world datasets -- i.e. images with real objects (see my suggestion 2 for prospective datasets). Perhaps in such settings, more sophisticated feature learning / entity abstraction methods would need to be used.\n- ablation study/analyses to study each module could enhance the evidence for the efficacy of the method\n\n*Suggestions*\n1. Increase font size in all figures. Tables are fine.\n2. References: there are two related papers that fall under the problem formulation given in Sec 2 Page 2:\n  - https://openaccess.thecvf.com/content_CVPRW_2019/papers/Vision_Meets_Cognition_Camera_Ready/Gokhale_Cooking_With_Blocks__A_Recipe_for_Visual_Reasoning_on_CVPRW_2019_paper.pdf This paper is about an object rearrangement task in the \"blocksworld\" setting but with real images. This one seems to be very similar to _\"block-stack\"_ from this paper.\n  - https://www.ecva.net/papers/eccv_2020/papers_ECCV/papers/123560324.pdf procedural planning in videos -- this paper involves complex actions beyond simple rearrangement, but in real-world videos like cooking etc.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster