Zero-Shot Generalization in Reinforcement Learning: An Information-Based Approach

Achieving zero-shot generalization in reinforcement learning is an open problem that poses challenges that differ greatly from standard supervised learning settings. A possible approach involves learning representations that are invariant to environment-specific context while retaining task-relevant information, making agents compatible with unseen environments at test time. While mutual information regularization has been applied to representation learning in RL, the critical choice of the regularization method itself has not been comprehensively investigated. This research shifts the focus from the RL algorithm alone to learning invariant representations, treating the policy optimization itself as a black box. We examine various methods both theoretically and empirically, using techniques based on information theory. We evaluate these regularization strategies by applying them in several RL frameworks, comparing training stability and generalization performance. Initial experiments in controlled environments reveal practical challenges in balancing invariance with task performance, and we discuss ongoing work to address these limitations. We demonstrate that the choice of information-theoretic objective can influence an agent's training dynamics and offer insights for designing more principled methods for zero-shot generalization.