Abstract
In view of its power in extracting feature representation, contrastiveself-supervised learning has been successfully integrated into the practice of(deep) reinforcement learning (RL), leading to efficient policy learning invarious applications. Despite its tremendous empirical successes, theunderstanding of contrastive learning for RL remains elusive. To narrow such agap, we study how RL can be empowered by contrastive learning in a class ofMarkov decision processes (MDPs) and Markov games (MGs) with low-ranktransitions. For both models, we propose to extract the correct featurerepresentations of the low-rank model by minimizing a contrastive loss.Moreover, under the online setting, we propose novel upper confidence bound(UCB)-type algorithms that incorporate such a contrastive loss with online RLalgorithms for MDPs or MGs. We further theoretically prove that our algorithmrecovers the true representations and simultaneously achieves sample efficiencyin learning the optimal policy and Nash equilibrium in MDPs and MGs. We alsoprovide empirical studies to demonstrate the efficacy of the UCB-basedcontrastive learning method for RL. To the best of our knowledge, we providethe first provably efficient online RL algorithm that incorporates contrastivelearning for representation learning. Our codes are available athttps://github.com/Baichenjia/Contrastive-UCB.