InfoBot: Transfer and Exploration via the Information Bottleneck

TL;DR

This work tackles sparse-reward reinforcement learning by identifying decision states through a variational information bottleneck between goals and actions. It introduces InfoBot, a goal-conditioned policy that minimizes $I(A;G|S)$ to cultivate useful default habits while enabling goal-driven deviations at decision states, and uses a KL-based exploration bonus derived from the learned structure to transfer exploration strategies to new tasks. Empirical results on MiniGrid and goal-based navigation demonstrate superior direct policy transfer and robust, task-tuned exploration compared to standard baselines. The approach also draws connections to neuroscience by framing the cost of control as a mechanism that promotes efficient habit formation and selective goal-directed behavior.

Abstract

A central challenge in reinforcement learning is discovering effective policies for tasks where rewards are sparsely distributed. We postulate that in the absence of useful reward signals, an effective exploration strategy should seek out {\it decision states}. These states lie at critical junctions in the state space from where the agent can transition to new, potentially unexplored regions. We propose to learn about decision states from prior experience. By training a goal-conditioned policy with an information bottleneck, we can identify decision states by examining where the model actually leverages the goal state. We find that this simple mechanism effectively identifies decision states, even in partially observed settings. In effect, the model learns the sensory cues that correlate with potential subgoals. In new environments, this model can then identify novel subgoals for further exploration, guiding the agent through a sequence of potential decision states and through new regions of the state space.

Figures (11)

• Figure 1: Policy architecture.
• Figure 2: MultiRoomN$X$S$Y$ and FindObj$SY$ MiniGrid environments. See text for details.
• Figure 3: Policy generalization on MultiRoomN$X$S$Y$. Success is measured by the percent of time the agent can find the goal in an unseen maze. Error bars are standard deviations across runs. Baseline is a vanilla goal-conditioned A2C agent.
• Figure 4: Transferable exploration strategies on MultiRoomN$X$S$Y$. As the number of rooms increases (from left to right), a count-based exploration bonus alone cannot solve the task, whereas the proposed exploration bonus, by being tuned to task structure, enables success on these more difficult tasks.
• Figure 5: Goal based MiniPacMan navigation task: We train on a $6 \times 6$ environment, and evaluate the generalization performance in a $11 \times 11$ maze. The agent is represented by white color and has to reach the goal (light green marker).