Source-linked AI summary
InfoBot: Transfer and Exploration via the Information Bottleneck
Anirudh Goyal, Riashat Islam, Daniel Strouse, Zafarali Ahmed, Matthew Botvinick, Hugo Larochelle, Yoshua Bengio, Sergey Levine
TL;DR
Sparse rewards make exploration and transfer difficult, so the paper learns task structure from goal-conditioned experience using an information bottleneck. The learned policy identifies decision states where goal information matters, and its KL-based signal guides exploration in new environments. The authors report better direct policy transfer and more effective exploration than task-agnostic alternatives, including an 85% success rate as environments grow larger.
Problem
Sparse-reward reinforcement learning lacks useful signals for quickly discovering effective policies, especially as partially observable environments become larger.
Method
The paper trains a goal-conditioned policy with an information bottleneck to learn default behaviors and identify decision states, then uses their KL divergence as an exploration bonus.
Results
The method achieves better direct policy transfer and more effective exploration than task-agnostic alternatives; success remains 85% as the number of rooms increases.
Takeaways & Limitations
Decision states provide learned subgoals for guiding agents through new regions of the state space and across related tasks.
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
Sparse-reward reinforcement learning remains difficult because agents lack useful signals for rapid exploration. The paper uses an information bottleneck to learn default behaviors, identify when goals require deviations, and turn those decision states into exploration guidance.
- 1 INTRODUCTION: The paper targets efficient exploration and transfer in environments where rewards are sparse and training signals are limited.It focuses on learning useful signals from multi-goal task structure rather than relying only on environmental reward.
- 1 INTRODUCTION: Goal-conditioned policies receive a sampled goal that provides information about the episode’s reward structure.The policy is represented as πθ(A | S, G), where S is the state and A is the action.
- 1 INTRODUCTION: The method minimizes I(A; G | S) while maximizing reward, encouraging behavior that performs across varied goals without overfitting to one goal.This information-bottleneck regularization is inspired by information-bottleneck methods for supervised learning.
- 1 INTRODUCTION: The resulting policy learns default behavior and deviates from it at decision states where goal-dependent actions are needed.The KL formulation interprets the regularizer as encouraging a goal-marginalized default policy.
- 1 INTRODUCTION: The exploration procedure freezes the trained policy and uses its KL divergence from the default policy as a bonus for seeking decision states.The authors report that this strategy is more effective than task-agnostic methods including count-based exploration, VIME, and curiosity.
- 1 INTRODUCTION: The contributions are information-bottleneck regularization for multi-goal RL and decision-state-based exploration for new environments.The authors report improved task generalization and more effective exploration than the listed baselines.
2 OUR APPROACH
The approach trains a goal-conditioned policy whose encoded goal information is penalized relative to a goal-marginalized default, then transfers this learned signal to exploration in new tasks. A variational approximation makes the information objective tractable, while a second policy is trained with the resulting bonus.
- 2 OUR APPROACH: The training objective maximizes environmental reward while penalizing the KL divergence between the goal-conditioned policy and a goal-marginalized default policy.The tradeoff is controlled by β > 0.
- 2 OUR APPROACH: The policy uses an encoder p_enc(Z | S, G) and decoder p_dec(A | S, Z) to convert goal information into actions.The encoder represents goal information considered useful in the current state, and the decoder produces an action distribution.
- 2 OUR APPROACH: Because I(Z; G | S) ≥ I(A; G | S), minimizing encoded goal information also minimizes the policy’s goal dependence.This connection follows from the data processing inequality.
- 2.1 TRACTABLE BOUNDS ON INFORMATION: The intractable goal-marginalized distribution is replaced with a variational approximation q(Z | S), yielding a lower bound optimized in practice.The implementation fixes q(Z | S) to a unit Gaussian, though it could also be learned.
- 2 OUR APPROACH: The policy-gradient update combines a modified return with a direct KL-divergence optimization term.The modified reward includes the KL divergence between the encoder distribution and q(Z | S).
- 2.2 POLICY AND EXPLORATION TRANSFER: After training, the learned policy identifies decision states, while a new policy is trained on test tasks using the KL divergence as an exploration reward.Algorithm 1 samples training and test tasks separately and applies the corresponding policy updates.
3 RELATED WORK
The paper relates decision states to prior information-based subgoals, bottleneck states, and exploration methods, while distinguishing its use of information bottleneck and KL regularization. It argues that decision states provide a more selective basis for subgoals and exploration.
- Information-based decision states: Prior work used action-goal information to identify decision states as subgoals, whereas this paper applies the idea to deep reinforcement learning and exploration bonuses.The paper emphasizes scalability and exploration in new domains rather than options-based subgoals alone.
- Bottleneck states: Decision states differ from bottleneck states because they incorporate reward structure and avoid treating every state in some environments as a subgoal.In the T-maze example, only the junction where the upper segments meet is a decision state, although many states qualify as bottlenecks.
- Information bottleneck: The paper uses the information-bottleneck terminology because it limits goal information in the policy, but it maximizes reward rather than predictive information.The authors therefore describe the correspondence with supervised information bottleneck methods as non-exact.
- KL regularization: Its KL divergence term differs from prior policy-distillation work because it involves goals and supports exploration and transfer to new domains.The related method uses KL regularization between different policies without a goal or exploration objective.
- Exploration methods: The approach is positioned against noise-based, curiosity-driven, and count-based exploration methods.These categories respectively perturb action selection, seek prediction-related signals, or incentivize rarely visited states.
4 EXPERIMENTAL RESULTS
Experiments show that information-bottleneck goal-conditioned policies transfer better to larger environments, while decision-state bonuses support exploration as task structure becomes more difficult.
- Experimental scope: The experiments test direct policy transfer and transferable exploration in partially observable, sparsely rewarded environments.The benchmark includes MiniGrid and goal-based MiniPacMan tasks, with training on smaller environments and evaluation on larger ones.
- Direct policy generalization: Goal-bottlenecked policies generalize much better than standard goal-conditioned A2C across larger MultiRoomNXSY and FindObjSY environments.On MultiRoomNXSY, the bottlenecked policy solves larger tasks near perfectly when trained on two-room tasks, whereas the baseline solves fewer than 50% of mazes.
- Direct policy generalization: The learned policies develop environment-specific default behaviors, including wall following in MultiRoomNXSY and doorway-to-doorway navigation in FindObjSY.These behaviors are reported as effective because they remain useful in the corresponding larger environments.
- Transferable exploration: As maze size increases, the proposed decision-state exploration maintains 85% success while curiosity falls to 54% and count-based exploration and VIME fail completely.On three-room mazes, count-based and curiosity-driven methods slightly outperform the proposed strategy, but the ordering reverses as the number of rooms grows.
- Goal-based navigation: In partially observed MiniPacMan, agents train in 6 × 6 mazes and are evaluated in 11 × 11 mazes, where the goal-bottlenecked agent solves 64% of trials.Other compared agents solve fewer than 50% of the larger mazes.
5 CONNECTIONS TO NEUROSCIENCE AND COGNITIVE SCIENCE
InfoBot is framed as a pair of automatic and controlled action pathways. Its goal-to-action route parallels cognitive-neuroscience accounts in which controlled behavior overrides habitual responses using task context.
- Automatic responses link perceptual inputs directly to actions through habitual stimulus-response associations.
- Controlled behavior overrides automatic responses to align actions with a more complete representation of task context.
- InfoBot contains pathways corresponding to automatic and controlled action selection.The paper relates these pathways to cognitive-neuroscience models of habitual and context-sensitive behavior.
- The paper connects InfoBot’s goal-to-action route with the prefrontal-cortex role proposed for task-context representation and overriding automatic responses.
6 CONCLUSION
The paper concludes that an information bottleneck trains default behaviors while preserving goal-dependent deviations when worthwhile, supporting transfer and exploration. Its formal framework uses variational KL regularization, and experiments extend across control, Atari, and maze settings.
- The method trains default behaviors and identifies when to break them through an information bottleneck between goals and policy.
- The learned decision-state exploration strategy is reported to outperform other task-agnostic exploration methods, including count-based bonuses, VIME, and curiosity.
- The framework assumes a Markov chain G → Z → A and uses the Data Processing Inequality to bound information between goals and latent variables or actions.
- The regularizer constrains goal information in the policy by penalizing divergence from a goal-marginalized or fixed prior representation.
- In continuous-control experiments, high-value states serve as sampled proxies for goals, while Atari experiments use high-value-state buffers and A2C.
B TRANSFERRED EXPLORATION STRATEGY IN ATARI
The Atari experiments test whether a decision-state encoder trained on Seaquest can transfer an exploration strategy to other games. The reported comparisons include Pong, Qbert, Freeway, and related Atari tasks against A2C baselines.
- An encoder trained on Seaquest identifies decision states and is reused to provide an exploration bonus in new Atari environments.
- The transferred strategy is evaluated on Pong and Qbert, with Figure 8 additionally comparing Pong, Qbert, and Freeway using an egocentric encoder.
- Figure 8 compares InfoBot, defined as A2C plus a KL regularizer, with a baseline A2C agent across transferred Atari games.
- The experiments average results over four random seeds in Figure 8, while the described Pong and Qbert evaluations report transfer to new games.
C ALGORITHM IMPLEMENTATION DETAILS
The implementation evaluates InfoBot across partially observable mazes, Atari, continuous control, communication, and visitation-count settings. These experiments use A2C or PPO, goal or high-value-state conditioning, and task-specific training setups.
- MINIGRID ENVIRONMENTS: Maze experiments use goal-conditioned A2C on MiniGrid tasks, representing goals by the agent’s relative distance to the absolute goal position.
- ALGORITHM IMPLEMENTATION DETAILS: Maze training uses 48 parallel workers, fully connected actor and critic networks, a one-layer encoder, and RMSProp with learning rate 0.0007.
- ALGORITHM IMPLEMENTATION DETAILS: Atari experiments use convolutional A2C policies conditioned on high-value states, while MuJoCo experiments use PPO with high-value states as goal proxies.
- ALGORITHM IMPLEMENTATION DETAILS: The grid-world hyperparameter search varies only the KL-loss weight across five candidate values for each environment.
- ALGORITHM IMPLEMENTATION DETAILS: The control hyperparameter search varies only the KL-loss weight across three candidate values and selects the setting with the most improvement over PPO.
- MULTI-AGENT COMMUNICATION: The communication experiment uses identical policies, shared rewards, identical action and observation spaces, and a cooperative multi-agent setting.
- VISITATION COUNT: The visitation-count experiment compares a 6 × 6 training maze with an 11 × 11 evaluation maze and uses blueness to quantify visited states.
F COMPARISON WITH OFF POLICY ALGORITHMS (SAC)
The proposed method is evaluated against SAC in sparse-reward MuJoCo tasks and extended to language-conditioned interactive worlds. The section also describes the task formulation, data, and comparison setup for generalization across environments and instructions.
- MuJoCo comparison: The proposed method is compared with SAC on four sparse-reward MuJoCo tasks.The tasks are HalfCheetah-v2, Walker2d-v2, Hopper-v2, and Swimmer-v2.
- MuJoCo comparison: The method improves over SAC on HalfCheetah-v2, Walker2d-v2, Hopper-v2, and Swimmer-v2.Performance is evaluated every 50K steps through 500K steps and averaged over two random seeds.
- Interactive worlds: The interactive-world task requires agents to reach language-specified goals while reasoning jointly over instructions and environment configurations.This joint reasoning is motivated by the need to generalize to unseen worlds and context-dependent language instructions.
- Task formulation: The task is modeled as an MDP with states, actions, natural-language goals, transitions, and goal-dependent rewards.The transition distribution and reward function are conditioned on the current state and goal.
- Generalization setup: The evaluation includes puddle-world environments with varied maps, objects, and linguistic inputs, using a UVFA-style baseline with MLP and LSTM goal-state embeddings.The baseline combines state and language embeddings by dot product to produce value estimates.
H MINIGRID ENVIRONMENTS FOR OPENAI GYM
The MiniGrid environments are parameterizable OpenAI Gym-compatible gridworlds used to study reinforcement-learning tasks with adjustable difficulty.
- Environment family: FindObj and MultiRoom are MiniGrid environments compatible with the OpenAI Gym framework.MiniGrid includes parameterizable environments in which task difficulty can be adjusted, such as by changing room sizes.
- Environment family: The environments support adjustable difficulty through parameters such as the sizes of rooms.This makes the package suitable for constructing tasks with different levels of difficulty.
- Environment family: The two environments used in this research are FindObj and MultiRoom.Both belong to the MiniGrid family of gridworld environments.
H.1 THE WORLD
MiniGrid represents each world as an N×N grid whose tiles contain at most one object, with objects drawn from a fixed set of types and colors.
- Grid representation: Each MiniGrid world is an N×N grid, and every tile contains zero or one object.This defines the basic spatial representation used by the environments.
- Object vocabulary: Possible object types are wall, door, key, ball, box, and goal.These object categories specify the contents that can appear on grid tiles.
- Object vocabulary: Each object has a discrete color selected from red, green, blue, purple, yellow, or grey.Walls are grey by default and goal squares are green by default.
H.2 REWARD FUNCTION
MiniGrid tasks use sparse rewards, discrete navigation actions, and partial egocentric observations encoded as compact tensors. The environments vary their visual inputs and network encoders across task types.
- Reward function: MultiRoom terminates with positive reward when the agent reaches the green goal; otherwise the episode ends with zero reward at the time limit.FindObj similarly gives positive reward for reaching the target object and zero reward when the time limit is reached.
- Action space: MiniGrid exposes seven actions, but only turning, moving forward, and environment interaction are relevant to navigation.The listed actions are turn left, turn right, move forward, pick up, drop, toggle, and done; pick up, drop, and done are irrelevant here.
- Observations: The agent receives partial, egocentric observations covering a 7×7 view oriented in its facing direction.The agent cannot see through walls or closed doors, and the observation includes the tile where it stands.
- Observations: Each observation tile is encoded by object type, color, and door-open status rather than RGB values.The compact tensor encoding is chosen for space efficiency and faster training, while RGB views are provided for human viewing.
- Task implementations: MultiRoom levels place the agent randomly in the first room and the goal randomly in the last room after generating a mult room layout.The MultiRoomNXSY variant uses a 3×3-pixel egocentric view processed by an MLP.