Revisiting Self-Attention in Deep Reinforcement Learning: A Dimensional Perspective

Improving the sample efficiency of deep reinforcement learning (DRL) agents has been an ongoing challenge in research and real-world applications. Self-attention, a mechanism originally popularized in natural language processing, has shown great potential to enhance sample efficiency when integrated with DRL algorithms. However, effective integration techniques remain underexplored. In this paper, we investigate the dimensional bias of self-attention in visual-based DRL settings. Specifically, we design various self-attention modules by altering the dimensions where the scaled dot-product attention is applied. Then, we integrate the proposed self-attention modules with the proximal policy optimization algorithm and evaluate the performances of the DRL agents in the Arcade Learning Environment. Our study reveals that each self-attention module possesses a unique bias and generates attention patterns oriented along specific axes. Similarly, each game environment exhibits a distinctive bias regarding its main dynamic axis. It is shown that games are more likely to be won by the agent whose attention axis aligns with the main dynamic axis. We demonstrate that this “alignment of biases” enables agents to learn better state representations by capturing important dynamics from the environment and serves as a practical guide in designing sample-efficient agents for specific tasks. The codes and supplementary are available at https://github.com/JianxiaoSun/self-attention-rl.