Representation Learning in Deep RL via Discrete Information Bottleneck

TL;DR

This paper introduces RepDIB, a novel approach using variational and discrete information bottlenecks to learn structured, factorized latent representations in deep reinforcement learning, effectively filtering out irrelevant information and improving performance.

Contribution

The work presents a new architecture, RepDIB, that integrates discrete information bottlenecks with existing RL objectives to enhance latent state learning in noisy environments.

Findings

RepDIB improves performance on RL benchmarks.

Compressed representations focus on relevant state information.

Effective in both online and offline RL settings.

Abstract

Several self-supervised representation learning methods have been proposed for reinforcement learning (RL) with rich observations. For real-world applications of RL, recovering underlying latent states is crucial, particularly when sensory inputs contain irrelevant and exogenous information. In this work, we study how information bottlenecks can be used to construct latent states efficiently in the presence of task-irrelevant information. We propose architectures that utilize variational and discrete information bottlenecks, coined as RepDIB, to learn structured factorized representations. Exploiting the expressiveness bought by factorized representations, we introduce a simple, yet effective, 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 strong performance improvements, as the learned bottlenecks help predict only the relevant state while ignoring irrelevant information.