Mask-based Latent Reconstruction for Reinforcement Learning

• DOI: 10.48550/arxiv.2201.12096
• Publication Date: 2022-01-01
• Authors: Changyuan Yu; Zhizheng Zhang; Cuiling Lan; Zhibo Chen; Yan Lü
• Keywords: Code (set theory); Representation; Feature Learning; Sample (material)

For deep reinforcement learning (RL) from pixels, effective state representations is crucial for achieving high performance. However, in practice, limited experience and high-dimensional inputs prevent representation learning. To address this, motivated by the success of mask-based modeling other research fields, we introduce reconstruction to promote RL. Specifically, propose a simple yet self-supervised method, Mask-based Latent Reconstruction (MLR), predict complete latent space observations with spatially temporally masked pixels. MLR enables better use context information when make them more informative, which facilitates training RL agents. Extensive experiments show that our significantly improves sample efficiency outperforms state-of-the-art sample-efficient methods on multiple continuous discrete control benchmarks. Our code available at https://github.com/microsoft/Mask-based-Latent-Reconstruction.