Source-linked AI summary
Variational Information Maximisation for Intrinsically Motivated Reinforcement Learning
Shakir Mohamed, Danilo Jimenez Rezende
TL;DR
The paper addresses the lack of scalable mutual-information optimisation, whose standard Blahut-Arimoto computation is exponential and unsuitable for high-dimensional learning. It combines variational inference with deep learning to optimise empowerment, matching exact-computation behaviours while extending reasoning to high-dimensional pixel inputs. The resulting system provides scalable mutual-information maximisation without requiring an environment generative model.
Problem
Mutual-information optimisation is generally intractable, while Blahut-Arimoto enumerates states with exponential complexity and existing intrinsic-reward approaches do not scale to high-dimensional problems.
Method
The paper uses a variational lower bound with stochastic optimisation and neural state representations, including convolutional networks for visual inputs, to maximise mutual information and empowerment.
Results
Empowerment-based behaviours from variational information maximisation match those obtained by exact mutual-information computation, including when learning from pixel information in high-dimensional problems.
Takeaways & Limitations
Variational information maximisation expands empowerment-based reasoning to high-dimensional settings where exact computation is not possible.
Takeaways & Limitations
The choice of behaviour policy remains an important line of ongoing research because using the source distribution as the behaviour policy has large variance.
Abstract
from arXiv · showhide
The mutual information is a core statistical quantity that has applications in all areas of machine learning, whether this is in training of density models over multiple data modalities, in maximising the efficiency of noisy transmission channels, or when learning behaviour policies for exploration by artificial agents. Most learning algorithms that involve optimisation of the mutual information rely on the Blahut-Arimoto algorithm --- an enumerative algorithm with exponential complexity that is not suitable for modern machine learning applications. This paper provides a new approach for scalable optimisation of the mutual information by merging techniques from variational inference and deep learning. We develop our approach by focusing on the problem of intrinsically-motivated learning, where the mutual information forms the definition of a well-known internal drive known as empowerment. Using a variational lower bound on the mutual information, combined with convolutional networks for handling visual input streams, we develop a stochastic optimisation algorithm that allows for scalable information maximisation and empowerment-based reasoning directly from pixels to actions.
1 Introduction
Mutual information is central to many machine-learning applications but is generally intractable to compute with existing methods. The paper develops variational information maximisation for scalable empowerment-based learning, including from visual inputs.
- Mutual information underlies applications including rate distortion, information bottlenecks, population coding, curiosity-driven exploration, model selection, and intrinsically motivated reinforcement learning.
- Blahut-Arimoto estimates mutual information by enumeration, giving exponential complexity unsuitable for modern machine-learning applications.
- Intrinsic motivation lets agents reason about the value of information in action-observation sequences without relying solely on external rewards.
- The paper introduces stochastic variational information maximisation for scalable mutual-information and channel-capacity estimation in discrete and continuous settings.
- Combining variational optimisation with deep learning supports empowerment-based reinforcement learning directly from pixel information.
2 Intrinsically-motivated Reinforcement Learning
Intrinsically motivated reinforcement learning addresses reward scarcity by providing internal rewards based on information or prediction. The paper focuses on scalable information-theoretic rewards for vision-based agents.
- Intrinsic motivation addresses environments where external rewards are absent or sparsely distributed.
- The perception-action framework separates internal and external environment components while an agent models incoming data and interacts online.
- Vision-based self-motivated systems use a convolutional neural network for state representation and a critic that supplies intrinsic rewards.
- Existing intrinsic reward measures include Bayesian surprise, prediction errors, salient-event prediction, and information-theoretic quantities.
- Available information-theoretic approaches require transition knowledge or summation over state configurations, making them intractable for complex environments and large images.
3 Mutual Information and Empowerment
Mutual information measures dependence between random variables, while empowerment applies it to the information action sequences contain about future states. Maximising empowerment assigns values to states and supports exploration of reachable state space.
- Mutual information is a general measure of dependence between random variables based on their joint and marginal distributions.
- In reinforcement learning, the variables can be the current state and a future action-conditioned state or observation.
- Empowerment maximises mutual information conditioned on a starting state between a K-action sequence and the final state reached.
- The action-sequence distribution defines the K-step source policy, while transition and marginal distributions specify the resulting action-state channel.
- Empowerment is channel capacity and assigns each state a task-independent value representing information in action sequences about future states.
- Optimising the source distribution yields an exploration policy that supports uniform exploration of states reachable at horizon K, though it is used for internal planning.
4 Scalable Information Maximisation
The paper makes mutual-information maximisation tractable by replacing direct computation with a variational lower bound optimised through stochastic deep-learning procedures. The resulting method avoids explicit state enumeration, supports visual inputs, and provides scalable empowerment estimation.
- Motivation: The mutual information is difficult to compute because it requires expectations over unknown transition dynamics and lacks an efficient scalable algorithm.Existing Blahut–Arimoto computation enumerates states and is limited to small-scale and discrete problems.
- Variational bound: A variational lower bound makes mutual-information estimation possible from Monte Carlo samples collected through direct environment interaction.The transition probability appears linearly in the expectation and need not be evaluated explicitly.
- Optimisation: The bound is optimised by alternating updates of a decoder q and an action-sequence source ω, with an entropy constraint preventing divergent solutions.The decoder update is supervised maximum likelihood, while the source distribution is approximated to avoid direct normalisation and sampling difficulties.
- Optimisation: The decoder learns action sequences from start and termination states using an autoregressive distribution implemented with nonlinear neural networks.The action distributions are categorical, and stochastic updates optimise the variational parameters.
- Deep-learning implementation: Raw pixel observations are transformed into state representations with convolutional networks before subsequent mutual-information computations.The full procedure uses recent experience mini-batches and stochastic gradient ascent.
- Complexity: The approach has at most quadratic complexity in hidden-layer size and is linear in other listed quantities, contrasting with state-enumeration costs that grow exponentially.For visual inputs, representation complexity is O(D^2P + LP^2), while autoregressive distributions have complexity O(H^2KN).
5 Results
The variational approach matches exact empowerment behaviour in static environments and supports empowerment-based planning in dynamic and continuous settings, including visual-input scenarios.
- 5.1 Effectiveness of the MI Bound: The variational empowerment landscape matches the exact mutual-information solution and therefore selects the same agent behaviour in static environments.The comparison uses environments with at most 400 discrete states, exact Blahut–Arimoto computation, and 20 × 20 pixel inputs.
- 5.1 Effectiveness of the MI Bound: The method closely approximates true empowerment in the two-rooms environment, with correlation coeff = 1.00 and R2=0.90.Maximum empowerment occurs centrally in large and cross-shaped rooms and near both doors in the two-rooms environment.
- 5.2 Dynamic Environments: In dynamic room environments, the approach reproduces empowered behaviour while reasoning directly from pixel information.Dynamic settings include movable objects and changing environmental features that make enumerative state-based methods quickly infeasible.
- 5.2 Dynamic Environments: Empowerment planning selects actions that preserve future options, such as building a wall to stem lava in a dynamic maze.At each time step, the agent evaluates expected empowerment after each available action and takes the action with the highest value.
- 5.2 Dynamic Environments: After collecting a key, the two-rooms scenario shifts maximum empowerment from the key region toward the door, while brick-building successfully protects the agent from lava.The same empowerment planning approach supports both behaviours.
- 5.3 Predator-Prey Scenario: In a continuous 3D predator-prey simulation, maximum empowerment lies away from the predator, and the agent learns to escape.The agent and predator are represented as spheres; actions are 2D force vectors and the state includes position, velocity, and angular momentum.
6 Conclusion
The paper develops a scalable mutual-information estimator by combining variational inference with deep learning, focusing on empowerment for intrinsically motivated learning. The resulting system supports high-dimensional, visual, continuous, and interactive settings with complexity more favourable than Blahut–Arimoto.
- 6 Conclusion: The paper develops scalable mutual-information estimation using a variational lower bound together with deep-learning techniques.The approach is developed around empowerment, an intrinsic reward measure for intrinsically motivated reinforcement learning.
- 6 Conclusion: The resulting algorithm expands empowerment to high-dimensional problems with complexity extremely favourable compared with Blahut–Arimoto.The system does not require a generative model and learns through interactions with the environment.
- 6 Conclusion: The system supports learning directly from visual information and in continuous state-action spaces.The conclusion also identifies applications for scalable mutual-information optimisation beyond intrinsic motivation.
A Empowerment as Path-counting
In deterministic discrete environments, empowerment can be understood through action-path counts: its optimal source distribution makes reachable terminal states uniformly likely. A tree-structured environment illustrates this path-counting interpretation.
- Empowerment as Path-counting: For deterministic transitions, each action sequence maps an initial state to a resulting state through the transition function T(s, a).The transition probability is represented as a delta distribution concentrated at T(s, a).
- Empowerment as Path-counting: The second Blahut-Arimoto iteration assigns action paths weights inversely proportional to the number of alternative paths reaching the same state.This path-count quantity is related to the number of distinct action paths connecting s to a terminal state s′.
- Empowerment as Path-counting: The quantity n(s) counts the different states reachable from state s at horizon K, and the paper demonstrates this reasoning with a tree-structured environment.Figure 8 depicts reachable states as nodes connected by up or down transitions.
- Empowerment as Path-counting: Unlike uniform action selection, the optimal distribution ω(∞)(a|s) visits terminal states uniformly.The resulting marginal distribution over terminal states motivates interpreting the distribution as an efficient exploration policy.
- Empowerment as Path-counting: The optimized distribution ω(∞)(a|s) allows an agent to explore all reachable states uniformly.This interpretation is presented as an intuitive property of empowerment.
B Model-based Empowerment
The main-text empowerment approach is model-free and does not use a transition-dynamics model. Model-based variants are discussed as desirable for transfer and task-independent reasoning but were less efficient and excluded from the main text.
- Model-based Empowerment: The main-text approach is model-free because it does not use a model of the environment’s transition dynamics.The paper contrasts this with model-based empowerment approaches described separately.
- Model-based Empowerment: Model-based methods could support reasoning about task-independent aspects of the world, transfer learning across domains, and potentially faster learning.These are stated motivations for considering a model-based approach.
- Model-based Empowerment: The developed model-based empowerment methods were less efficient, so they were not included in the main text.The paper says the reasons for their lower efficiency are described below.
B.1 Importance Sampling Estimator
The importance-sampling estimator approximates Blahut-Arimoto updates using sampled action sequences and future states from a transition model. It adapts normalized importance weights, but its distortion computation scales quadratically with the number of samples and may require many samples.
- Importance Sampling Estimator: The estimator assumes an available environment model p(s′|a_i, s) and approximates empowerment with a generic importance-sampling estimator.The model’s acquisition method is not specified at this stage.
- Importance Sampling Estimator: It samples S action sequences from ω_t(a|s) with positive normalized weights α_t,i and keeps the samples fixed during optimization.Only the importance weights are adapted to maximize mutual information.
- Importance Sampling Estimator: For each action-sequence sample, it generates J future-state samples from the transition model p(s′|a_i, s).These samples support the estimator’s approximation of the required mutual-information quantities.
- Importance Sampling Estimator: Rather than directly approximating all mutual-information quantities, the method approximates the Blahut-Arimoto iteration through a distortion D_t,i.The resulting normalized weights are chosen so the next source distribution approximates the Blahut-Arimoto update.
- Importance Sampling Estimator: The update has algorithmic complexity O(S), while computing the distortion scales as O(JS^2).Although applicable to continuous domains, the method typically requires many samples for accurate empowerment estimation.
B.2 Efficient Optimisation
For smooth transition models and policies, reparameterization expresses both as differentiable functions of Gaussian noise. The resulting variational objective can be optimized with stochastic gradient ascent, while the Blahut-Arimoto algorithm emerges as a special derivation of the variational bound.
- Efficient Optimisation: Smooth transition models and policies can be reparameterized as differentiable functions of action, state, and Gaussian noise.The model uses s′ = f(a_i, s, ξ_m), while the policy uses a = h_θ(s, ξ_p).
- Efficient Optimisation: The reparameterized variational bound is written as an expectation over the model and policy noise variables.The objective compares the variational distribution’s log probability with the policy’s log probability.
- Efficient Optimisation: The policy parameters θ are optimized by stochastic gradient ascent on the reparameterized bound.This provides a more efficient algorithm for the smooth-model setting.
- Efficient Optimisation: The Blahut-Arimoto algorithm can be derived from the variational bound by choosing the variational distribution as the posterior over actions given present and future states.Substitution into the variational objective and rearrangement yields the Blahut-Arimoto form.
D Neural Network Description
The implementation uses shared neural-network components to represent image observations and parameterize action distributions and directed models. Gaussian action distributions use two-layer networks with rectified nonlinearities.
- Image state representation: Image observations are converted into state representations using the same convolutional network across all experiments.Each convolution applies a rectified non-linearity and uses 10 filters per layer.
- Image state representation: The convolutional encoder uses 4 × 4 kernels with stride 1 followed by 3 × 3 kernels with stride 2.
- Action-model parameterization: Neural networks parameterize both the decoder distribution qξ(a|s, s′) and the directed model hθ(a|s).
- Action-model parameterization: Per-action distributions q(ak) are Gaussian, with their means and variances parameterized by two-layer neural networks.
- Neural-network nonlinearities: The shared distribution component uses a two-layer network with a rectified non-linearity, while ψθ(s) is also specified by a two-layer network.The rectified non-linearity is Rect(x) = max(0, x).