Energy-based Models are Zero-Shot Planners for Compositional Scene Rearrangement

I Introduction

We consider the scene arrangement task shown in Figure 1. Given a visual scene and an instruction regarding object spatial relations, the robot is tasked to rearrange the objects to their instructed configuration. Our focus is on strong generalization to longer instructions with novel predicate compositions, as well as to scene arrangements that involve novel objects and backgrounds.

We propose generating goal scene configurations corresponding to language instructions by minimizing a composition of energy functions over object spatial locations, where each energy function corresponds to a language concept (predicate) in the instruction. We represent each language concept as an n-ary energy function over relative object poses and other static attributes, such as object size. We train these predicate energy functions to optimize object poses starting from randomly sampled object arrangements through Langevin dynamics minimization , using a handful of examples of visual scenes paired with single predicate captions. Energy functions can be binary for two-object concepts such as left of and in front of, or multi-ary for concepts that describe arrangements for sets of objects, such as line or circle. We show that gradient descent on the sum of predicate energy functions, each one involving different subsets of objects, generates a configuration that jointly satisfies all predicates, if this configuration exists, as shown in Figure 1.

We propose a robot learning framework that harnesses minimization of compositions of energy functions to generate instruction-compatible object configurations for robot scene rearrangement. A neural semantic parser is trained to map the input instruction to a set of predicates and corresponding energy functions, and an open-vocabulary visual-language grounding model grounds their arguments to objects in the scene, as shown in Figure 2. Gradient descent on the sum of energies with respect to the objects’ spatial coordinates computes the final object locations that best satisfy the set of spatial constraints expressed in the instruction. Given the predicted object goal locations, we use vision-based pick-and-place policies that condition on the visual patch around the predicted pick and place locations to rearrange the objects . We call our framework Scene Rearrangement via Energy Minimization (SREM).

We test SREM in scene rearrangement of tabletop environments on simulation benchmarks of previous works , as well as on new benchmarks we contribute that involve compositional instructions. We curate multiple train and test splits to test out-of-distribution generalization with respect to (i) longer instructions with more predicates, (ii) novel objects and (iii) novel background colors. We show SREM generalizes zero-shot to complex predicate compositions, such as “put all red blocks in a circle in the plate” while trained from single predicate examples, such as “an apple inside the plate” and “a circle of blocks”. We show SREM generalizes to real-world scene rearrangement without any fine-tuning, thanks to the object abstractions it operates on. We compare our model against state-of-the-art language-to-action policies as well as Large Language Model planners and show it dramatically outperforms both, especially for long complicated instructions. We ablate each component of our model and evaluate contributions of perception, semantic parsing, goal generation and low-level policy modules to performance.

In summary, our contributions are: (i) A novel energy-based object-centric planning framework for zero-shot compositional language-conditioned goal scene generation. (ii) A modular system for instruction-guided robot scene rearrangement that uses semantic parsers, vision-language grounding models, energy-based models for scene generation, and vision-based policies for object manipulation. (iii) A new instruction-guided scene rearrangement benchmark in simulation with compositional language instructions. (iv) Comparisons against state-of-the-art language-to-action policies and LLM planners, and extensive ablations.

Simulation and real-world robot execution videos, as well as our code are publicly available on our website: https://ebmplanner.github.io.

II Related Work

Following instructions for rearranging scenes: Language is a natural means of communicating goals and can easily describe compositions of actions and arrangements , providing more versatile goal descriptions compared to supplying one or more goal images. The latter requires the task to be executed beforehand, which defeats the purpose of instruction . We group methods in the literature in the following broad categories:

End-to-end language to action policies map instructions to actions or to object locations directly. We have found that these reactive policies, despite impressively effective within the training distribution, typically do not generalize to longer instructions, new object classes and attributes or novel backgrounds .

Symbolic planners such as PDDL (Planning Domain Definition Language) planners use predefined symbolic rules and known dynamics models, and infer discrete task plans given an instruction with lookahead logic search . Symbolic planners assume that each state of the world, scene goal and intermediate subgoal can be sufficiently represented in a logical form, using language predicates that describe object spatial relations. These methods predominantly rely on manually-specified symbolic transition rules, planning domains and grounding, which limits their applicability.

Large language models (LLMs) map instructions to language subgoals or program policies with appropriate plan-like prompts. The predicted subgoals interface with low-level short-term policies or skill controllers. LLMs trained from Internet-scale text have shown impressive zero-shot reasoning capabilities for a variety of downstream language tasks when prompted appropriately, without any weight fine-tuning . The scene description is usually provided in a symbolic form as a list of objects present, predicted by open-vocabulary detectors . Recent works of have also fed as input overhead pixel coordinates of objects to inform the LLM’s predictions. The prompts for these methods need to be engineered per family of tasks. It is yet to be shown how the composition of spatial concept functions can emerge in this way.

Language-conditioned scene generation: A large body of work has explored scene generation conditioned on text descriptions . The work of leverages web-scale pre-trained models to generate segmentation masks for each object in the generated goal image. Given an input image, their method generates a text prompt using a captioning model and feeds it to a generative model that outputs a goal image, which is then further parsed into segmentation masks. However, the prompt is limited to contain only names of objects and there is no explicit language-guided spatial reasoning. In this work, we seek to make scene generation useful as goal imagination for robotic spatial reasoning and instruction following. Instead of generating pixel-accurate images, we generate object configurations by abstracting the appearance of object entities. We show this abstraction suffices for a great number of diverse scene rearrangement tasks.

Energy-based models: Our work builds upon existing work on energy-based models (EBMs) . Most similar to our work is that of , which generates and detects spatial concepts with EBMs on images with dots, and , which demonstrate composability of image-centric EBMs for generating face images and images from CLEVR dataset . In this work, we demonstrate zero-shot composability of EBMs over object poses instead of images, and showcase their applicability on spatial reasoning and instruction following for robotic scene rearrangement.

III Method

The architecture of SREM is shown in Figure 2. The model takes as input an RGB-D image of the scene and a language instruction. A semantic parser maps the instruction to a set of spatial predicate energy functions and corresponding referential expressions for their object arguments. An open-vocabulary visual detector grounds the arguments of each energy function to actual objects in the scene. The goal object locations are predicted via gradient descent on the sum of energy functions. Lastly, short-term vision-based pick-and-place policies move the objects to their inferred goal locations. Below, we describe each component in detail.

A library of energy-based models for spatial concepts In our work, a spatial predicate is represented by an energy-based model (EBM) that takes as input $x$ the set of objects that participate in the spatial predicate and maps them to a scalar energy value $E_{\theta}(x)$. An EBM defines a distribution over configurations $x$ that satisfy its concept through the Boltzmann distribution $p_{\theta}(x)\propto e^{-E_{\theta}(x)}$. Low-energy configurations imply satisfaction of the language concept and have high probability. An example of the spatial concept can be generated by optimizing for a low-energy configuration through gradient descent on (part of) the input $x$. We represent each object entity by its 2D overhead centroid coordinates and box size. During gradient descent, we only update the center coordinates and leave box sizes fixed. We consider both binary spatial concepts (in, left of, right of, in front of, behind) as well as multi-ary spatial concepts (circle, line).

Using an EBM, we can sample configurations from $p_{\theta}$, by starting from an initial configuration $x^{0}$ and refining it using Langevin Dynamics :

where $z^{k}$ is random noise, $\lambda$ is an update rate hyperparameter and $\epsilon^{k}$ is a time-dependent hyperparameter that monotonically decreases as $k$ increases. The role of $z^{k}$ and decreasing $\epsilon^{k}$ is to induce noise in optimization and promote exploration, similar to Simulated Annealing . After $K$ iterations, we obtain $x^{-}=x^{K}$. During training, we iterate over Equation 1 $K=30$ times, using $\lambda=1$ and $\epsilon_{k}=5e-3$. During inference, we find that iterating for more, e.g. $K=50$ often leads to better solution. In this case we also linearly decay $\epsilon_{k}$ to 0 for $k>30$.

We learn the parameters $\theta$ of our EBM using a contrastive divergence loss that penalizes energies of examples sampled by the model being lower than energies of ground-truth configuration:

where $x^{+}$ a sample from the data distribution $p_{D}$ and $x^{-}$ a sample drawn from the learned distribution $p_{\theta}$. We additionally use the KL-loss and the L2 regularization proposed in for stable training. At test time, compositions of concepts can be created by simply summing energies of individual constituent concept, as shown in Figures 1 and 2.

We implement two sets of EBMs, a BinaryEBM and a MultiAryEBM for binary (e.g., left of) and multi-ary (e.g., circle) language concepts, respectively. The BinaryEBM expects two object arguments, each represented by its bounding box. We convert the object bounding box to (top-left corner, bottom-right corner) representation. Then we compute the difference between all corners of the two object arguments and concatenate and feed to a multi-layer perceptron (MLP) that outputs a scalar energy value. Note that the energy function only depends on the relative arrangement of the two objects, not their absolute locations. The MultiAryEBM is used for order-invariant concepts of multiple entities, such as shapes. The input is a set of objects, each represented as a point (box center). We subtract the centroid of the configuration from each point and then featurize each object using an MLP. We feed this set of object features to a sequence of four attention layers for contextualization. The refined features are averaged into an 1D vector which is then mapped to a scalar energy using an MLP. We train a separate EBM for each language concept in our vocabulary using corresponding annotated scenes in given demonstrations. Note that annotated scenes suffice to train the energy functions, kinesthetic demonstrations are not necessary, and in practice each EBM can be trained within a few minutes. We provide further implementation details and architecture diagrams for our EBMs in Section VII-A and Figure 4 of the Appendix. We also visualize the energy landscape for various concepts and combinations in VII-E and Figure 6 of the Appendix.

Semantic parsing of instructions into spatial concepts and their arguments. Our parser maps language instructions to instantiations of energy-based models and their arguments. It is a Sequence-to-Tree model with a copying mechanism which allows it to handle a larger vocabulary than the one seen during training. The input to the model is a natural language instruction and the output is a tree. Each tree node is an operation. The three operations supported are i) BinaryEBM which calls a BinaryEBM from our library, ii) MultiAryEBM and iii) VLMGround which calls the visual-language grounding module. Each node has a pointer to the arguments of the operation, language concepts for EBM calls, e.g., behind, and noun phrases for grounding model calls, e.g., “the green cube”. Nodes in the parsing tree may also have children nodes, which imply nested execution of the corresponding operations. The input utterance is encoded using a pre-trained RoBERTa encoder , giving a sequence of contextualized word embeddings and a global representation of the full utterance. Then, a decoder is iteratively employed to i) decode an operation, ii) condition on this operation to decode or copy the arguments for this operation, iii) add one (or more) children node(s). For example, the instruction “a circle of cubes inside the plate” is mapped to a sum of energy functions where each object of the multi-ary concept circle participates in the constraining binary concept in:

We train our semantic parser on the instructions of all training demonstrations of all tasks jointly, as well as on synthesized instructions paired with programs, each with 1-7 predicates, that we generate by sampling from a grammar, similar to previous works . For more details on the domain-specific language of our parser and the arguments for each operation see Section VII-A and Table VIII in our Appendix.

We ground noun phrases predicted by our parser with an off-the-shelf language grounding model , which operates as an open-vocabulary detector. The input is the noun phrase, e.g., “the blue cube” and the image, while the output is the boxes of all object instances that match the noun phrase. The open-vocabulary detector has been pre-trained for object detection and referential grounding on MS COCO , Flickr30k and Visual Genome . We finetune the publicly available code of on our training data of all tasks jointly.

Short-term vision-based manipulation skills We use short-term manipulation policies built upon Transporter Networks to move the obejcts to their predicted locations. Transporter Networks take as input one or more RGB-D images, reproject them to the overhead birds-eye-view, and predict two robot gripper poses: i) a pick pose and ii) a pick-conditioned placement pose. These networks can model any behaviour that can be effectively represented as two consecutive poses for the robot gripper, such as pushing, sweeping, rearranging ropes, folding, and so on – for more details please refer to .

We modify Transporter Networks to take as input a small image RGB-D patch, instead of a complete image view. Specifically, we consider as input the image patches around the object pick and object goal locations predicted by our visual grounding and energy-based minimization modules respectively. In this way, the low-level policies know roughly what to pick and where to place it, and only locally optimize over the best pick location, as well as the gripper’s relative rotation, within an object of interest, or placement location, at a particular part of the scene, respectively. We show in our ablations (Table VII) that using learning-based pick-and-place policies helps performance, even if the search space is limited thanks to grounding and goal imagination. We train Transporter Networks from scratch on all our pick-and-place demonstration datasets jointly.

Termination of execution: SREM generates a goal scene by optimizing the relative poses of the objects mentioned in the instruction. We estimate how many objects should be moved by comparing the detected bounding box (by the language grounding model) and the optimized bounding box (by the EBM). For non-compositional tasks that involve binary concepts, we inject the prior that one object is fixed. Then we take as many actions as the number of objects the EBM moved.

Closed-loop execution: SREM first generates a goal scene from the input instruction and then executes it. After execution, we re-detect all relevant objects using our VLM-grounder module to check if they are close to their predicted goal locations. If the re-detected object’s bounding box and initially predicted goal bounding box intersect over a certain IoU threshold, we consider the goal to be successfully executed. If we fail to reach the goal, we call again our vision based policies using the current scene configuration. Comparing the post-execution object configuration with the initially imagined goal scene allows to track progress and estimate goal completion as we show in the experimental section and in Section VII-C and Table IX of the Appendix.

IV Experiments

We test SREM in its ability to follow language instructions for rearrangement of tabletop scenes in simulation and in the real world. We compare our model against LLM planners and end-to-end language-to-action policies . Our experiments aim to answer the following questions:

How does SREM compare to LLM planners in predicting scene configurations from instructions? (Section IV-A)

How does SREM compare to state-of-the-art language-to-action policies for rearranging scenes? How does their relative performance change with varying instruction length and varying amount of training data? (Section IV-B)

How does SREM generalize to novel objects, object colors and background colors, compared to an end-to-end language-to-action model? (Section IV-C)

How much do different modules of our framework contribute to performance? (Section IV-D)

Benchmarks: Existing language-conditioned manipulation benchmarks are usually dominated by a single spatial concept like “inside” . To better illustrate the compositionality of spatial concepts, we introduce the following set of benchmarks, implemented with PyBullet:

spatial-relations, containing single pick-and-place instructions with referential expressions in cluttered scenes with distractors, e.g. “Put the cyan cube above the red cylinder”. We consider the relations left of, right of, in front of, behind.

comp-one-step, containing compositional instructions with referential expressions in cluttered scenes with distractors that require one object to be re-located to a particular location, e.g. “put the red bowl to the right of the yellow cube, to the left of the red cylinder, and above blue cylinder”.

comp-group, containing compositional instructions with referential expressions in cluttered scenes with distractors that require multiple objects to be re-located, e.g., “put the grey bowl above the brown cylinder, put the yellow cube to the right of the blue ring, and put the blue ring below the grey bowl”.

shapes, containing instructions for making multi-entity shapes (circles and lines), e.g. “rearrange all red cubes in a circle”.

We further evaluate our model and baselines on four tasks from the CLIPort benchmark , namely put-block-in-bowls, pack-google objects-seq, pack-google objects-group and assemble-kits-seq.

For all tasks we train on either 10 or 100 demos and use the same demos to train all our modules, as discussed in Section III. We test on 50 episodes per task, where we vary the instruction and the initial configuration of objects. For spatial-relations and shapes each concept corresponds to a task, while the composition benchmarks correspond to one task each.

Baselines: We compare SREM to the following baselines:

CLIPort , a model that takes as input an overhead RGB-D image and an instruction and uses pre-trained CLIP language and image encoders to featurize the instruction and RGB image, respectively; then fuses these with depth features to predict pick-and-place actions using the action parametrization of Transporter Networks . The model capitalizes on language-vision associations learnt by the CLIP encoders. We use the publicly available code of . We train one CLIPort model on all tasks of each benchmark, e.g., one model for spatial-relations, a different for comp-group etc. Note that the original CLIPort implementation assumes access to oracle success/failure information based on which the model can retry the task for a fixed budget of steps or stop the execution if oracle confirms that the task is completed. We evaluate the CLIPort model without this oracle retry but still with oracle information of how many minimum steps it needs to take to complete the task, so we force CLIPort to take exactly that number of actions.

LLMplanner, inspired by , an instruction-following scene-rearrangement model that uses an LLM to predict a sequence of subgoals in language form, e.g. “pick the red cube and place it to the right of the blue bowl”. The generated language subgoals are fed as input to language-to-action policies, such as CLIPort. Scene state description is provided as a list of objects in the scene. LLMplanner does not finetune the LLM but instead uses appropriate prompts so that the LLM adapts its behavior in-context and generates similar statements. The prompts include various previous successful interactions between a human user and the model. We design suitable prompts for our introduced benchmarks and use the LLM to decompose a long instruction into simpler ones (see Figure 5 in the Appendix for an example). Then, we feed each generated instruction to a CLIPort model, trained as described earlier. Lastly, for tabletop manipulation tasks in simulation, the LLMPlanner of assumes access to an oracle success/failure detector. The difference in our implementation is that we do not assume any success detector. The execution terminates when all language subgoals have been fed to and handled by CLIPort.

Note that LLMplanner boils down to CLIPort for non-compositional instructions. As such, we compare with LLMplanner only on comp-one-step and comp-group, both in simulation and real world.

Evaluation Metrics: We use the following two evaluation metrics: (i) Task Progress (TP) is the percentage of the referred objects placed in their goal location, e.g. 4/5 = 80.0% for rearranging 4 out of 5 objects specified in the instruction. (ii) Task Completion (TC) rewards the model only if the full rearrangement is complete. For the introduced benchmarks we have oracle reward functions that evaluate whether the task constraints are satisfied.

In this section, we compare spatial reasoning for predicting compositional scene subgoals in a language space versus in an abstract visually grounded space. In this section, to isolate this reasoning ability from nuisance factors of visually localizing the objects and picking them up effectively, we consider oracle object detection, referential grounding and low-level pick-and-place policies. Specifically, we carry out inferred language subgoals from LLMplanner using oracle controllers that relocate an object in the scene such that it satisfies the predicted subgoals. Note that SREM relies on pick-and-place policies that are not language-conditioned, while LLMplanner relies on language-conditioned policies for object re-location. Thus, the oracle control assumption is less realistic in the latter case. We forego this difference for the sake of comparison.

We show quantitative results of SREM and LLMplanner on the comp-one-step and comp-group benchmarks in Table I. Our model outperforms LLMplanner and the performance gap is larger in more complex instructions. To elucidate why an abstract visual space may be preferable for planning, we visualize steps of energy minimization for different instructions in Figure 1 and steps of the execution of the LLM prompted by us to the best of our capability in Figure 3. We can see that SREM trained on single-predicate scenes shows remarkable composability in case of multiple predicates. Language planning on the other hand suffers from the ambiguity of translating geometric concepts to language and vice versa: step-by-step execution of language subgoals does not suffice for the composition of the two subgoals to emerge (Figure 3).

IV-B Spatial scene rearrangement

Simulation: In this section, we compare our model and the baselines in the task of instruction-guided scene rearrangement. We first show results on spatial-relations and shapes in Table II. We largely outperform CLIPort, especially when less training demos are considered.

To evaluate generalization on longer instructions at test time, we show quantitative results in Table III for the benchmarks of comp-one-step and comp-group. We compare our model with CLIPort trained on atomic spatial relations and zero-shot evaluated on compositional benchmarks. We further fine-tune CLIPort on demos from the compositional benchmarks. SREM is not trained on these benchmarks, because the energy functions are already composable, meaning that we can jointly optimize over an arbitrary number of constraints by simply summing the different energy terms. Under all different settings, we significantly outperform all variants of CLIPort and LLMplanner. We also observe that closed-loop execution boosts our performance further.

We additionally show results on the CLIPort benchmark in Table IV. We largely outperform CLIPort on almost all tested tasks. Margins are significantly larger when i) less demos are used and ii) the robot has to interact with objects of unseen colors or classes. Most of the failure cases for our model are due to the language grounding mistakes - in particular for assemble-kits-seq we find that the grounder gets confused between letters and letter holes.

Real World: We test our model on a 7-DoF Franka Emika robot, equipped with a parallel jaw gripper and a top-down Azure Kinect RGB-D camera. We do not perform any real-world finetuning. Our test set contains 10 language-guided tabletop manipulation tasks per setting (Comp-one-step, Comp-group, Circles, Lines). We show quantitative results in Table VI. SREM generalizes to the real world without any real-world training or adaptation thanks to the open-vocabulary detector trained on real-world images, as well as the object abstractions in the predicate EBMs and low-level policy modules. We encourage readers to refer to our supplementary video and our website for more detailed results.

IV-C Generalization analysis

We conduct controlled studies of our model’s generalization across three axes: a) novel colors: we train the models with objects of 7 different colors and evaluate them on objects of 4 unseen colors; b) novel background colors: we train all models on black-colored tables and evaluate on tables of randomly sampled RGB colors; c) novel objects: we train the models on objects of 4 classes and evaluate on rearrangement of 11 novel classes. In each of these settings, we only change one attribute (i.e. object color, background color or object instance) while keeping everything else constant.

We evaluate our model and CLIPort trained on 10 or 100 demos per task on spatial-relations (average performance over all tasks) and composition (average performance over all tasks from comp-one-step and comp-group). The results are summarized in Table-V. We observe that our model maintains high performance across all axes of generalization, independently of the number of training demos.

Our model’s generalization capabilities rely on the open-vocabulary detector and the fact that EBMs and transporter-based low-level execution policy operate on abstracted space in a modular fashion. While CLIPort models can also generalize to novel scenarios by leveraging the CLIP model, the action prediction and perception are completely entangled and hence even if CLIP manages to identify the right objects based on the language, it has trouble predicting the correct pick and place locations.

IV-D Ablations

We show an error analysis of our model in Table-VII. First, we remove the goal generation from SREM (SREM w/o goal generation) by conditioning the place network on the language input instead of the EBM-generated goal image, while keeping the pick network and object grounders identical. We observe a drop of 35.1% in accuracy, underscoring the importance of goal generation. We then remove our executor policy (SREM w/o learnable policies) and instead randomly select pick/place locations inside the bounding box of the relevant object. This results in a drop of 16%, showing the importance of robust low-level policies. We do not remove the grounder and parser since they are necessary for goal generation. We then experiment with oracle visual language grounder (SREM w/ oracle language grounding) that perfectly detects the objects mentioned in the sentence, which results in a performance gain of 5.1%. We finally evaluate with perfect grounding, language parsing and low-level execution (SREM w/ everything oracle except goal) to test the error rate of our goal generator. We obtain an 88.3% accuracy, thus concluding that our goal generator fails in 11.7% cases. For a more detailed error analysis, please refer to VII-C in the Appendix.

IV-E Limitations

Our model presently has the following two limitations: First, it predicts the goal object scene configuration but does not have any knowledge regarding temporal ordering constraints on object manipulation execution implied by physics. For example, our model can predict a stack of multiple objects on top of one another but cannot suggest which object needs to be moved first. One solution to this problem is to heuristically pick the order based on objects that are closer to the floor in the predicted scene configuration. However, more explicit encoding of physics priors are important to also identify if the generated configuration is stable or not. A promising direction is to model physics-based constraints as additional energy constraints, and obtain optimization gradients by leveraging either differentiable physics simulators or learned dynamics models . Second, our EBMs are currently parametrized by object locations and sizes, but different tasks need different abstractions. Manipulation of articulated objects, fluids, deformable objects or granular materials, would require finer-grained parametrization in both space and time. Furthermore, even for rigid objects, many tasks would require finer in-space parametrization, e.g., it would be useful to know a set of points in the perimeter of a plate as opposed to solely representing its bounding box for accurately placing things inside it. Considering EBMs over keypoint or object part graphs is a direct avenue for future work.

V Conclusion

We introduce SREM, a modular robot learning framework for instruction-guided scene rearrangement that maps instructions to object scene configurations via compositional energy minimization over object spatial coordinates. We test our model in diverse tabletop manipulation tasks in simulation and in the real world. Our model outperforms state-of-the-art end-to-end language-to-action policies, and LLM-based instruction following methods both in in- and out-of-distribution settings, and across varying amount of supervision. We contribute a new scene rearrangement benchmark that contains more compositional language instructions than previous works, which we make publicly available to the community. Our work shows that a handful of visually-grounded examples suffice to learn energy-based spatial language concepts that can be composed to infer novel instructed scene arrangements, in long and complex compositional instructions.

VI Acknowledgements

This work is supported by Sony AI, NSF award No 1849287, DARPA Machine Common Sense, an Amazon faculty award, and an NSF CAREER award.

References

VII Appendix

In Section VII-A we give implementation details for the components of our method; in Section VII-B we present in more detail the evaluation metrics for the newly-introduced tasks; in Section VII-C we show the effectiveness of closed-loop execution for predicting the success or failure of execution and present a detailed error analysis; we show an example prompt used for our LLMPlanner baseline in Section VII-D; we visualize the learned energy landscapes in Section VII-E; in Section VII-F we include additional related work on constraint-guided layout optimization.

Energy-based Models: The architectures of the BinaryEBM and MultiAryEBM are shown in Figure 4a and b respectively. We train a separate network for each concept in our library using the same demos that are used to train the other modules. We augment and repeat the samples multiple times to create an artificially larger dataset during training. We use Adam optimizer, learning rate $1e-4$, batch size 128.

Buffer: During training, we fill a buffer of previously generated examples, following . The buffer is updated after each training iteration to store at most 100000 generated examples. When the buffer is full, we randomly replace older examples with incoming new ones. We initialize $x_{0}$ for Equation 1 by sampling from the buffer 70% of the times or loading from the data loader 30% of the times.

Extension to tasks with 3D information or pose: An important design choice is what parameters of the input we should be able to edit. We inject the prior knowledge that on our manipulation domain the objects move without deformations, so we fix their sizes and update only their positions. Our EBMs operate on boxes so that they can abstract relative placement without any need for object class or shape information. However, EBMs can be easily extended to optimize other types of representations, such as 3D bounding boxes or pose.

We train EBMs that optimize over 3D locations for relative placement. The architecture is shown in Figure 4c. We adapt the BinaryEBM to represent boxes as $(xyz_{min},xyz_{max})$ and then compute the relative representations as in the 2D case. We optimize for one 3D relation, “on”. For pose-aware EBMs, we adapt the MultiAryEBM to represent objects as $(x_{center},y_{center},\theta)$, where $\theta$ is the rotation wrt the world frame. We then simply change the first linear layer of the MultiAryEBM to map the new input tuple to a 128-d feature vector. The architecture is shown in Figure 4d. We show qualitative results that compose these EBMs into new concepts on our website.

Domain-Specific Language: We design a Domain-Specific Language (DSL) which extends the DSL of NS-CL (designed for visual question answering in CLEVR ) to further predict scene generations, e.g. “put all brown shoes in the green box”. Detailed description of our DSL can be found in Table VIII.

Semantic Parser We construct program annotations for the language instructions of the training demos by mapping them to our DSL (rule-based). We then train our parser on all instructions using Adam optimizer with learning rate $1e-3$ and batch size 32. The same parser weights are used across all tasks. For compositional tasks, we train the parser on the descriptions from the demos that are used to finetune our baselines. The parser is the only part of our framework that needs to be updated to handle longer instructions.

Visual-language Grounder: We finetune BEAUTY-DETR on the scenes of the training demos in simulation. BEAUTY-DETR is an encoder-decoder Detection Transformer that takes as input an image and a referential language expression and maps word spans to image regions (bounding boxes). The original BEAUTY-DETR implementation uses an additional box stream of object proposals generated by an off-the-shelf object detector. We use the variant without this box stream for simplicity. BEAUTY-DETR has been trained on real-world images. We finetune it using the weights and hyperparameters from the publicly available code of .

Short-term Manipulation Skills Our low-level policy network is based on Transporter Networks. Transporter Networks decompose a given task into a sequence of pick-and-place actions. Given an overhead image, the model predicts a pick location and then conditions on it to predict a place location and gripper pose. The original implementation of Transporter Networks supports training with batch size 1 only. We implement a batch-supporting version and find it more stable. We use batch size 8 and follow the original paper in other hyperparameter values.

VII-B Benchmark Generation

We extend the Ravens benchmark for spatial reasoning in the PyBullet simulator. For each benchmark, we write a template sentence (e.g. “Arrange OBJ1 into a circle”) and then randomly select valid objects and colors from a pre-defined list. To test generalization, we include novel colors or novel objects in the evaluation set. Once the sentence is generated, we programatically define valid regions which satisfy the relation and then sample empty locations from it to specify object goal locations. We start by placing all objects randomly in the scene. Then, an oracle hand-designed policy picks and places the objects to the desired locations and returns a demo trajectory which consists of raw RGB-D images and pick-and-place locations. These can be used then to train a behaviour cloning policy similar to CLIPort.

Evaluation Metrics for Rearrangement Tasks To evaluate make-a-circle task, we fit a best-fit circle for the final configuration predicted by SREM. To do this, we consider the centers of the bounding boxes as points. Then, we compute the centroid of those points and the distance of each point from the centroid. This is an estimate of “radius”. We compute the standard deviation of this radius. If this is lower than 0.03, then we assign a perfect reward. The reward linearly decreases when the std increases from 0.03 to 0.06. Beyond that, we give zero reward. We tuned these thresholds empirically by generating and distorting circle configurations.

We follow a similar evaluation strategy for make-a-line. Here we compute the average slope and fit a line to our data. Then we measure the standard deviation of the distance of each point from the line. We found that the same thresholds we use for circles work well for lines as well.

VII-C Additional Experiments

Details on Generalization Experiments: We conduct controlled studies of our model’s generalization across three axes: a) Novel Colors b) Novel background color of the table c) Novel Objects. In each of these settings, we only change one attribute (i.e. object color, background color or object instance) while keeping everything else constant.

Novel colors: We train the models with [“blue”, “red”, “green”, “yellow”, “brown”, “gray”, “cyan”] colors and evaluate them with unseen [“orange”, “purple”, “pink”, “white”].

Novel background colors: All models are trained with black colored tables and evaluated with randomly sampled RGB color for each instruction.

Novel objects: We train the models on [“ring”, “cube”, “cylinder”, “bowl”] and evaluate them with [“triangle”, “square”, “plus”, “diamond”, “pentagon”, “rectangle”, “flower”, “star”, “circle”, “hexagon”, “heart”].

Real-World Experiments In our real-world experiments, we also use novel objects or novel object descriptions like “bananas”, “strawberry”, “small objects” which the model hasn’t seen during training in simulation.

Closed-Loop Execution: As we show in the main paper, adding feedback roughly adds 3.5% performance boost on comp-group benchmark and 1% boost on comp-one-step benchmark. Retrying cannot always recover from failure, however, it would be still important to know if the execution failed so that we can request help from a human. Towards this, we evaluate the failure detection capabilities of our feedback mechanism in Table-IX. We observe that all metrics are high for comp-group benchmark. For comp-one-step, however, the recall is very low i.e. only 19.2%. This is because in this benchmark, we only need to move one object to complete the task and hence most failures are due to wrong goal generation. Thus, the failure classifier classifies those demos as success, because indeed the model managed to achieve its predicted goal and hence results in low recall. This motivates the use of external failure classifiers in conjunction with internal goal-checking classifiers. In contrast, comp-group benchmark is harder because it requires the model to move many objects and thus results in higher chances of robot failures. These failures can be better detected and fixed by our goal-checking classifier.

Also, note that while previous works like InnerMonologue show large improvements by adding closed-loop feedback, the improvements for us are smaller. This is because, by design, our model is more likely to reach its goal - since the EBM generates a visual goal and then the low-level policy predicts a pick-and-place location within the predicted goal, it is very likely to satisfy its predicted goal. Indeed, in the comp-one-step, the model satisfies its goal in 93% cases and in 70% cases in the comp-group benchmark. In contrast, InnerMonologue does not have any built-in mechanisms for promoting goal-reaching behaviours and thus the difference in performance with additional goal-satisfying constraints is larger for that method.

There is a lot of potential in designing better goal checkers. As already discussed, we can incorporate external success classifiers like those used by prior literature in conjunction with goal-checking classifiers. Explicit goal generation allows then to check validity of goals directly even before actual robot execution (which makes it safer and less expensive). We leave this for future work. Another promising direction is to add object trackers in the feedback mechanism to keep track of the objects and detect if a failure happened. This is important, if we have multiple objects that are visually similar and hence we would need to keep track of which object corresponds to which goal and retry if it failed to reach it.

Error Analysis: We conduct a detailed error analysis of our model on comp-group benchmark, shown in Table-X. We find that robot failures, i.e. collisions or failure to pick/place objects, result in 6.0% drop in accuracy. Goal generation adds 11.7% to the failure. Visual grounding leads to 5.1% errors while we find language parsing to be nearly perfect.

VII-D Prompt used for LLMPlanner

We show a prompt used for LLMPlanner for comp-one-step in Figure 5.

VII-E EBM Energy Landscape Visualization

We visualize the energy landscape of our EBMs for different concepts in Figure-6. For binary relations (A, rel, B), we fix B’s bounding box in the scene and then move the bounding box of A all over the scene and evaluate the energy of the configuration at each position. For shapes, it is impossible to represent the landscape in 2D or 3D because we need to jointly consider all possible combinations of objects in the scene. Hence, we fix all but one object in a valid circle/line location and move a free box in all possible regions. For compositions of relations, we move only one box and score the sum of energy for all constraints. We expect the energy to be low in regions which satisfy the described concept (relation(s) or shape) and high elsewhere. We observe that the energy landscape is usually smooth with low values in valid regions and high otherwise.

VII-F Additional Related Work

Constraint-Guided Layout Optimization: Automatic optimization for object rearrangement has been studied outside the field of robotics. and use few user-annotated examples of scenes to adapt the hyperparameters of task-specific cost functions, which are then minimized using standard optimization algorithms (hill climbing and/or simulated annealing). To learn those hyperparameters from data, these approaches fit statistical models, e.g. Mixtures of Gaussian, to the given samples. further employ such optimization constraints into an interactive environment, where the user can provide an initial layout and the algorithm suggests improvements. All these approaches require expert knowledge to manually design rules and cost function, namely identifies seven and eleven expert-suggested criteria for successful rearrangement. Since they are hand-crafted, these methods do not generalize beyond the domain of furniture arrangement. In contrast, energy optimization is purely data-driven and domain-agnostic: a neural network scores layouts, assigning high energy to those that do not satisfy the (implicit) constraints and low energy to those who do, essentially modeling the underlying distribution of valid layouts.