World Models with an Entity-Based Representation

Abstract

Reinforcement learning (RL) is a powerful way to solve sequential decision-making tasks. However training an RL agent in a complex environment requires a large amount of interactions, which is non-ideal when acting in an environment is costly or dangerous. One alternative is to learn an approximation of the real environment, referred to as a world model. This simulator can be used to train an agent and transfer the learned policy to the real environment. Unfortunately, training world models have traditionally required a significant number of interactions in the real environment. This brings us to the same problem when it is costly or dangerous to act in the real environment. To address this problem, we present an entity-based representation and corresponding architecture, which allows for greater data efficiency in world model training. Our approach outperforms other world model baselines in an initial application to the game Pong.