# A2C: Attention-Augmented Contrastive Learning for State Representation Extraction
## 1 Introduction  
Reinforcement learning (RL) has achieved great success in various domains [1], but it requires a large amount of interaction with the environment and is difficult to generalize to new situations. State representation used to train RL plays an important role in the sample efficiency and generalization. In general, those low-dimensional and robust representations are easier to associate with the reward signal [2], and they can help agents obtain better generalization with less training episodes. In this paper, we explore providing a more efficient state representation for RL. In vision-based RL, most existing work directly uses raw images (RGB usually) as the state representation [3; 4; 5; 6; 7]. Although RGB is efficient for accomplishing tasks, it results in two disadvantages for RL to some extent, i.e., sample inefficiency and poor generalization. RGB has a very powerful expression capability that can express the state of almost any visible scenes. However, it is this flexibility that makes its expression non-robust, because pixels in RGB are independent of each other, and it is difficult to maintain any robust features. For example, even if the camera is placed in the same place with the same perspective, the representations in the RGB space may be very different due to some changes of the environmental conditions such as light and object position. To overcome the shortcomings of RGB, researchers have sought other ways to obtain more effective state representations. Currently, there are mainly two kinds of methods to improve the state representation. (1) The first is joint optimization of state representation and policy. This method generally embeds raw data processing mechanisms in the RL model and simultaneously optimizes itwith a policy network. A suitable example can be found in [3] where a pair of shared Siamese networks was used to output a joint representation of the current/target state, which was immediately taken into the policy network as input. The parameters of the Siamese network also participate in gradient back-propagation when the policy network is updated. It is shown that this joint representation can indeed improve the learning efficiency and generalization of agents. Another example comes from [8]: instead of using RGB directly, it extracts agents' neighbour objects like obstacles as nodes and embeds them into a graph [9] that is able to make the representation more compact. In fact, we tend to believe that the effectiveness of the method of joint optimization should be attributed more to the innovation of the network structure rather than the state representation, because the parameters of the raw data processing mechanism actually participate in the policy learning directly. In other words, the state representation used by this method is still RGB. (2) The second is to decouple the representation learning from policy learning and specifically optimize it for the state representation. Recently, deep learning [10] technology has become the mainstream way to obtain effective representations for agents. There were various neural network-based representation learning methods described in [2], and they can be divided into supervised and unsupervised. It has been attempted to use the state representations extracted from the supervised vision tasks, such as segmentation, edge detection, etc., to accomplish the robot navigation task [11]. Compared with RGB, these compact features are more robust, that is less affected by changes in ambient light and the scene. Therefore, it is conducive to improving the learning efficiency and generalization of agents. However, as labels are not always available, researchers seek to distil state representations through unsupervised learning, such as the generative adversarial network (GAN) [12; 13; 14], autoencoder [15], contrastive learning [16; 17; 18], etc.