Efficient and Interpretable Robot Manipulation with Graph Neural Networks
Manipulation tasks, like loading a dishwasher, can be seen as a sequence of spatial constraints and relationships between different objects. We aim to discover these rules from demonstrations by posing manipulation as a classification problem over a graph, whose nodes represent task-relevant entities like objects and goals, and present a graph neural network (GNN) policy architecture for solving this problem from demonstrations. In our experiments, a single GNN policy trained using imitation learning (IL) on 20 expert demos can solve blockstacking, rearrangement, and dishwasher loading tasks; once the policy has learned the spatial structure, it can generalize to a larger number of objects, goal configurations, and from simulation to the real world. These experiments show that graphical IL can solve complex long-horizon manipulation problems without requiring detailed task descriptions. Videos can be found at: https://youtu.be/POxaTDAj7aY.