Deep Reinforcement Learning with Swin Transformers

Transformers are neural network models that utilize multiple layers of self-attention heads and have exhibited enormous potential in natural language processing tasks. Meanwhile, there have been efforts to adapt transformers to visual tasks of machine learning, including Vision Transformers and Swin Transformers. Although some researchers use Vision Transformers for reinforcement learning tasks, their experiments remain at a small scale due to the high computational cost. Experiments conducted at a large scale, on the other hand, have to rely on techniques to cut the costs of Vision Transformers, which also yield inferior results. To address this challenge, this article presents the first online reinforcement learning scheme that is based on Swin Transformers: Swin DQN. Swin Transformers are promising as a backbone in neural networks by splitting groups of image pixels into small patches and applying local self-attention operations inside the (shifted) windows of fixed sizes. They have demonstrated state-of-the-art performances in benchmarks. In contrast to existing research, our novel approach is reducing the computational costs, as well as significantly improving the performance. We demonstrate the superior performance with experiments on 49 games in the Arcade Learning Environment. The results show that our approach, using Swin Transformers with Double DQN, achieves significantly higher maximal evaluation scores than the baseline method in 45 of all the 49 games ~92%, and higher mean evaluation scores than the baseline method in 40 of all the 49 games ~82%.