Deep Reinforcement and InfoMax Learning
This work addresses the challenge of enhancing adaptability and performance in reinforcement learning agents, particularly for continual learning and procedurally generated environments, though it is incremental as it builds on existing methods like C51.
The paper tackled the problem of improving model-free reinforcement learning agents by making their internal representations predictive of future states beyond just rewards, and demonstrated that augmenting the C51 baseline with a temporal Deep InfoMax objective led to improved performance on a continual learning task and the Procgen environment.
We begin with the hypothesis that a model-free agent whose representations are predictive of properties of future states (beyond expected rewards) will be more capable of solving and adapting to new RL problems. To test that hypothesis, we introduce an objective based on Deep InfoMax (DIM) which trains the agent to predict the future by maximizing the mutual information between its internal representation of successive timesteps. We test our approach in several synthetic settings, where it successfully learns representations that are predictive of the future. Finally, we augment C51, a strong RL baseline, with our temporal DIM objective and demonstrate improved performance on a continual learning task and on the recently introduced Procgen environment.