🤖 machine learning

Distributional Active Inference

This paper presents a formal abstraction that integrates active inference into the distributional reinforcement learning framework, enabling sample-efficient optimal control for robotic systems without requiring explicit transition dynamics modeling.

Original authors: Abdullah Akgül, Gulcin Baykal, Manuel Haußmann, Mustafa Mert Çelikok, Melih Kandemir

Published 2026-08-06

📖 3 min read☕ Coffee break read

CC BY 4.0

Original authors: Abdullah Akgül, Gulcin Baykal, Manuel Haußmann, Mustafa Mert Çelikok, Melih Kandemir

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot to navigate a chaotic, unpredictable world. To do this well, the robot needs two superpowers: it must be able to organize a flood of sensory data (like what it sees and hears) into something useful, and it must be able to plan far into the future to make smart moves. For decades, the field of Reinforcement Learning (RL) has been the go-to method for teaching machines how to plan. Think of RL as a student who learns by trial and error, trying to maximize a score by guessing what actions lead to the best rewards. However, this student is often inefficient; it needs to try things millions of times to learn, and it struggles when the world is too complex to simulate perfectly in its head.

Enter Active Inference, a theory inspired by how biological brains work. Instead of just guessing, the brain is seen as a prediction machine that constantly tries to minimize the surprise of what it experiences. It plans by imagining different futures and picking the one that feels most "certain" or least surprising. While this sounds powerful, traditional Active Inference has been stuck in a corner: it usually requires the robot to build a perfect, detailed map of how the world works (a "world model") before it can plan. Building these maps is computationally expensive and often impossible for complex tasks. The big question scientists have been asking is: Can we get the smart planning benefits of Active Inference without the heavy burden of building a perfect map of the world?

This paper introduces a new method called Distributional Active Inference (DAIF) that answers "yes." The researchers propose a clever shortcut. Instead of trying to predict exactly what will happen next (like a specific temperature or position), DAIF focuses on predicting the entire range of possible futures and their probabilities all at once. They call this "distributional" learning. Imagine instead of guessing if it will rain tomorrow, you learn the full weather forecast: a 20% chance of drizzle, a 50% chance of sun, and a 30% chance of a storm. By learning this full picture of possibilities, the robot can plan effectively without needing to know the exact laws of physics governing the rain.

The paper suggests that by combining the "full picture" approach of distributional learning with the "surprise-minimizing" logic of Active Inference, robots can learn much faster and more efficiently. The authors tested this idea in computer simulations. In simple grid-world games, they found that their new method, DAIF, could solve problems that other methods failed to solve, especially when the problems got harder and required looking further ahead. In more complex, real-world-like simulations involving soft robots and video-based control, DAIF consistently outperformed the current best methods, learning to control these systems with fewer attempts. The results suggest that this approach allows agents to handle complex environments with limited computing power, much like a human brain does, by organizing information efficiently rather than brute-forcing a perfect simulation of reality.

Technical Summary: Distributional Active Inference

Problem Statement

Optimal control in complex robotic environments faces two intertwined challenges: the efficient organization of sensory state information and far-sighted action planning. While Reinforcement Learning (RL) addresses planning, it often suffers from sample inefficiency because it typically focuses only on the latter. Active Inference (AIF), a process theory explaining biological intelligence, addresses both by minimizing Expected Free Energy (EFE) through an action-perception cycle. However, existing AIF applications in artificial intelligence are largely restricted to model-based approaches that require learning explicit transition dynamics, which is computationally expensive and often infeasible for high-fidelity forward simulation.

The core problem this paper addresses is how to integrate the performance advantages of Active Inference into a framework that does not require learning an explicit transition model, thereby making AIF accessible to model-free and distributional settings.

Methodology

The authors propose a three-step theoretical and algorithmic framework to bridge Active Inference and Distributional RL:

1. Rigorous Formulation of Active Inference

The authors re-derive the AIF objective from first principles of Bayesian and causal inference. By applying do-calculus to the standard variational inference setup, they demonstrate that the standard AIF objective admits a simpler equivalent form.

They show that intervening on the observable state $X$ with a desired distribution $P_R(X)$ (acting as the reward) disconnects the observable from the latent variables $(Y, S)$ under the product rule.
This leads to a simplified ELBO (Evidence Lower Bound) where the objective consists of a reconstruction term, a policy entropy term, and a term maximizing the log-probability of the desired state distribution under the predicted trajectory.
Crucially, this formulation reveals that AIF can be viewed as a form of predictive coding where the agent optimizes a world model to fit desired outcomes, without necessarily requiring a complex inference of posterior latents for every step.

2. Push-Forward Reinforcement Learning

To integrate AIF without explicit transition modeling, the authors introduce Push-Forward RL, a theoretical framework that generalizes Distributional RL.

Trajectory Measures: They define the return distribution not just as an expectation, but as a "push-forward" of the trajectory measure induced by a policy through a return functional.
State Abstraction: They formalize the relationship between state abstractions (latent spaces) and distributional RL. By defining an encoder $S$ and a decoder $P_D$ , they show that performing distributional RL on a latent space is equivalent to applying a composite kernel operator.
Contraction Modulus: A key theoretical result (Theorem 3.5) establishes that the contraction modulus of the distributional Bellman operator in a latent space is scaled by the product of the Lipschitz constants of the encoder ( $L_E$ ) and decoder ( $L_D$ ). If the latent space provides efficient compression (small $L_E$ ) and captures reward-relevant structure (small $L_D$ ), the convergence of the Bellman backups is accelerated.
Algorithm Template: They propose Push-Forward Policy Iteration (PPI), which minimizes the Bellman residual on the space of encoding measures rather than raw states, effectively unifying model-based and model-free views.

3. Distributional Active Inference (DAIF)

The authors instantiate the above theory into a practical algorithm called Distributional Active Inference (DAIF).

Mechanism: DAIF learns a state-action amortized parametric distribution (an encoder) that maps observations to a latent space. It then performs temporal-difference quantile matching on this latent space.
Implementation Details:
- The quantile regression problem is framed as Maximum Likelihood Estimation (MLE) of an Asymmetric Laplace Distribution (ALD).
- The uncertainty around the quantile estimates is modeled using an Inverse-Gamma prior on the scale parameter ( $\sigma_\tau$ ). This posterior variance acts as an implicit epistemic uncertainty, inducing exploration similar to Thompson sampling without explicit exploration bonuses.
- The algorithm utilizes a deep actor-critic architecture with twin critics, min-clipping for Bellman targets, and delayed actor updates (similar to TD3), but replaces standard value targets with the distributional returns derived from the AIF objective.
Key Distinction: Unlike traditional AIF, DAIF does not learn a forward transition model ( $P(X'|X, A)$ ). Instead, it learns the return distribution directly from sampled transitions in the latent space, inheriting the benefits of AIF's state abstraction while avoiding the cost of transition dynamics modeling.

Key Contributions

Theoretical Unification: The paper provides a formal abstraction spanning model-based, distributional, and model-free RL, demonstrating how AIF can be seamlessly integrated into the distributional framework.
Simplified AIF Objective: By using do-calculus, the authors derive a simplified AIF objective that eliminates the need for complex inference of posterior latents in the standard formulation, aligning it more closely with variational inference machinery used in modern RL.
Push-Forward Framework: The introduction of "Push-Forward RL" provides a rigorous mathematical link between state abstractions and distributional Bellman operators, showing how latent space compression affects convergence rates.
DAIF Algorithm: The proposal of DAIF, a practical algorithm that achieves AIF-like performance gains (efficient planning via state abstraction) without the computational overhead of learning explicit transition models.

Experimental Results

The authors evaluate DAIF across tabular and continuous control tasks:

Tabular Experiments (Latent RiverSwim): In a grid world where reward-relevant dynamics lie on a 1D latent manifold, DAIF significantly outperforms both model-based (PSRL-PI) and distributional (IQQL) baselines as the planning horizon increases. When state abstraction offers no advantage (plain RiverSwim), DAIF performs comparably to distributional RL.
Continuous Control: DAIF was tested on three benchmark suites:
- EvoGym: Soft robot locomotion and manipulation.
- DeepMind Control Suite (DMC): Continuous control with varying morphology.
- DMC Vision: Control from raw pixel observations.
Performance: DAIF consistently achieved state-of-the-art performance across all suites, particularly in environments with complex dynamics or high-dimensional state spaces (e.g., "Catcher-v0", "Dog Run", "Quadruped Run").
Efficiency: While DAIF requires approximately 12% more wall-clock time than standard distributional actor-critics, this overhead is lower than that of other model-based or exploration-heavy methods like DSAC (26%) and DRND (37%).

Significance and Claims

The paper claims that Active Inference is particularly powerful when the agent has limited computational capabilities, mirroring the conditions of biological brains. By casting AIF as a simple extension of distributional RL, the authors argue that the performance gains of AIF (efficient organization of sensory information for planning) can be realized without the prohibitive cost of learning explicit forward models.

The work suggests that the "dual problem" of sensory organization and planning is best solved by learning a latent representation where the return distribution is compressed and structured, rather than by modeling the full transition dynamics of the environment. The authors position DAIF not as a replacement for all RL, but as a robust method for controlling challenging environments where sample efficiency and computational constraints are critical. They acknowledge limitations, noting that DAIF does not explicitly separate epistemic and instrumental value terms (relying on implicit uncertainty instead) and that the latent bottleneck size is currently architecture-dependent rather than adaptively learned.