Masked Contrastive Representation Learning for Reinforcement Learning

In pixel-based reinforcement learning (RL), the states are raw video frames, which are mapped into hidden representation before feeding to a policy network. To improve sample efficiency of state representation learning, recently, the most prominent work is based on contrastive unsupervised represent...

Full description

Saved in:
Bibliographic Details
Published in:IEEE transactions on pattern analysis and machine intelligence 2023-03, Vol.45 (3), p.3421-3433
Main Authors: Zhu, Jinhua, Xia, Yingce, Wu, Lijun, Deng, Jiajun, Zhou, Wengang, Qin, Tao, Liu, Tie-Yan, Li, Houqiang
Format: Article
Language:English
Subjects:
Algorithms
Coders
Contrastive learning
Convolutional neural networks
Correlation
Frames (data processing)
Games
Image reconstruction
Machine learning
Modules
reinforcement learning
Representation learning
Representations
Task analysis
Training
transformer
Transformers
Citations: Items that this one cites
Items that cite this one
Online Access:Get full text
Tags: Add Tag
No Tags, Be the first to tag this record!
Description
Summary:In pixel-based reinforcement learning (RL), the states are raw video frames, which are mapped into hidden representation before feeding to a policy network. To improve sample efficiency of state representation learning, recently, the most prominent work is based on contrastive unsupervised representation. Witnessing that consecutive video frames in a game are highly correlated, to further improve data efficiency, we propose a new algorithm, i.e. , masked contrastive representation learning for RL (M-CURL), which takes the correlation among consecutive inputs into consideration. In our architecture, besides a CNN encoder for hidden presentation of input state and a policy network for action selection, we introduce an auxiliary Transformer encoder module to leverage the correlations among video frames. During training, we randomly mask the features of several frames, and use the CNN encoder and Transformer to reconstruct them based on context frames. The CNN encoder and Transformer are jointly trained via contrastive learning where the reconstructed features should be similar to the ground-truth ones while dissimilar to others. During policy evaluation, the CNN encoder and the policy network are used to take actions, and the Transformer module is discarded. Our method achieves consistent improvements over CURL on 14 out of 16 environments from DMControl suite and 23 out of 26 environments from Atari 2600 Games. The code is available at https://github.com/teslacool/m-curl .
ISSN:0162-8828
1939-3539
2160-9292
DOI:10.1109/TPAMI.2022.3176413