A Simple Framework for Self-Supervised Learning of Sample-Efficient World Models

Deep reinforcement learning algorithms suffer from low sample efficiency, which is addressed in recent approaches by building a world model and learning behaviors in imagination. We present a simple framework for self-supervised learning of world models inspired by VICReg, requiring neither image reconstructions nor specific neural network architectures. The learned representations are temporally consistent, which facilitates next state prediction and leads to good generalization properties for the policy and the value function. We build a world model for Atari consisting only of feedforward layers that is easy to implement and allows fast training and inference. By learning behaviors in imagination, we evaluate our method on the Atari 100k benchmark.