[Paper Review] InfoBot: Transfer and Exploration via the Information Bottleneck
InfoBot regularizes goal-conditioned policies with an information bottleneck to discover decision states, enabling better transfer to new tasks and a transferable exploration bonus that outperforms standard exploration methods.
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.
Motivation & Objective
- Encourage learning of task structure by reducing dependency on individual goals in multi-goal RL.
- Identify decision states where goal-dependent decisions are needed.
- Leverage learned decision states to provide a transfer-friendly exploration bonus in new environments.
- Promote generalization and efficient exploration over diverse, sparsely rewarded tasks.
Proposed method
- Goal-conditioned policy with an information bottleneck regularizer: maximize expected reward while penalizing dependence on the goal via I(A;G|S).
- Policy factorization using encoder p_enc(Z|S,G) and decoder p_dec(A|S,Z) to approximate πθ(A|S,G).
- Use variational bound with q(Z|S) to avoid intractable marginalization, resulting in a tractable objective including KL[p_enc(Z|S,G)||q(Z|S)].
- On-policy policy gradient updates (REINFORCE) with modified reward r̃t = rt + β KL[p_enc(Z|st,gt)||q(Z|st)].
- After training on train tasks, freeze encoder, use KL[p_enc(Z|S,G)||q(Z|S)] as an exploration bonus when training a new policy on test tasks, modulated by a count-based term to discourage over-exploration.
- Optional: demonstrate transfer by applying the encoder from training tasks to guide exploration in unseen but related tasks.
Experimental results
Research questions
- RQ1Can a goal-conditioned policy with an information bottleneck improve policy transfer to related, unseen tasks?
- RQ2Can the decision states learned via the information bottleneck serve as an effective, task-tuned exploration bonus in new environments?
- RQ3Does regularizing goal information promote generalization in multi-goal RL under partial observability?
- RQ4How does InfoBot's transferable exploration compare to count-based, VIME, and curiosity-based exploration methods?
Key findings
- Goal-bottlenecked policies generalize better than vanilla goal-conditioned baselines on MiniGrid tasks.
- InfoBot yields higher transfer performance to larger or more complex environments (e.g., MultiRoom and FindObj tasks) than standard RL approaches.
- Using decision states as an exploration bonus maintains high success in transfer settings, outperforming count-based, VIME, and curiosity-driven methods as task complexity increases.
- In goal-based MiniPacMan, InfoBot achieves 64% success on 11x11 mazes after training on 6x6, exceeding several baselines.
- Across tasks, the InfoBot controller learns habits that generalize, while deviations (decision states) capture task-specific adaptations for exploration.
Better researchstarts right now
From reading papers to final review, dramatically reduce your research time.
No credit card · Free plan available
This review was created by AI and reviewed by human editors.