Representation Learning in Deep RL via Discrete Information Bottleneck

Several representation learning methods using self-supervised objectives have been proposed for rich observation space reinforcement learning (RL). For real world applications of RL, recovering underlying latent states is of importance, especially when sensory inputs can contain irrelevant and exogenous information. A central object of this work is to understand effectiveness of compressed representations, through use of information bottlenecks, to learn robust representations in presence of task irrelevant information. We propose an algorithm that utilizes variational and discrete information bottleneck, coined as RepDIB, to learn structured factorized representations. By exploiting the expressiveness in representation space due to factorized representations, we introduce a simple, yet effective, representation bottleneck that can be integrated with any existing self supervised objective for RL. We demonstrate this across several online and offline RL benchmarks, along with a real robot arm task, where we find that compressed representations with RepDIB can lead to effective performance improvements, since the learnt bottlenecks can help with predicting only the relevant state, while ignoring irrelevant information.