MCLF: A Max-Margin Contrastive Learning Framework for Reinforcement Learning
摘要
Contrastive self-supervised learning gives new avenues for improving sample efficiency in Reinforcement Learning (RL), especially for high-dimensional observations by extracting latent features from raw inputs. However, previous studies typically base RL on standard contrastive learning for self-supervised learning, ignoring that contrastive learning implicitly assumes that positives and negatives belong to different classes in downstream tasks. This requirement is difficult to enforce in the context of RL, and ignoring this assumption may negatively affect policy learning due to the omission of beneficial discriminative cues. In this paper, we approach this problem systematically by developing a Max-Margin Contrastive Learning Framework (MCLF), which introduces a novel contrastive learning objective and improves RL efficiency in a self-supervised manner. Our method is inspired by the principle of Support Vector Machine (SVM), which learns powerful deep representations by maximizing the decision margin between positive and negative state samples. Furthermore, to ensure the efficiency of SVM optimization in the RL setting, we propose a sample splitting method to accelerate computations. We apply MCLF to several base RL algorithms and evaluate the proposed method on DeepMind Control Suite and Atari Games. In low sample regimes, MCLF outperforms state-of-the-art baselines in terms of asymptotic performance and sample efficiency.