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

TL;DR

Problem: enable robots to rearrange scenes from language with strong generalization to longer, novel predicate compositions. Approach: Scene Rearrangement via Energy Minimization (SREM) uses energy-based models to represent spatial predicates, grounded by an open-vocabulary detector, with a semantic parser translating instructions into a sum of predicate energies and gradient-based goal generation, followed by vision-based manipulation. Contributions: a modular EBMs-based planning framework, a new compositional instruction benchmark, and extensive comparisons showing zero-shot compositional generalization and real-world transfer. Impact: enables robust, scalable instruction-following for complex scene rearrangements beyond training distributions.

Abstract

Language is compositional; an instruction can express multiple relation constraints to hold among objects in a scene that a robot is tasked to rearrange. Our focus in this work is an instructable scene-rearranging framework that generalizes to longer instructions and to spatial concept compositions never seen at training time. We propose to represent language-instructed spatial concepts with energy functions over relative object arrangements. A language parser maps instructions to corresponding energy functions and an open-vocabulary visual-language model grounds their arguments to relevant objects in the scene. We generate goal scene configurations by gradient descent on the sum of energy functions, one per language predicate in the instruction. Local vision-based policies then re-locate objects to the inferred goal locations. We test our model on established instruction-guided manipulation benchmarks, as well as benchmarks of compositional instructions we introduce. We show our model can execute highly compositional instructions zero-shot in simulation and in the real world. It outperforms language-to-action reactive policies and Large Language Model planners by a large margin, especially for long instructions that involve compositions of multiple spatial concepts. Simulation and real-world robot execution videos, as well as our code and datasets are publicly available on our website: https://ebmplanner.github.io.

Figures (6)

• Figure 1: Energy-based Models are Zero-Shot Planners for Compositional Scene Rearrangement. We represent language concepts with energy functions over object locations and sizes. Gradient descent on the sum of energy functions, one per predicate in the instruction, iteratively updates the object spatial coordinates and generates a goal scene configuration that satisfies the instruction, if one exists.
• Figure 2: Scene rearrangement through energy minimization. Given an image and a language instruction, a semantic parser maps the language into a set of energy functions (BinaryEBM, MultiAryEBM), one for each spatial predicate in the instruction, and calls to an open-vocabulary visual language grounder (VLMGround) to localize the object arguments of each energy function mentioned in the instruction, here "fruits" and "plate". Gradient descent on the sum of energy functions with respect to object spatial coordinates generates the goal scene configuration. Vision-based neural policies condition on the predicted pick and place visual image crops and predict accurate pick and place locations to manipulate the objects.
• Figure 3: Planning in language space with Large Language Models (LLMs). LLM Planners predict language subgoals that decompose the initial instruction to simpler-to-execute subtasks. Predicted language subgoals are fed to reactive language-to-action policies for execution. In cases where concept intersection is needed, the predicted sequential language subgoal decomposition of instructions can fail. Here, the LLM predicts the first subgoal of putting the strawberry to the right of the apple. The reactive policy can succeed if it places the strawberry anywhere within the shaded region. During execution of the next issued language subgoal of putting the strawberry in front of the bowl, the policy violates the first constraint. Placing the strawberry in the intersection of the two shaded regions may not be achieved by decomposing the two predicates sequentially, as opposed to composing them. Then the burden of handling the compositional instruction is outsourced to the language-to-action policy, which often fails to generalize. Instead, SREM directly addresses compositionality of multiple spatial language predicates.
• Figure 4: (a): Architecture of the EBM used for binary concepts such as "right of". The inputs are two boxes $O_1$ and $O_2$ and the output is the energy of their relative placement. (b): Architecture of the EBM used for multi-ary concepts such as "circle". The input is a set of $n$ entities $O_k, k=1, \dots, n$. The output is the energy of this set of entities wrt the concept. (c): Architecture of the EBM used for 3D binary concepts such as "on". Each object is now represented by a 3D bounding box. (d): Architecture of the EBM used for concepts that involve pose optimization (rotation). Each object is represented with its center and rotation wrt the global coordinate frame.
• Figure 5: Example prompt used for LLMPlanner.