Masked Contrastive Representation Learning for Reinforcement Learning.

Area of Science:

• Artificial Intelligence
• Machine Learning
• Computer Vision

Background:

• Pixel-based reinforcement learning (RL) relies on mapping raw video frames to hidden states for policy networks.
• Current state representation learning often uses contrastive unsupervised learning to improve sample efficiency.
• Consecutive video frames in RL environments exhibit high correlation, offering potential for enhanced data efficiency.

Purpose of the Study:

• To introduce a novel algorithm, Masked Contrastive Representation Learning for RL (M-CURL), to improve data efficiency in RL.
• To leverage the inherent correlations among consecutive video frames for better state representation learning.
• To enhance the performance of reinforcement learning agents through improved sample efficiency.

Main Methods:

• M-CURL employs a CNN encoder for state representation and a policy network for action selection.
• An auxiliary Transformer encoder module is introduced to capture correlations among video frames.
• During training, features of randomly masked frames are reconstructed using context frames via joint CNN and Transformer training with contrastive learning.

Main Results:

• M-CURL demonstrated consistent improvements over the standard CURL algorithm.
• The proposed method achieved superior performance on 14 out of 16 DMControl suite environments.
• Significant gains were observed across 23 out of 26 Atari 2600 Games.

Conclusions:

• M-CURL effectively utilizes frame correlations to enhance state representation learning in RL.
• The algorithm offers a significant improvement in sample efficiency and overall performance.
• The Transformer module can be discarded during policy evaluation, simplifying the final agent architecture.