Transformer-based World Models Are Happy With 100k Interactions
This work addresses the data-hungry nature of deep neural networks in reinforcement learning, offering a more sample-efficient approach for training agents in environments like Atari games.
The paper tackles the problem of sample inefficiency in reinforcement learning by introducing a transformer-based world model (TWM) that learns from real-world episodes in an autoregressive manner, achieving state-of-the-art performance on the Atari 100k benchmark.
Deep neural networks have been successful in many reinforcement learning settings. However, compared to human learners they are overly data hungry. To build a sample-efficient world model, we apply a transformer to real-world episodes in an autoregressive manner: not only the compact latent states and the taken actions but also the experienced or predicted rewards are fed into the transformer, so that it can attend flexibly to all three modalities at different time steps. The transformer allows our world model to access previous states directly, instead of viewing them through a compressed recurrent state. By utilizing the Transformer-XL architecture, it is able to learn long-term dependencies while staying computationally efficient. Our transformer-based world model (TWM) generates meaningful, new experience, which is used to train a policy that outperforms previous model-free and model-based reinforcement learning algorithms on the Atari 100k benchmark.