Keypoints into the Future: Self-Supervised Correspondence in Model-Based Reinforcement Learning

TL;DR

This paper introduces a self-supervised visual correspondence learning approach for model-based reinforcement learning, significantly improving generalization and transferability in vision-based robotic manipulation tasks.

Contribution

It presents a novel method combining self-supervised correspondence learning with predictive models, outperforming autoencoder-based methods in vision-based RL.

Findings

Better generalization in 3D scenes and occlusions

Improved transfer from simulation to real-world robots

Enhanced category-generalization capabilities

Abstract

Predictive models have been at the core of many robotic systems, from quadrotors to walking robots. However, it has been challenging to develop and apply such models to practical robotic manipulation due to high-dimensional sensory observations such as images. Previous approaches to learning models in the context of robotic manipulation have either learned whole image dynamics or used autoencoders to learn dynamics in a low-dimensional latent state. In this work, we introduce model-based prediction with self-supervised visual correspondence learning, and show that not only is this indeed possible, but demonstrate that these types of predictive models show compelling performance improvements over alternative methods for vision-based RL with autoencoder-type vision training. Through simulation experiments, we demonstrate that our models provide better generalization precision, particularly in 3D scenes, scenes involving occlusion, and in category-generalization. Additionally, we validate that our method effectively transfers to the real world through hardware experiments. Videos and supplementary materials available at https://sites.google.com/view/keypointsintothefuture