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: We present a unified framework for self-supervised learning of 3D representations for visual reinforecment learning. Our framework consists of two phases: a pretraining phase where a deep voxel-based 3D autoencoder is pretrained on a large object-centric dataset, and a 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, better sim-to-real transfer, and robustness compared to 2D representation learning methods. Videos are available at https://3d4rl.github.io/ .