Data Dreaming for Object Detection: Learning Object-Centric State Representations for Visual Imitation

We present a visual imitation learning method that enables robots to imitate demonstrated skills by learning a perceptual reward function based on object-centric feature representations. Our method uses the background configuration of the scene to compute object masks for the objects present. The robotic agent then trains a detector for the relevant objects in the scene via a process we call data dreaming, generating a synthetic dataset of images of various object occlusion configurations using only a small amount of background-subtracted ground truth images. We use the output of the object detector to learn an object-centric visual feature representation. We show that the resulting factorized feature representation comprised of per-object appearance features and cross-object relative locations enables efficient real world reinforcement learning that can teach a robot a policy based on a single demonstration after few minutes of training.