Incorporating Attention in World Models for Improved Dynamics Modeling

Recurrent world models have recently been shown to perform well on reinforcement learning tasks, where spatio-temporal representations are learned in an unsupervised manner in simulated game environments. This work is an attempt to improve the world model by incorporating a content-based attention mechanism in the dynamics modeling module which is used to simulate a training environment. The proposed attention mechanism uses the dynamics from the recurrent neural network (RNN) with the encoded visual information to create a focused spatial representation of the input. Further, a gating network is used during training of the visual encoding module which helps in predicting descriptive factorized latent feature vectors, which tend to generate more realistic states when combined with attention in an unsupervised manner. The efficacy of the approach is demonstrated through experiments in VizDoom game environment.