Dual-head prediction and reconstruction with coarse-to-fine masks for visual reinforcement learning

In situations of limited experience and high-dimensional input data, effective representation learning plays a vital role in enabling visual reinforcement learning (RL) to excel in diverse tasks. To better leverage the agent's sampled trajectory during the training process, we introduce the DPRM approach, which involves a Dual-head Prediction and Reconstruction task with coarse-to-fine Masks in RL. The DPRM method tackles these challenges through integration of coarse-to-fine masks with a dual-head prediction-reconstruction (DHPR) architecture, complemented by a coordinate-based spatial coding strategy (CSCS). The CSCS enhances the spatial information of the observation state, facilitating the capture of motion changes between continuous context states. Furthermore, the coarse-to-fine masks gradually refine, guiding the following DHPR model to learn essential features and semantics more effectively. Built on a transformer architecture, DHPR introduces a novel triplet input token comprising two consecutive actions paired with an observation state. This design facilitates bidirectional prediction of past and future states from temporal extremities while efficiently reconstructing masked latent features throughout state sequences. Experimental results on both multiple continuous control (DeepMind Control Suite benchmarks) and discrete control (Atari) tasks demonstrate that the DPRM algorithm significantly enhances performance, leading to higher reward accumulation and faster convergence. Code is available athere.