Visual Reinforcement Learning with Self-Supervised 3D Representations

Yanjie Ze, Nicklas Hansen, Yinbo Chen, Mohit Jain, Xiaolong Wang

Keywords: Reinforcement Learning, 3D Representation Learning

Abstract: A prominent approach to visual Reinforcement Learning (RL) is to learn an internal state representation using self-supervised methods, which has the potential benefit of improved sample-efficiency and generalization through additional learning signal and inductive biases. However, while the real world is inherently 3D, prior efforts have largely been focused on leveraging 2D computer vision techniques as auxiliary self-supervision. In this work, we present a unified framework for self-supervised learning of 3D representations for motor control. Our proposed framework consists of two phases: a \textit{pretraining} phase where a deep voxel-based 3D autoencoder is pretrained on a large object-centric dataset, and a \textit{finetuning} phase where the representation is jointly finetuned together with RL on in-domain data. We empirically show that our method enjoys improved sample efficiency in simulated manipulation tasks compared to 2D representation learning methods. Additionally, our learned policies transfer zero-shot to a real robot setup with only approximate geometric correspondence, and successfully solve motor control tasks that involve grasping and lifting from \textit{a single, uncalibrated RGB camera}. Videos are available at https://yanjieze.com/3d4rl .