Abstract:In view of its power in extracting feature representation, contrastive self-supervised learning has been successfully integrated into the practice of (deep) reinforcement learning (RL), leading to efficient policy learning in various applications. Despite its tremendous empirical successes, the understanding of contrastive learning for RL remains elusive. To narrow such a gap, we study how RL can be empowered by contrastive learning in a class of Markov decision processes (MDPs) and Markov games (MGs) with low-rank transitions. For both models, we propose to extract the correct feature representations 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 RL algorithms for MDPs or MGs. We further theoretically prove that our algorithm recovers the true representations and simultaneously achieves sample efficiency in learning the optimal policy and Nash equilibrium in MDPs and MGs. We also provide empirical studies to demonstrate the efficacy of the UCB-based contrastive learning method for RL. To the best of our knowledge, we provide the first provably efficient online RL algorithm that incorporates contrastive learning for representation learning. Our codes are available at <a class="link-external link-https" href="https://github.com/Baichenjia/Contrastive-UCB" rel="external noopener nofollow">this https URL</a>.
What problem does this paper attempt to address?
The problem that this paper attempts to solve is how to improve the sample efficiency of representation learning through contrastive self - supervised learning in online reinforcement learning (RL), and further enhance the effects of policy learning and Nash equilibrium learning. Specifically, the paper focuses on how to use contrastive loss to extract correct state - action representations in Markov decision processes (MDPs) under low - rank transition models and zero - sum Markov games (MGs), and how to combine the upper confidence bound (UCB) algorithm for effective exploration.
### Main Contributions of the Paper
1. **Effectiveness of Representation Learning**:
- The paper proves that contrastive self - supervised learning can recover the true transition dynamics in low - rank MDPs, thus revealing the theoretical advantages of integrating representation learning into RL.
2. **Efficient Exploration Strategy**:
- The first efficient exploration strategy combining contrastive self - supervised learning is proposed. The proposed UCB - based method can be easily adapted to existing RL representation learning methods, and the experimental results show that its performance is better than previous methods.
3. **Extension to Multi - Agent Environments**:
- The algorithm and theory are extended to zero - sum MGs. An upper and lower confidence bounds (ULCB) - based algorithm is proposed, and it is proved that the sample complexity of this method when reaching an ε - approximate Nash equilibrium is \(\tilde{O}(1/\varepsilon^2)\).
### Key Technical Details
- **Contrastive Loss**:
- A contrastive loss function \(L_h(\psi, \phi; D_k^h)\) is defined for learning low - rank representations \(\tilde{\phi}_k^h\) and \(\tilde{\psi}_k^h\). The specific form is:
\[
L_h(\psi, \phi; D_k^h) := E_{D_k^h} \left[ y \log \left(1 + \frac{1}{\psi(s')^\top \phi(z)} \right) + (1 - y) \log \left(1 + \psi(s')^\top \phi(z) \right) \right]
\]
- where \(z=(s, a)\) or \(z=(s, a, b)\), \(y\) is the label, and \(D_k^h\) is the dataset for contrastive learning.
- **UCB and ULCB**:
- For single - agent MDPs, a UCB - based algorithm is proposed, and the UCB reward term is constructed to guide exploration.
- For multi - agent MGs, a ULCB - based algorithm is proposed, and the upper and lower confidence bounds value functions are constructed to guide exploration.
- **Sample Complexity**:
- It is proved that in single - agent MDPs, after \(K \geq \tilde{O}(1/\varepsilon^2)\) rounds, the generated policy \(\hat{\pi}\) is an ε - approximate optimal policy.
- It is proved that in multi - agent MGs, after \(K \geq \tilde{O}(1/\varepsilon^2)\) rounds, the generated policy pair \((\pi_k, \nu_k)\) is an ε - approximate Nash equilibrium.
### Related Work
- **Low - Rank Transition Models**:
- Compared with linear MDPs, low - rank MDPs assume that the transition dynamics is the inner product of two unknown feature vectors, which makes representation learning more challenging.
- **Application of Contrastive Learning in RL**:
- Previous work has mainly focused on empirical aspects.