Efficient and Interpretable Robot Manipulation with Graph Neural Networks

Many manipulation tasks can be naturally cast as a sequence of spatial relationships and constraints between objects. We aim to discover and scale these task-specific spatial relationships by representing manipulation tasks as operations over graphs. To do this, we pose manipulating a large, variable number of objects as a probabilistic classification problem over actions, objects and goals, learned using graph neural networks (GNNs). Our formulation first transforms the environment into a graph representation, then applies a trained GNN policy to predict which object to manipulate towards which goal state. Our GNN policies are trained using very few expert demonstrations on simple tasks, and exhibit generalization over number and configurations of objects in the environment and even to new, more complex tasks, while providing interpretable explanations for their decision-making. We present experiments which show that a single learned GNN policy can solve a variety of long-horizon blockstacking and rearrangement tasks.